torch_harmonics.distributed.copy_to_polar_region

torch_harmonics.distributed.copy_to_polar_region(input_)

Identity in the forward pass; all-reduce across polar ranks in the backward pass.

Use this to broadcast a replicated tensor into a region where each polar rank will compute a partial result. The backward pass sums the partial gradients so that the replicated parameter receives the correct total gradient.

Parameters:

input (torch.Tensor) – Replicated tensor (same value on every polar rank).

Returns:

Same tensor (forward is a no-op).

Return type:

torch.Tensor