Distributed helpers#

Distributed (multi-GPU) setup, process group accessors, and data partitioning utilities. These are available in the torch_harmonics.distributed subpackage. See the distributed guide for a complete walkthrough.

Setup and teardown#

`init`	Initialize the torch-harmonics distributed backend.
`finalize`	Tear down the torch-harmonics distributed backend.
`is_initialized`	Return `True` if `init()` has been called and `finalize()` has not.
`is_distributed_polar`	Return `True` if a polar process group has been registered.
`is_distributed_azimuth`	Return `True` if an azimuth process group has been registered.

Process group accessors#

`polar_group`	Return the polar (latitudinal) process group registered by `init()`, or `None`.
`polar_group_rank`	Return this rank's index within the polar group (0 if not distributed).
`polar_group_size`	Return the number of ranks in the polar group (1 if not distributed).
`azimuth_group`	Return the azimuth (longitudinal) process group registered by `init()`, or `None`.
`azimuth_group_rank`	Return this rank's index within the azimuth group (0 if not distributed).
`azimuth_group_size`	Return the number of ranks in the azimuth group (1 if not distributed).

Data partitioning#

`compute_split_shapes`	Compute balanced chunk sizes for distributing a dimension across ranks.
`split_tensor_along_dim`	Split a tensor along a given dimension into balanced chunks.