torch_harmonics.distributed.reduce_from_scatter_to_azimuth_region#

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

Fused reduce-scatter across azimuth ranks along dim_.

Azimuth counterpart of reduce_from_scatter_to_polar_region().

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:

torch.Tensor