torch_harmonics.distributed.gather_from_copy_to_polar_region#

torch_harmonics.distributed.gather_from_copy_to_polar_region(input_, dim_, shapes_)[source]#

All-gather along dim_ across polar ranks; reduce-scatter in the backward pass.

Similar to gather_from_polar_region(), but the backward pass uses reduce-scatter instead of split, making it the adjoint of a copy-then-gather pattern.

Parameters:

  • input (torch.Tensor) – Local partition of the tensor.

  • dim (int) – Dimension along which to gather.

  • shapes (list[int]) – Per-rank sizes along dim_.

Returns:

Fully gathered tensor.

Return type:

torch.Tensor