torch_harmonics.distributed.gather_from_polar_region#

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

All-gather along dim_ across polar ranks to reconstruct the full tensor.

The backward pass is a split (scatter) that distributes gradients back to owning ranks.

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