Distributed layers#
Distributed (multi-GPU) counterparts of the serial layers. These are
available in the torch_harmonics.distributed subpackage.
Coverage#
The table below shows which serial layers have a distributed counterpart and which do not.
Serial layer |
Distributed counterpart |
Notes |
|---|---|---|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
— |
Global attention; no distributed version |
|
— |
Sampling utility |
Layer reference#
Distributed version of the forward (real-valued) SHT. |
|
Distributed version of the inverse (real-valued) SHT. |
|
Distributed version of the forward (real) vector SHT. |
|
Distributed version of the inverse (real-valued) vector SHT. |
|
Distributed spectral convolution layer on \(S^2\) implemented with distributed real SHT (Driscoll-Healy formulation, see https://api.semanticscholar.org/CorpusID:122817218). |
|
Distributed version of Discrete-continuous convolutions (DISCO) on the 2-Sphere as described in [1]. |
|
Distributed version of discrete-continuous transpose convolutions (DISCO) on the 2-Sphere as described in [1]. |
|
Distributed neighborhood attention on the 2-sphere using a ring exchange strategy for the longitude dimension and halo exchange for the latitude dimension. |
|
Distributed resampling module for spherical data on the 2-sphere. |
|
Distributed scalar quadrature on \(S^2\) for integrating spherical fields on a latitude/longitude grid, with data and weights split across polar and azimuth communicator groups. |
Note
The custom C++/CUDA kernels are an implementation detail invoked from these Python modules; they have no separately documented API.