torch_harmonics.distributed.scatter_to_polar_region#

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

Split input_ along dim_ and keep only the local polar rank’s chunk.

The backward pass is an all-gather that reconstructs the full tensor.

Parameters:

  • input (torch.Tensor) – Full (non-partitioned) tensor.

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

Returns:

The local rank’s slice of the input.

Return type:

torch.Tensor