torch_harmonics.distributed.DistributedDiscreteContinuousConvS2#

class torch_harmonics.distributed.DistributedDiscreteContinuousConvS2(
in_channels,
out_channels,
in_shape,
out_shape,
kernel_shape,
basis_type='piecewise linear',
basis_norm_mode='nodal',
groups=1,
grid_in='equiangular',
grid_out='equiangular',
bias=True,
theta_cutoff=None,
optimized_kernel=True,
fused=False,
polar_mode='halo-exchange',
)[source]#

Bases: DiscreteContinuousConv

Distributed version of Discrete-continuous convolutions (DISCO) on the 2-Sphere as described in [2]. We assume the data can be split in polar and azimuthal directions.

See also

torch_harmonics.DiscreteContinuousConvS2

Serial counterpart with full mathematical description and parameter documentation.

The azimuth direction is all-to-all: a channel swap so the sparse psi contraction runs against the full nlon_in row, then back to channel-distributed. The polar direction is chosen by polar_mode=, which decides whether a rank computes only the output latitudes it owns (borrowing a halo) or all of them (and reduce-scatters the partial sums). That choice governs how the K-expanded intermediate scales, so it is usually the one that decides whether a large model fits. The fused= flag is orthogonal to it and mirrors the serial conv:

fused=False (default) — standard a2a: einsum after the

transpose-back; the K-expanded intermediate is saved for backward.

fused=True — reordered a2a: the weight einsum runs before the

collectives on the local azimuth channel shard, via the fused contraction+einsum op that recomputes the K-expanded in backward instead of saving it. K× lower activation memory and K× less collective volume, at the cost of one extra contraction in backward. CUDA + optimized kernels only.

Parameters:

  • in_channels (int) – Number of input channels

  • out_channels (int) – Number of output channels

  • in_shape (Tuple[int]) – Shape of the input tensor

  • out_shape (Tuple[int]) – Shape of the output tensor

  • kernel_shape (Union[int, Tuple[int], Tuple[int, int]]) – Shape of the kernel

  • basis_type (Optional[str]) – Type of basis to use

  • basis_norm_mode (Optional[str]) – Normalization mode for the filter basis

  • groups (Optional[int]) – Number of groups

  • grid_in (Optional[str]) – Grid type for the input tensor

  • grid_out (Optional[str]) – Grid type for the output tensor

  • bias (Optional[bool]) – Whether to use bias

  • theta_cutoff (Optional[float]) – Theta cutoff for the filter basis

  • optimized_kernel (Optional[bool]) – Use the optimized CUDA contraction kernel. Required when fused=True.

  • fused (bool) – Mirrors the serial conv. False (default): standard all-to-all (the K-expanded intermediate is saved for backward). True: reordered all-to-all — the weight einsum runs before the collectives on the local azimuth channel shard and the K-expanded is recomputed in backward instead of saved, for K× lower activation memory and K× less collective volume (CUDA + optimized kernels only).

  • polar_mode (Optional[str]) –

    How the polar ranks obtain their output latitudes.

    "halo-exchange" (default): each rank borrows the input rows its own output rows reach into and computes them outright. The K-expanded intermediate is then (B, C, K, nlat_out / polar_group_size, nlon_out), i.e. it shrinks as polar ranks are added. Requires the filter support to reach no further than the neighbouring rank; when it does not, construction raises and says so, rather than quietly selecting the other mode and leaving the memory problem in place.

    "reduce-scatter": every rank computes a partial sum over all output latitudes and a reduce-scatter completes and splits it. No bound on the angular reach, so this is the mode for cutoffs wide enough that a halo cannot serve them – at the cost of an intermediate that stays at the full nlat_out however many polar ranks are used.

Returns:

Output tensor

Return type:

torch.Tensor

References

[2]