torch_harmonics.distributed.distributed_transpose_polar#

torch_harmonics.distributed.distributed_transpose_polar(input_, dims_, shapes_)[source]#

All-to-all transpose across the polar process group.

Same semantics as distributed_transpose_azimuth() but operates on the polar (latitudinal) process group.

Parameters:

  • input (torch.Tensor) – Input tensor, partitioned along dims_[0].

  • dims (tuple[int, int]) – (source_dim, target_dim).

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

Returns:

Transposed tensor, now partitioned along dims_[1].

Return type:

torch.Tensor