Teaching PyTorch that the world isn't flat

What it does torch-harmonics implements differentiable signal processing on the sphere for PyTorch. It provides torch.nn.Modules for spherical harmonic transforms (SHT), vector SHTs, and discrete-continuous convolutions, using quadrature rules for Legendre projections and FFTs for longitudinal modes. The package was originally built to enable Spherical Fourier Neural Operators (SFNO) and differentiable PDE solvers on spherical domains.

The interesting bit Because every operation is built from PyTorch primitives, the entire transform pipeline remains fully differentiable—you can backpropagate through a spectral climate simulation or a neural operator that respects spherical geometry. The quadrature computation can also be distributed across multiple ranks, making the transform spatially parallel for large-scale model training.

Key highlights

• Differentiable SHT and inverse SHT implemented as composable torch.nn.Modules.
• Discrete-continuous (DISCO) convolutions and spherical attention mechanisms that use geodesic distance to define neighborhoods.
• Distributed quadrature support in torch_harmonics.distributed for spatial model parallelism.
• Bundled Jupyter notebooks covering spherical harmonics visualization, spectral fitting, and solving the Helmholtz and shallow water equations.
• Prebuilt CUDA wheels hosted on NVIDIA’s PyPI registry, though the vanilla PyPI release is CPU-only.

Caveats

• Custom CUDA extensions require GPU compute capability >= 7.0.
• The default PyPI wheel is CPU-only; GPU support requires installing a CUDA-specific build from NVIDIA’s index.

Verdict Reach for this if you are building neural operators, climate models, or any differentiable pipeline where the data natively lives on a sphere. If your work is strictly Euclidean, standard PyTorch FFT and convolution layers are the simpler path.