torch_harmonics.distributed.reduce_from_scatter_to_polar_region#
- torch_harmonics.distributed.reduce_from_scatter_to_polar_region(input_, dim_)[source]#
Fused reduce-scatter across polar ranks along
dim_.Equivalent to an all-reduce followed by keeping only the local rank’s chunk, but performed in a single collective for efficiency. The backward pass is an all-gather.
- Parameters:
input (torch.Tensor) – Tensor with partial contributions from the local rank.
dim (int) – Dimension along which to scatter after reducing.
- Returns:
Reduced and scattered tensor (local chunk only).
- Return type: