Upload folder using huggingface_hub

.deps/torch_harmonics-0.8.0.dist-info/INSTALLER

@@ -0,0 +1 @@
+pip

.deps/torch_harmonics-0.8.0.dist-info/METADATA

@@ -0,0 +1,324 @@
+Metadata-Version: 2.4
+Name: torch_harmonics
+Version: 0.8.0
+Summary: Differentiable signal processing on the sphere for PyTorch.
+Author: Boris Bonev, Thorsten Kurth, Max Rietmann, Mauro Bisson, Andrea Paris, Alberto Carpentieri, Massimiliano Fatica, Jean Kossaifi, Nikola Kovachki, Christian Hundt
+Maintainer-email: Boris Bonev <bbonev@nvidia.com>, Thorsten Kurth <tkurth@nvidia.com>
+Classifier: Development Status :: 3 - Alpha
+Classifier: Programming Language :: Python :: 3.9
+Classifier: License :: OSI Approved :: BSD License
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+License-File: LICENSE
+License-File: AUTHORS
+Requires-Dist: torch>=2.4.0
+Requires-Dist: numpy>=1.22.4
+Provides-Extra: dev
+Requires-Dist: pytest>=6.0.0; extra == "dev"
+Requires-Dist: coverage>=6.5.0; extra == "dev"
+Provides-Extra: 2d3ds
+Requires-Dist: requests; extra == "2d3ds"
+Requires-Dist: tarfile; extra == "2d3ds"
+Requires-Dist: tqdm; extra == "2d3ds"
+Requires-Dist: PIL; extra == "2d3ds"
+Requires-Dist: h5py; extra == "2d3ds"
+Dynamic: license-file
+
+<!--
+SPDX-FileCopyrightText: Copyright (c) 2022 The torch-harmonics Authors. All rights reserved.
+
+SPDX-License-Identifier: BSD-3-Clause
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this
+   list of conditions and the following disclaimer.
+2. Redistributions in binary form must reproduce the above copyright notice,
+   this list of conditions and the following disclaimer in the documentation
+   and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its
+   contributors may be used to endorse or promote products derived from
+   this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+-->
+
+<!-- <div align="center">
+    <img src="https://raw.githubusercontent.com/NVIDIA/torch-harmonics/main/images/logo/logo.png"  width="568">
+    <br>
+    <a href="https://github.com/NVIDIA/torch-harmonics/actions/workflows/tests.yml"><img src="https://github.com/NVIDIA/torch-harmonics/actions/workflows/tests.yml/badge.svg"></a>
+    <a href="https://pypi.org/project/torch_harmonics/"><img src="https://img.shields.io/pypi/v/torch_harmonics"></a>
+</div> -->
+
+<!--
+[![pypi](https://img.shields.io/pypi/v/torch_harmonics)](https://pypi.org/project/torch_harmonics/)
+-->
+
+<!-- # spherical harmonic transforms -->
+
+# torch-harmonics
+
+[**Overview**](#overview) | [**Installation**](#installation) | [**More information**](#more-about-torch-harmonics) | [**Getting started**](#getting-started) | [**Contributors**](#contributors) | [**Cite us**](#cite-us) | [**References**](#references)
+
+[![tests](https://github.com/NVIDIA/torch-harmonics/actions/workflows/tests.yml/badge.svg)](https://github.com/NVIDIA/torch-harmonics/actions/workflows/tests.yml)
+[![pypi](https://img.shields.io/pypi/v/torch_harmonics)](https://pypi.org/project/torch_harmonics/)
+
+## Overview
+
+torch-harmonics implements differentiable signal processing on the sphere. This includes differentiable implementations of the spherical harmonic transforms, vector spherical harmonic transforms and discrete-continuous convolutions on the sphere. The package was originally implemented to enable Spherical Fourier Neural Operators (SFNO) [1].
+
+The SHT algorithm uses quadrature rules to compute the projection onto the associated Legendre polynomials and FFTs for the projection onto the harmonic basis. This algorithm tends to outperform others with better asymptotic scaling for most practical purposes [2].
+
+torch-harmonics uses PyTorch primitives to implement these operations, making it fully differentiable. Moreover, the quadrature can be distributed onto multiple ranks making it spatially distributed.
+
+torch-harmonics has been used to implement a variety of differentiable PDE solvers which generated the animations below. Moreover, it has enabled the development of Spherical Fourier Neural Operators  [1].
+
+<div align="center">
+<table border="0" cellspacing="0" cellpadding="0">
+    <tr>
+        <td><img src="https://media.githubusercontent.com/media/NVIDIA/torch-harmonics/main/images/sfno.gif"  width="240"></td>
+        <td><img src="https://media.githubusercontent.com/media/NVIDIA/torch-harmonics/main/images/zonal_jet.gif"  width="240"></td>
+        <td><img src="https://media.githubusercontent.com/media/NVIDIA/torch-harmonics/main/images/allen-cahn.gif"  width="240"></td>
+    </tr>
+<!--     <tr>
+        <td style="text-align:center; border-style : hidden!important;">Shallow Water Eqns.</td>
+        <td style="text-align:center; border-style : hidden!important;">Ginzburg-Landau Eqn.</td>
+        <td style="text-align:center; border-style : hidden!important;">Allen-Cahn Eqn.</td>
+    </tr>  -->
+</table>
+</div>
+
+
+## Installation
+A simple installation can be directly done from PyPI:
+
+```bash
+pip install torch-harmonics
+```
+If you are planning to use spherical convolutions, we recommend building the corresponding custom CUDA kernels. To enforce this, you can set the `FORCE_CUDA_EXTENSION` flag. You may also want to set appropriate architectures with the `TORCH_CUDA_ARCH_LIST` flag. Finally, make sure to disable build isolation via the `--no-build-isolation` flag to ensure that the custom kernels are built with the existing torch installation.
+```bash
+export FORCE_CUDA_EXTENSION=1
+export TORCH_CUDA_ARCH_LIST="7.0 7.2 7.5 8.0 8.6 8.7 9.0+PTX"
+pip install --no-build-isolation torch-harmonics
+```
+:warning: Please note that the custom CUDA extensions currently only support CUDA architectures >= 7.0.
+
+If you want to actively develop torch-harmonics, we recommend building it in your environment from github:
+
+```bash
+git clone git@github.com:NVIDIA/torch-harmonics.git
+cd torch-harmonics
+pip install -e .
+```
+
+Alternatively, use the Dockerfile to build your custom container after cloning:
+
+```bash
+git clone git@github.com:NVIDIA/torch-harmonics.git
+cd torch-harmonics
+docker build . -t torch_harmonics
+docker run --gpus all -it --rm --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 torch_harmonics
+```
+
+## More about torch-harmonics
+
+### Spherical harmonics
+
+The [spherical harmonics](https://en.wikipedia.org/wiki/Spherical_harmonics) are special functions defined on the two-dimensional sphere $S^2$ (embedded in three dimensions). They form an orthonormal basis of the space of square-integrable functions defined on the sphere $L^2(S^2)$ and are comparable to the harmonic functions defined on a circle/torus. The spherical harmonics are defined as
+
+$$
+Y_l^m(\theta, \lambda) = \sqrt{\frac{(2l + 1)}{4 \pi} \frac{(l - m)!}{(l + m)!}} P_l^m(\cos \theta)  \exp(im\lambda),
+$$
+
+where $\theta$ and $\lambda$ are colatitude and longitude respectively, and $P_l^m$ the normalized, [associated Legendre polynomials](https://en.wikipedia.org/wiki/Associated_Legendre_polynomials).
+
+<div align="center">
+<img src="https://media.githubusercontent.com/media/NVIDIA/torch-harmonics/main/images/spherical_harmonics.gif" width="432">
+<br>
+Spherical harmonics up to degree 5
+</div>
+
+### Spherical harmonic transform
+
+The spherical harmonic transform (SHT)
+
+$$
+f_l^m = \int_{S^2}  \overline{Y_{l}^{m}}(\theta, \lambda) f(\theta, \lambda) \mathrm{d} \mu(\theta, \lambda)
+$$
+
+realizes the projection of a signal $f(\theta, \lambda)$ on $S^2$ onto the spherical harmonics basis. The SHT generalizes the Fourier transform on the sphere. Conversely, a truncated series expansion of a function $f$ can be written in terms of spherical harmonics as
+
+$$
+f (\theta, \lambda) = \sum_{m=-M}^{M} \exp(im\lambda) \sum_{l=|m|}^{M} \hat f_l^m  P_l^m (\cos \theta),
+$$
+
+where $\hat{f}_l^m$, are the expansion coefficients associated to the mode $m$, $n$.
+
+The implementation of the SHT follows the algorithm as presented in [2]. A direct spherical harmonic transform can be accomplished by a Fourier transform
+
+$$
+\hat f^m(\theta) = \frac{1}{2 \pi} \int_{0}^{2\pi} f(\theta, \lambda) \exp(-im\lambda) \mathrm{d} \lambda
+$$
+
+in longitude and a Legendre transform
+
+$$
+\hat f_l^m = \frac{1}{2} \int^{\pi}_0 \hat f^{m} (\theta) P_l^m (\cos \theta) \sin \theta \mathrm{d} \theta
+$$
+
+in latitude.
+
+### Discrete Legendre transform
+
+The second integral, which computed the projection onto the Legendre polynomials is realized with quadrature. On the Gaussian grid, we use Gaussian quadrature in the $\cos \theta$ domain. The integral
+
+$$
+\hat f_l^m = \frac{1}{2} \int_{-1}^1 \hat{f}^m(\arccos x) P_l^m (x) \mathrm{d} x
+$$
+
+is obtained with the substitution $x = \cos \theta$ and then approximated by the sum
+
+$$
+\hat f_l^m = \sum_{j=1}^{N_\theta}  \hat{f}^m(\arccos x_j) P_l^m(x_j) w_j.
+$$
+
+Here, $x_j \in [-1,1]$ are the quadrature nodes with the respective quadrature weights $w_j$.
+
+### Discrete-continuous convolutions on the sphere
+
+torch-harmonics now provides local discrete-continuous (DISCO) convolutions as outlined in [5] on the sphere. These are use in local neural operators [2] to generalize convolutions to structured and unstructured meshes on the sphere.
+
+### Spherical (neighborhood) attention
+
+torch-harmonics introducers spherical attention mechanisms which correctly generalize the attention mechanism to the sphere. The use of quadrature rules makes the resulting operations approximately equivariant and equivariant in the continuous limit. Moreover, neighborhood attention is correctly generalized onto the sphere by using the geodesic distance to determine the size of the neighborhood.
+
+## Getting started
+
+The main functionality of `torch_harmonics` is provided in the form of `torch.nn.Modules` for composability. A minimum example is given by:
+
+```python
+import torch
+import torch_harmonics as th
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+nlat = 512
+nlon = 2*nlat
+batch_size = 32
+signal = torch.randn(batch_size, nlat, nlon, device=device)
+
+# transform data on an equiangular grid
+sht = th.RealSHT(nlat, nlon, grid="equiangular").to(device)
+
+coeffs = sht(signal)
+```
+
+To enable scalable model-parallelism, `torch-harmonics` implements a distributed variant of the SHT located in `torch_harmonics.distributed`.
+
+Detailed usage of torch-harmonics, alongside helpful analysis provided in a series of notebooks:
+
+1. [Getting started](./notebooks/getting_started.ipynb)
+2. [Quadrature](./notebooks/quadrature.ipynb)
+3. [Visualizing the spherical harmonics](./notebooks/plot_spherical_harmonics.ipynb)
+4. [Spectral fitting vs. SHT](./notebooks/gradient_analysis.ipynb)
+5. [Conditioning of the Gramian](./notebooks/conditioning_sht.ipynb)
+6. [Solving the Helmholtz equation](./notebooks/helmholtz.ipynb)
+7. [Solving the shallow water equations](./notebooks/shallow_water_equations.ipynb)
+8. [Training Spherical Fourier Neural Operators (SFNO)](./notebooks/train_sfno.ipynb)
+9. [Resampling signals on the sphere](./notebooks/resample_sphere.ipynb)
+
+## Examples and reproducibility
+
+The `examples` folder contains training scripts for three distinct tasks:
+
+* [solution of the shallow water equations on the rotating sphere](./examples/shallow_water_equations/train.py)
+* [depth estimation on the sphere](./examples/depth/train.py)
+* [semantic segmentation on the sphere](./examples/segmentation/train.py)
+
+Results from the papers can generally be reproduced by running `python train.py`. In the case of some older results the number of epochs and learning-rate may need to be adjusted by passing the corresponding command line argument.
+
+## Remarks on automatic mixed precision (AMP) support
+
+Note that torch-harmonics uses Fourier transforms from `torch.fft` which in turn uses kernels from the optimized `cuFFT` library. This library supports fourier transforms of `float32` and `float64` (i.e. `single` and `double` precision) tensors for all input sizes. For `float16` (i.e. `half` precision) and `bfloat16` inputs however, the dimensions which are transformed are restricted to powers of two. Since data is converted to one of these reduced precision floating point formats when `torch.autocast` is used, torch-harmonics will issue an error when the input shapes are not powers of two. For these cases, we recommend disabling autocast for the harmonics transform specifically:
+
+```python
+import torch
+import torch_harmonics as th
+
+sht = th.RealSHT(512, 1024, grid="equiangular").cuda()
+
+with torch.autocast(device_type="cuda", enabled = True):
+   # do some AMP converted math here
+   x = some_math(x)
+   # convert tensor to float32
+   x = x.to(torch.float32)
+   # now disable autocast specifically for the transform,
+   # making sure that the tensors are not converted
+   # back to reduced precision internally
+   with torch.autocast(device_type="cuda", enabled = False):
+      xt = sht(x)
+
+   # continue operating on the transformed tensor
+   xt = some_more_math(xt)
+```
+
+Depending on the problem, it might be beneficial to upcast data to `float64` instead of `float32` precision for numerical stability.
+
+## Contributors
+
+[Boris Bonev](https://bonevbs.github.io) (bbonev@nvidia.com), [Thorsten Kurth](https://github.com/azrael417) (tkurth@nvidia.com), [Max Rietmann](https://github.com/rietmann-nv), [Mauro Bisson](https://scholar.google.com/citations?hl=en&user=f0JE-0gAAAAJ), [Andrea Paris](https://github.com/apaaris), [Alberto Carpentieri](https://github.com/albertocarpentieri), [Massimiliano Fatica](https://scholar.google.com/citations?user=Deaq4uUAAAAJ&hl=en), [Nikola Kovachki](https://kovachki.github.io), [Jean Kossaifi](http://jeankossaifi.com), [Christian Hundt](https://github.com/gravitino)
+
+## Cite us
+
+If you use `torch-harmonics` in an academic paper, please cite [1]
+
+```bibtex
+@misc{bonev2023spherical,
+      title={Spherical Fourier Neural Operators: Learning Stable Dynamics on the Sphere},
+      author={Boris Bonev and Thorsten Kurth and Christian Hundt and Jaideep Pathak and Maximilian Baust and Karthik Kashinath and Anima Anandkumar},
+      year={2023},
+      eprint={2306.03838},
+      archivePrefix={arXiv},
+      primaryClass={cs.LG}
+}
+```
+
+## References
+
+<a id="1">[1]</a>
+Bonev B., Kurth T., Hundt C., Pathak, J., Baust M., Kashinath K., Anandkumar A.;
+Spherical Fourier Neural Operators: Learning Stable Dynamics on the Sphere;
+International Conference on Machine Learning, 2023. [arxiv link](https://arxiv.org/abs/2306.03838)
+
+<a id="1">[2]</a>
+Liu-Schiaffini M., Berner J., Bonev B., Kurth T., Azizzadenesheli K., Anandkumar A.;
+Neural Operators with Localized Integral and Differential Kernels;
+International Conference on Machine Learning, 2024. [arxiv link](https://arxiv.org/abs/2402.16845)
+
+<a id="1">[3]</a>
+Schaeffer N.;
+Efficient spherical harmonic transforms aimed at pseudospectral numerical simulations;
+G3: Geochemistry, Geophysics, Geosystems, 2013.
+
+<a id="1">[4]</a>
+Wang B., Wang L., Xie Z.;
+Accurate calculation of spherical and vector spherical harmonic expansions via spectral element grids;
+Adv Comput Math, 2018.
+
+<a id="1">[5]</a>
+Ocampo, Price, McEwen, Scalable and equivariant spherical CNNs by discrete-continuous (DISCO) convolutions, ICLR (2023), arXiv:2209.13603
+
+<a id="1">[6]</a>
+Bonev B., Rietmann M., Paris A., Carpentieri A., Kurth T.; Attention on the Sphere; [arxiv link](https://arxiv.org/abs/2505.11157)

.deps/torch_harmonics-0.8.0.dist-info/RECORD

@@ -0,0 +1,86 @@
+torch_harmonics-0.8.0.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4
+torch_harmonics-0.8.0.dist-info/METADATA,sha256=R5upGTlvdXzHVZsns-Tvrf-DYjeVrKfZAEMbmcXwEGU,16077
+torch_harmonics-0.8.0.dist-info/RECORD,,
+torch_harmonics-0.8.0.dist-info/REQUESTED,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+torch_harmonics-0.8.0.dist-info/WHEEL,sha256=K260EYznzXsJYBQGqmI8VTxEdiZYNvDZwW9cBh9-_MA,91
+torch_harmonics-0.8.0.dist-info/licenses/AUTHORS,sha256=Lee9XR9hIxhBwbjp8RS0dZMwPDqtfoWnxxAGe1zLR3o,409
+torch_harmonics-0.8.0.dist-info/licenses/LICENSE,sha256=vexnFPd7j3gnhT4EuEB5Fhgbpm_5DmSY5g5y2lDSXUI,1576
+torch_harmonics-0.8.0.dist-info/top_level.txt,sha256=fPyCWy2b6BcyqFoikrfbgESNiG0oABgOFU2r5hgFhgw,16
+torch_harmonics/__init__.py,sha256=MW9sMszgWYufNlBaHgkmCc6bicejzRsaF_W8F3WAsQw,1983
+torch_harmonics/__pycache__/__init__.cpython-311.pyc,,
+torch_harmonics/__pycache__/_disco_convolution.cpython-311.pyc,,
+torch_harmonics/__pycache__/_neighborhood_attention.cpython-311.pyc,,
+torch_harmonics/__pycache__/attention.cpython-311.pyc,,
+torch_harmonics/__pycache__/cache.cpython-311.pyc,,
+torch_harmonics/__pycache__/convolution.cpython-311.pyc,,
+torch_harmonics/__pycache__/filter_basis.cpython-311.pyc,,
+torch_harmonics/__pycache__/legendre.cpython-311.pyc,,
+torch_harmonics/__pycache__/plotting.cpython-311.pyc,,
+torch_harmonics/__pycache__/quadrature.cpython-311.pyc,,
+torch_harmonics/__pycache__/random_fields.cpython-311.pyc,,
+torch_harmonics/__pycache__/resample.cpython-311.pyc,,
+torch_harmonics/__pycache__/sht.cpython-311.pyc,,
+torch_harmonics/_disco_convolution.py,sha256=znDWSirhJpBWX1yNJH6Z9vP81FmEBhXHcRuhNW7dxN0,9744
+torch_harmonics/_neighborhood_attention.py,sha256=eglziOF2yfpt4-1du2hFOqJ2ZaVHkwWVEhuhyUIdZuc,24120
+torch_harmonics/attention.py,sha256=OmuY8HKpcMbLbzzGBP9dtvT0uBNTkFBeHLJg94tNSok,16296
+torch_harmonics/cache.py,sha256=2W_ug1_skw4BP880CazlELzpx2ASxE_bopWHJ1uKjts,3023
+torch_harmonics/convolution.py,sha256=zr6dL8SKnmygU-h3rnX79PK5scK40t8C6yadltMj6n4,24911
+torch_harmonics/csrc/attention/attention.cuh,sha256=a4B-Ejs-kXMuAPnrmS9Mujy-m5Z48yj9BlkEOXNMIRw,2959
+torch_harmonics/csrc/attention/attention_bwd_cuda.cu,sha256=V8Y-7s6aHPfB0Wklks_BEUQ9iYo4EDMuwzPDJ2cxpac,40904
+torch_harmonics/csrc/attention/attention_fwd_cuda.cu,sha256=MnDys8h_ZgSEDbrb-QvcnHL3RTOA6NYp-gdnNuFuevk,21851
+torch_harmonics/csrc/attention/attention_interface.cu,sha256=lFaHi8u5spNwt_3f-NclhcQKEaoKNUGXEG3ikrL22p4,1944
+torch_harmonics/csrc/attention/attention_utils.cu,sha256=5kLzIlUfr45yb9bo5Xac6f1NbJKBrT7BMoxFHYNmIB8,7036
+torch_harmonics/csrc/attention/attention_utils.cuh,sha256=1o9FshZSYbwCFPiFR1InHUgFPZ5hJJg2JpJ8n9mgyiY,11771
+torch_harmonics/csrc/attention/cudamacro.h,sha256=rcc3li9Z2Fy4eOxPm0vRITzsHUsziaNtBX2mXnAJ06g,2633
+torch_harmonics/csrc/disco/disco.h,sha256=_fu-JrHMJTC4KnMQ1hj9dcnWIXSbXaiFiPbKIJoIJFo,1907
+torch_harmonics/csrc/disco/disco_cuda.cuh,sha256=VEowGAdWhaHUx4UeKVC1714sXPGKpcFKBXvsEnUrlXE,2699
+torch_harmonics/csrc/disco/disco_cuda_bwd.cu,sha256=18UGV2g6ORO-3aW_h_1SgeKwsqeMig4I9P8p5fqmD1E,12598
+torch_harmonics/csrc/disco/disco_cuda_fwd.cu,sha256=vFWpIR-GOvJjTzHBSuvU4N97JfmIjyENcbQuGlFq5iU,11588
+torch_harmonics/csrc/disco/disco_helpers.cpp,sha256=uLG4wYUPNFCHeGlRLWYhS-73yt0iayq0DYiiIDadSyM,5310
+torch_harmonics/csrc/disco/disco_interface.cu,sha256=57xKiwItqwA1pEsM-VglLGb-_rAsI1JHs2nvHw4B7e4,1875
+torch_harmonics/distributed/__init__.py,sha256=tusU6vOpKMwSV8OkIEqAL4jfUcuzceYhkh3wV0jjsgg,2693
+torch_harmonics/distributed/__pycache__/__init__.cpython-311.pyc,,
+torch_harmonics/distributed/__pycache__/distributed_convolution.cpython-311.pyc,,
+torch_harmonics/distributed/__pycache__/distributed_resample.cpython-311.pyc,,
+torch_harmonics/distributed/__pycache__/distributed_sht.cpython-311.pyc,,
+torch_harmonics/distributed/__pycache__/primitives.cpython-311.pyc,,
+torch_harmonics/distributed/__pycache__/utils.cpython-311.pyc,,
+torch_harmonics/distributed/distributed_convolution.py,sha256=4feNPcchZ8BglBLr6nnJbZ39TEiZ_WHpacPzBiiAs2M,19651
+torch_harmonics/distributed/distributed_resample.py,sha256=1fSAveUH1YOzVWgPYSaV__UD6OV3EEjXrmanWUsX5wY,11138
+torch_harmonics/distributed/distributed_sht.py,sha256=J7m9fcxhSW3YE51MrR7Ej6TUY0kr0YPYL6IqMzFvbGg,24988
+torch_harmonics/distributed/primitives.py,sha256=TlBJSgEGPCuzhIrLb_ChBbmz_I1xos7va-UJwetFtGc,15052
+torch_harmonics/distributed/utils.py,sha256=orUloyD9taV11TkRjtOLuJCpI49ObeybjvqYgxYRpGg,3301
+torch_harmonics/examples/__init__.py,sha256=FgSFWGxp4m0oPdzyJYS5D6FO7AaXlX6bp1NbJd6P6FQ,1910
+torch_harmonics/examples/__pycache__/__init__.cpython-311.pyc,,
+torch_harmonics/examples/__pycache__/losses.cpython-311.pyc,,
+torch_harmonics/examples/__pycache__/pde_dataset.cpython-311.pyc,,
+torch_harmonics/examples/__pycache__/pde_sphere.cpython-311.pyc,,
+torch_harmonics/examples/__pycache__/shallow_water_equations.cpython-311.pyc,,
+torch_harmonics/examples/__pycache__/stanford_2d3ds_dataset.cpython-311.pyc,,
+torch_harmonics/examples/losses.py,sha256=MST2ImGVjvu9cfueURM3GOIO44MY29tWDE_PxlPVD9s,15480
+torch_harmonics/examples/metrics.py,sha256=QOHU2cfY7oDmsnTNhgRWfODAIu95_w6IUPf_OZ3wqmo,10914
+torch_harmonics/examples/models/__init__.py,sha256=ovMQpqcnFo1HG3uwiXS0EoErOSRbENazuDrOjkcVzPg,1850
+torch_harmonics/examples/models/__pycache__/__init__.cpython-311.pyc,,
+torch_harmonics/examples/models/__pycache__/_layers.cpython-311.pyc,,
+torch_harmonics/examples/models/__pycache__/lsno.cpython-311.pyc,,
+torch_harmonics/examples/models/__pycache__/s2segformer.cpython-311.pyc,,
+torch_harmonics/examples/models/__pycache__/s2transformer.cpython-311.pyc,,
+torch_harmonics/examples/models/__pycache__/s2unet.cpython-311.pyc,,
+torch_harmonics/examples/models/__pycache__/sfno.cpython-311.pyc,,
+torch_harmonics/examples/models/_layers.py,sha256=2B5WqBGmQ3YWk9i7CngB9auqErmh4utbX79uyQX8I2o,22696
+torch_harmonics/examples/models/lsno.py,sha256=K8TA7Jvnt13OszNdCpcnh9NtIRNSIvGuobvgwi50huA,22295
+torch_harmonics/examples/models/s2segformer.py,sha256=AKhEHH6Pbw3hjFPWpp77ni5NB_W2qQc5YSsiFQj3R2M,27198
+torch_harmonics/examples/models/s2transformer.py,sha256=FzxsB2Z8JDpzdlee9zw4XWYaBZMoufNYPb_AQ9yrf5Q,20331
+torch_harmonics/examples/models/s2unet.py,sha256=s_oWn1bUoWswRTWBTOMM9hKLLSpwxhd6K5uIBU4yGNI,21564
+torch_harmonics/examples/models/sfno.py,sha256=eRYTaUA6VoFtzAWdafIY7C6GoyjEdsCLDJSH64QptgI,16005
+torch_harmonics/examples/pde_dataset.py,sha256=mvY6FBcY7PSd7eLRGr2bmXbz7JPVXeTO2X5CnNpClHk,5564
+torch_harmonics/examples/pde_sphere.py,sha256=cupdwimLpVco0bghk4VWaI3bbY4oV8eOIxfnUdYshcg,8256
+torch_harmonics/examples/shallow_water_equations.py,sha256=yfZUw_Fnn-N40vvom6lYNv_Nr1lSQGr0yH9eLbc5JhA,16103
+torch_harmonics/examples/stanford_2d3ds_dataset.py,sha256=WMvjEGxn7Vr697n3iBdOhmG5ox-xKqomwMbPmG8pQ3w,30650
+torch_harmonics/filter_basis.py,sha256=7104x7Fvof_G9FjwJJdhPk0BbR3mpPp-lVVspVqtrjM,13140
+torch_harmonics/legendre.py,sha256=wDog2ViC3JIU-bQqFzRXCbWitTD9Xoh7Riu99wmNA6k,9008
+torch_harmonics/plotting.py,sha256=7vYee0mv1lVa0ztl1uUNYmjpVNggTWHC-9ZqR5q3uWA,8822
+torch_harmonics/quadrature.py,sha256=a7emSieECNDL6e8StLxorZYTXSPOkidHKSuz5JaoL1s,10536
+torch_harmonics/random_fields.py,sha256=yGOfM24RUdHWMFCwjmfG2Tsxa6CftdIiXNgQYl08y-E,4399
+torch_harmonics/resample.py,sha256=lNX-6F0XFMkSdOqJaTiO6bUJTF1d-0b9m9cHz_CSIcg,8409
+torch_harmonics/sht.py,sha256=ig1J0qpHw-6nlcdkVIo3-UdptU_BkwZGqqtPJnXqcVQ,19118

.deps/torch_harmonics-0.8.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (83.0.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

.deps/torch_harmonics-0.8.0.dist-info/licenses/AUTHORS

@@ -0,0 +1,12 @@
+The code was authored by the following people:
+
+Boris Bonev - NVIDIA Corporation
+Thorsten Kurth - NVIDIA Corporation
+Max Rietmann - NVIDIA Corporation
+Mauro Bisson - NVIDIA Corporation
+Andrea Paris - NVIDIA Corporation
+Alberto Carpentieri - NVIDIA Corporation
+Massimiliano Fatica - NVIDIA Corporation
+Jean Kossaifi - NVIDIA Corporation
+Nikola Kovachki - NVIDIA Corporation
+Christian Hundt - NVIDIA Corporation

.deps/torch_harmonics-0.8.0.dist-info/licenses/LICENSE

@@ -0,0 +1,27 @@
+SPDX-FileCopyrightText: Copyright (c) 2022 The torch-harmonics Authors. All rights reserved.
+
+SPDX-License-Identifier: BSD-3-Clause
+ 
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this
+   list of conditions and the following disclaimer.
+2. Redistributions in binary form must reproduce the above copyright notice,
+   this list of conditions and the following disclaimer in the documentation
+   and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its
+   contributors may be used to endorse or promote products derived from
+   this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

.deps/torch_harmonics-0.8.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+torch_harmonics

.deps/torch_harmonics/__init__.py

@@ -0,0 +1,40 @@
+# coding=utf-8
+
+# SPDX-FileCopyrightText: Copyright (c) 2022 The torch-harmonics Authors. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+#
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+#
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+#
+
+__version__ = "0.8.0"
+
+from .sht import RealSHT, InverseRealSHT, RealVectorSHT, InverseRealVectorSHT
+from .convolution import DiscreteContinuousConvS2, DiscreteContinuousConvTransposeS2
+from .resample import ResampleS2
+from .attention import AttentionS2, NeighborhoodAttentionS2
+from . import quadrature
+from . import random_fields
+from . import examples

.deps/torch_harmonics/_disco_convolution.py

@@ -0,0 +1,212 @@
+# coding=utf-8
+
+# SPDX-FileCopyrightText: Copyright (c) 2022 The torch-harmonics Authors. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+#
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+#
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+#
+
+from typing import Optional
+import math
+
+import torch
+from torch.amp import custom_fwd, custom_bwd
+
+try:
+    import disco_cuda_extension
+except ImportError as err:
+    disco_cuda_extension = None
+
+# some helper functions
+def _get_psi(kernel_size: int, psi_idx: torch.Tensor, psi_vals: torch.Tensor, nlat_in: int, nlon_in: int, nlat_out: int, nlon_out: int, nlat_in_local: Optional[int] = None, nlat_out_local: Optional[int] = None, semi_transposed: Optional[bool] = False):
+    """Creates a sparse tensor for spherical harmonic convolution operations."""
+    nlat_in_local = nlat_in_local if nlat_in_local is not None else nlat_in
+    nlat_out_local = nlat_out_local if nlat_out_local is not None else nlat_out
+    
+    if semi_transposed:
+        # do partial transpose
+        # we do a semi-transposition to faciliate the computation
+        tout = psi_idx[2] // nlon_out
+        pout = psi_idx[2] % nlon_out
+        # flip the axis of longitudes
+        pout = nlon_out - 1 - pout
+        tin = psi_idx[1]
+        idx = torch.stack([psi_idx[0], tout, tin * nlon_out + pout], dim=0)
+        psi = torch.sparse_coo_tensor(idx, psi_vals, size=(kernel_size, nlat_out_local, nlat_in_local * nlon_out)).coalesce()
+    else:
+        psi = torch.sparse_coo_tensor(psi_idx, psi_vals, size=(kernel_size, nlat_out_local, nlat_in_local * nlon_in)).coalesce()
+    return psi
+
+
+class _DiscoS2ContractionCuda(torch.autograd.Function):
+    @staticmethod
+    @custom_fwd(device_type="cuda")
+    def forward(ctx, x: torch.Tensor, roff_idx: torch.Tensor, ker_idx: torch.Tensor,
+                row_idx: torch.Tensor, col_idx: torch.Tensor, vals: torch.Tensor,
+                kernel_size: int, nlat_out: int, nlon_out: int):
+        
+        ctx.save_for_backward(roff_idx, ker_idx, row_idx, col_idx, vals)
+        ctx.kernel_size = kernel_size
+        ctx.nlat_in = x.shape[-2]
+        ctx.nlon_in = x.shape[-1]
+        xtype = x.dtype
+        x = x.to(torch.float32).contiguous()
+        output = disco_cuda_extension.forward(x, roff_idx, ker_idx, row_idx, col_idx, vals, kernel_size, nlat_out, nlon_out)
+        output = output.to(xtype)
+
+        return output
+
+    @staticmethod
+    @custom_bwd(device_type="cuda")
+    def backward(ctx, grad_output):
+
+        roff_idx, ker_idx, row_idx, col_idx, vals = ctx.saved_tensors
+        gtype =	grad_output.dtype
+        grad_output = grad_output.to(torch.float32).contiguous()
+        grad_input = disco_cuda_extension.backward(grad_output, roff_idx, ker_idx, row_idx, col_idx, vals,
+                                         ctx.kernel_size, ctx.nlat_in, ctx.nlon_in)
+        grad_input = grad_input.to(gtype)
+
+        return grad_input, None, None, None, None, None, None, None, None
+
+
+class _DiscoS2TransposeContractionCuda(torch.autograd.Function):
+    @staticmethod
+    @custom_fwd(device_type="cuda")
+    def forward(ctx, x: torch.Tensor, roff_idx: torch.Tensor, ker_idx: torch.Tensor,
+                row_idx: torch.Tensor, col_idx: torch.Tensor, vals: torch.Tensor,
+                kernel_size: int, nlat_out: int, nlon_out: int):
+        
+        ctx.save_for_backward(roff_idx, ker_idx, row_idx, col_idx, vals)
+        ctx.kernel_size = kernel_size
+        ctx.nlat_in = x.shape[-2]
+        ctx.nlon_in = x.shape[-1]
+        xtype =	x.dtype
+        x = x.to(torch.float32).contiguous()
+        output = disco_cuda_extension.backward(x, roff_idx, ker_idx, row_idx, col_idx, vals, kernel_size, nlat_out, nlon_out)
+        output = output.to(xtype)
+
+        return output
+
+    @staticmethod
+    @custom_bwd(device_type="cuda")
+    def backward(ctx, grad_output):
+       
+        roff_idx, ker_idx, row_idx, col_idx, vals = ctx.saved_tensors
+        gtype = grad_output.dtype
+        grad_output = grad_output.to(torch.float32).contiguous()
+        grad_input = disco_cuda_extension.forward(grad_output, roff_idx, ker_idx, row_idx, col_idx, vals,
+                                        ctx.kernel_size, ctx.nlat_in, ctx.nlon_in)
+        grad_input = grad_input.to(gtype)
+
+        return grad_input, None, None, None, None, None, None, None, None
+
+# CUDA
+def _disco_s2_contraction_cuda(x: torch.Tensor, roff_idx: torch.Tensor, ker_idx: torch.Tensor,
+                               row_idx: torch.Tensor, col_idx: torch.Tensor, vals: torch.Tensor,
+                               kernel_size: int, nlat_out: int, nlon_out: int) -> torch.Tensor:
+    return _DiscoS2ContractionCuda.apply(x, roff_idx, ker_idx, row_idx, col_idx, vals,
+                                         kernel_size, nlat_out, nlon_out)
+
+def _disco_s2_transpose_contraction_cuda(x: torch.Tensor, roff_idx: torch.Tensor, ker_idx: torch.Tensor,
+                                         row_idx: torch.Tensor, col_idx: torch.Tensor, vals: torch.Tensor,
+                                         kernel_size: int, nlat_out: int, nlon_out: int) -> torch.Tensor:
+    return _DiscoS2TransposeContractionCuda.apply(x, roff_idx, ker_idx, row_idx, col_idx, vals,
+                                                  kernel_size, nlat_out, nlon_out)
+
+
+def _disco_s2_contraction_torch(x: torch.Tensor, psi: torch.Tensor, nlon_out: int):
+    """
+    Reference implementation of the custom contraction as described in [1]. This requires repeated
+    shifting of the input tensor, which can potentially be costly. For an efficient implementation
+    on GPU, make sure to use the custom kernel written in CUDA.
+    """
+    
+    assert len(psi.shape) == 3
+    assert len(x.shape) == 4
+    psi = psi.to(x.device)
+
+    batch_size, n_chans, nlat_in, nlon_in = x.shape
+    kernel_size, nlat_out, _ = psi.shape
+
+    assert psi.shape[-1] == nlat_in * nlon_in
+    assert nlon_in % nlon_out == 0
+    assert nlon_in >= nlat_out
+    pscale = nlon_in // nlon_out
+
+    # add a dummy dimension for nkernel and move the batch and channel dims to the end
+    x = x.reshape(1, batch_size * n_chans, nlat_in, nlon_in).permute(0, 2, 3, 1)
+    x = x.expand(kernel_size, -1, -1, -1)
+
+    y = torch.zeros(nlon_out, kernel_size, nlat_out, batch_size * n_chans, device=x.device, dtype=x.dtype)
+
+    for pout in range(nlon_out):
+        # sparse contraction with psi
+        y[pout] = torch.bmm(psi, x.reshape(kernel_size, nlat_in * nlon_in, -1))
+        # we need to repeatedly roll the input tensor to faciliate the shifted multiplication
+        x = torch.roll(x, -pscale, dims=2)
+
+    # reshape y back to expose the correct dimensions
+    y = y.permute(3, 1, 2, 0).reshape(batch_size, n_chans, kernel_size, nlat_out, nlon_out)
+
+    return y
+
+
+def _disco_s2_transpose_contraction_torch(x: torch.Tensor, psi: torch.Tensor, nlon_out: int):
+    assert len(psi.shape) == 3
+    assert len(x.shape) == 5
+    psi = psi.to(x.device)
+
+    batch_size, n_chans, kernel_size, nlat_in, nlon_in = x.shape
+    kernel_size, nlat_out, n_out = psi.shape
+
+    assert n_out % nlon_out == 0
+    assert nlon_out >= nlon_in
+    pscale = nlon_out // nlon_in
+
+    # interleave zeros along the longitude dimension to allow for fractional offsets to be considered
+    x_ext = torch.zeros(kernel_size, nlat_in, nlon_out, batch_size * n_chans, device=x.device, dtype=x.dtype)
+    x = x.reshape(batch_size * n_chans, kernel_size, nlat_in, nlon_in).permute(1, 2, 3, 0)
+
+    # x has shape kernel_size x nlat_in x nlon_in x batch_size * n_chans
+    # we only need to apoply the nlon stride here, since nlat stride is taken care of by the kernel
+    x_ext[:, :, ::pscale, :] = x[...]
+
+    # create output tensor
+    y = torch.zeros(kernel_size, nlon_out, nlat_out, batch_size * n_chans, device=x.device, dtype=x.dtype)
+
+    for pout in range(nlon_out):
+        # we need to repeatedly roll the input tensor to faciliate the shifted multiplication
+        # TODO: double-check why this has to happen first
+        x_ext = torch.roll(x_ext, -1, dims=2)
+        # sparse contraction with the modified psi
+        y[:, pout, :, :] = torch.bmm(psi, x_ext.reshape(kernel_size, nlat_in * nlon_out, -1))
+
+    # sum over the kernel dimension and reshape to the correct output size
+    y = y.sum(dim=0).permute(2, 1, 0).reshape(batch_size, n_chans, nlat_out, nlon_out).contiguous()
+
+    return y
+

.deps/torch_harmonics/_neighborhood_attention.py

@@ -0,0 +1,589 @@
+# coding=utf-8
+
+# SPDX-FileCopyrightText: Copyright (c) 2025 The torch-harmonics Authors. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+#
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+#
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+#
+
+import math
+from typing import Union
+
+import torch
+import torch.nn.functional as F
+from torch.amp import custom_fwd, custom_bwd
+
+try:
+    import attention_cuda_extension
+    _cuda_extension_available = True
+except ImportError as err:
+    attention_cuda_extension = None
+    _cuda_extension_available = False
+
+# s2 neighborhood attention forward pass
+# uses qdotk_max update trick to avoid two loops when computing the softmax
+# see e.g., https://arxiv.org/abs/1805.02867
+# and https://alexdremov.me/understanding-flash-attention-writing-the-algorithm-from-scratch-in-triton/
+def _neighborhood_attention_s2_fwd_torch(kx: torch.Tensor, vx: torch.Tensor, qy: torch.Tensor,
+                                            quad_weights: torch.Tensor, col_idx: torch.Tensor, row_off: torch.Tensor,
+                                            nlon_in: int, nlat_out: int, nlon_out: int) -> torch.Tensor:
+    # prepare result tensor
+    y = torch.zeros_like(qy)
+
+    for ho in range(nlat_out):
+        
+	# get number of nonzeros
+        zstart = row_off[ho]
+        zend = row_off[ho+1]
+
+        for wo in range(nlon_out):
+
+            alpha_sum = torch.zeros((y.shape[0],), dtype=y.dtype, device=y.device)
+            qdotk_max = torch.zeros((y.shape[0],), dtype=y.dtype, device=y.device)
+
+            for idz in range(zstart, zend):
+                nz_col_idx = col_idx[idz]
+
+                # compute input indices from psi datastructure
+                hi = nz_col_idx // nlon_in
+                # account for output shift and ensure positive index due to circular condition
+                wi = nz_col_idx % nlon_in
+                wip = (wi + wo) % nlon_in
+
+                # compute correlation & softmax numerator
+                q_ho_wo = qy[:, :, ho, wo]
+                k_hi_wip = kx[:, :, hi, wip]
+                qdotk = torch.sum(q_ho_wo * k_hi_wip, dim=1)
+
+                # tmp max
+                qdotk_max_tmp = torch.maximum(qdotk_max, qdotk)
+
+                # alpha sum update
+                alpha = torch.exp(qdotk - qdotk_max_tmp) * quad_weights[hi]
+                alpha_sum = alpha + alpha_sum * torch.exp(qdotk_max - qdotk_max_tmp)
+                # update output
+                y[:,:,ho,wo] = y[:,:,ho,wo] * torch.exp(qdotk_max - qdotk_max_tmp).unsqueeze(1) + alpha[:, None] * vx[:,:,hi,wip]
+
+                # define new max
+                qdotk_max = qdotk_max_tmp
+
+            y[:,:,ho,wo] = y[:,:,ho,wo] / alpha_sum[:, None]
+
+    return y
+
+
+# Explicit gradient w.r.t. vx: dM/dv
+# provided as a reference for CUDA & other hand-written gradients
+def _neighborhood_attention_s2_bwd_dv_torch(kx: torch.Tensor, vx: torch.Tensor, qy: torch.Tensor, dy: torch.Tensor,
+                                            quad_weights: torch.Tensor, col_idx: torch.Tensor, row_off: torch.Tensor,
+                                            nlon_in: int, nlat_out: int, nlon_out: int):
+
+    # shapes:
+    # input
+    # kx: B, C, Hi, Wi
+    # vx: B, C, Hi, Wi
+    # qy: B, C, Ho, Wo
+    # quad_weights: Hi
+    # output
+    # dvx: B, C, Hi, Wi
+
+    dvx = torch.zeros_like(vx)
+
+    for ho in range(nlat_out):
+
+        # get number of nonzeros
+        zstart = row_off[ho]
+        zend = row_off[ho+1]
+
+        for wo in range(nlon_out):
+
+            alpha_nz = torch.zeros((dy.shape[0], zend-zstart), dtype=dy.dtype, device=dy.device)
+            qdotk_nz = torch.zeros((dy.shape[0], zend-zstart), dtype=dy.dtype, device=dy.device)
+            alpha_sum = torch.zeros((dy.shape[0],), dtype=dy.dtype, device=dy.device)
+            for idz in range(zstart, zend):
+                nz_col_idx = col_idx[idz]
+
+                # compute input indices from psi datastructure
+                hi = nz_col_idx // nlon_in
+                # account for output shift and ensure positive index due to circular condition
+                wi = nz_col_idx % nlon_in
+                wip = (wi+wo) % nlon_in
+
+                # compute correlation & softmax numerator
+                q_ho_wo = qy[:, :, ho, wo]
+                k_hi_wi = kx[:, :, hi, wip]
+                qdotk_nz[:,idz-zstart] = torch.sum(q_ho_wo * k_hi_wi, dim=1)
+
+            qdotk_max, _ = torch.max(qdotk_nz, dim=1)
+
+            for idz in range(zstart, zend):
+                nz_col_idx = col_idx[idz]
+
+                # compute input indices from psi datastructure
+                hi = nz_col_idx // nlon_in
+                # account for output shift and ensure positive index due to circular condition
+                wi = nz_col_idx % nlon_in
+                wip = (wi+wo) % nlon_in
+                alpha_nz[:,idz-zstart] = torch.exp(qdotk_nz[:,idz-zstart] - qdotk_max) * quad_weights[hi]
+                alpha_sum[:] += alpha_nz[:,idz-zstart]
+
+            for idz in range(zstart, zend):
+                nz_col_idx = col_idx[idz]
+
+                # compute input indices from psi datastructure
+                hi = nz_col_idx // nlon_in
+                # account for output shift and ensure positive index due to circular condition
+                wi = nz_col_idx % nlon_in
+                wip = (wi+wo) % nlon_in
+                dvx[:,:,hi, wip] += (alpha_nz[:, None, idz-zstart] / alpha_sum[:, None]) * dy[:,:,ho,wo]
+
+    return dvx
+
+
+# Explicit gradient w.r.t. kx: dM/dk
+# provided as a reference for CUDA & other hand-written gradients
+def _neighborhood_attention_s2_bwd_dk_torch(kx: torch.Tensor, vx: torch.Tensor, qy: torch.Tensor, dy: torch.Tensor,
+                                            quad_weights: torch.Tensor, col_idx: torch.Tensor, row_off: torch.Tensor,
+                                            nlon_in: int, nlat_out: int, nlon_out: int):
+    # shapes:
+    # input
+    # kx: B, C, Hi, Wi
+    # vx: B, C, Hi, Wi
+    # qy: B, C, Ho, Wo
+    # quad_weights: Hi
+    # output
+    # dkx: B, C, Hi, Wi
+
+    dkx = torch.zeros_like(kx)
+
+    for ho in range(nlat_out):
+
+        # get number of nonzeros
+        zstart = row_off[ho]
+        zend = row_off[ho+1]
+
+        for wo in range(nlon_out):
+
+            qdotk_nz = torch.zeros((dy.shape[0], zend-zstart), dtype=dy.dtype, device=dy.device)
+            integral = torch.zeros((dy.shape[0],), dtype=dy.dtype, device=dy.device)
+            alpha = torch.zeros((dy.shape[0], zend-zstart), dtype=dy.dtype, device=dy.device)
+            alpha_sum = torch.zeros((dy.shape[0],), dtype=dy.dtype, device=dy.device)
+            for idz in range(zstart, zend):
+                nz_col_idx = col_idx[idz]
+
+                # compute input indices from psi datastructure
+                hj = nz_col_idx // nlon_in
+                # account for output shift and ensure positive index due to circular condition
+                wj = nz_col_idx % nlon_in
+                wjp = (wj+wo) % nlon_in
+
+                # compute correlation & softmax numerator
+                q_ho_wo = qy[:, :, ho, wo]
+                k_hj_wjp = kx[:, :, hj, wjp]
+                qdotk_nz[:,idz-zstart] = torch.sum(q_ho_wo * k_hj_wjp, dim=1)
+
+            qdotk_max, _ = torch.max(qdotk_nz, dim=1)
+
+            for idz in range(zstart, zend):
+                nz_col_idx = col_idx[idz]
+
+                # compute input indices from psi datastructure
+                hj = nz_col_idx // nlon_in
+                # account for output shift and ensure positive index due to circular condition
+                wj = nz_col_idx % nlon_in
+                wjp = (wj+wo) % nlon_in
+
+                alpha[:, idz-zstart] = torch.exp(qdotk_nz[:,idz-zstart] - qdotk_max) * quad_weights[hj]
+                alpha_sum[:] += alpha[:, idz-zstart]
+
+                # input dot
+                gdotv = torch.sum(dy[:,:,ho, wo] * vx[:,:,hj, wjp], dim=1)
+
+                # integral term
+                integral[:] += alpha[:, idz-zstart] * gdotv[:]
+
+            integral[:] = integral[:] / alpha_sum[:]
+
+            for idz in range(zstart, zend):
+                nz_col_idx = col_idx[idz]
+
+                # compute input indices from psi datastructure
+                hi = nz_col_idx // nlon_in
+                # account for output shift and ensure positive index due to circular condition
+                wi = nz_col_idx % nlon_in
+                wip = (wi+wo) % nlon_in
+
+                # compute correlation & softmax numerator
+                gdotv = torch.sum(dy[:,:,ho, wo] * vx[:,:,hi, wip], dim=1)
+
+                dkx[:,:,hi,wip] += qy[:, :, ho, wo] * (alpha[:, None, idz-zstart] / alpha_sum[:, None]) * (gdotv[:, None] - integral[:, None])
+
+    return dkx
+
+# Explicit gradient w.r.t. qy: dM/dq
+# provided as a reference for CUDA & other hand-written gradients
+def _neighborhood_attention_s2_bwd_dq_torch(kx: torch.Tensor, vx: torch.Tensor, qy: torch.Tensor, dy: torch.Tensor,
+                                            quad_weights: torch.Tensor, col_idx: torch.Tensor, row_off: torch.Tensor,
+                                            nlon_in: int, nlat_out: int, nlon_out: int):
+
+
+    # shapes:
+    # input
+    # kx: B, C, Hi, Wi
+    # vx: B, C, Hi, Wi
+    # qy: B, C, Ho, Wo
+    # quad_weights: Hi
+    # output
+    # dvx: B, C, Hi, Wi
+
+    dqy = torch.zeros_like(qy)
+
+    for ho in range(nlat_out):
+
+        # get number of nonzeros
+        zstart = row_off[ho]
+        zend = row_off[ho+1]
+
+        for wo in range(nlon_out):
+
+            alpha = torch.zeros((dy.shape[0], zend-zstart), dtype=dy.dtype, device=dy.device)
+            qdotk_nz = torch.zeros((dy.shape[0], zend-zstart), dtype=dy.dtype, device=dy.device)
+            alpha_k = torch.zeros((dy.shape[0], dy.shape[1]), dtype=dy.dtype, device=dy.device)
+            alpha_vw = torch.zeros((dy.shape[0], dy.shape[1]), dtype=dy.dtype, device=dy.device)
+            alpha_kvw = torch.zeros((dy.shape[0], dy.shape[1]), dtype=dy.dtype, device=dy.device)
+            alpha_sum = torch.zeros((dy.shape[0],), dtype=dy.dtype, device=dy.device)
+            alpha_sum2 = torch.zeros((dy.shape[0],), dtype=dy.dtype, device=dy.device)
+            for idz in range(zstart, zend):
+                nz_col_idx = col_idx[idz]
+
+                # compute input indices from psi datastructure
+                hi = nz_col_idx // nlon_in
+                # account for output shift and ensure positive index due to circular condition
+                wi = nz_col_idx % nlon_in
+                wip = (wi+wo) % nlon_in
+
+                idz_i = idz-zstart
+
+                # compute correlation & softmax numerator
+                q_ho_wo = qy[:, :, ho, wo]
+                k_hi_wi = kx[:, :, hi, wip]
+                qdotk_nz[:,idz-zstart] = torch.sum(q_ho_wo * k_hi_wi, dim=1)
+
+            qdotk_max,_ = qdotk_nz.max(dim=1)
+
+            for idz in range(zstart, zend):
+                nz_col_idx = col_idx[idz]
+
+                # compute input indices from psi datastructure
+                hi = nz_col_idx // nlon_in
+                # account for output shift and ensure positive index due to circular condition
+                wi = nz_col_idx % nlon_in
+                wip = (wi+wo) % nlon_in
+
+                q_ho_wo = qy[:, :, ho, wo]
+                k_hi_wi = kx[:, :, hi, wip]
+                idz_i = idz-zstart
+                alpha[:, idz_i] = torch.exp(qdotk_nz[:,idz-zstart] - qdotk_max) * quad_weights[hi]
+                alpha_sum[:] += alpha[:, idz_i]
+
+                gdotv = torch.sum(dy[:,:,ho, wo] * vx[:,:,hi, wip], dim=1)
+                alpha_k[:,:] += alpha[:, None, idz_i] * k_hi_wi
+                alpha_vw[:,:] += alpha[:, None, idz_i] * gdotv[:,None]
+                alpha_kvw[:,:] += alpha[:, None, idz_i] * k_hi_wi * gdotv[:,None]
+
+            dqy[:,:,ho,wo] = (alpha_kvw*alpha_sum[:,None] - alpha_vw*alpha_k) / (alpha_sum[:,None]*alpha_sum[:,None])
+
+    return dqy
+
+class _NeighborhoodAttentionS2(torch.autograd.Function):
+
+    @staticmethod
+    @custom_fwd(device_type="cpu")
+    def forward(ctx, k: torch.Tensor, v: torch.Tensor, q: torch.Tensor,
+                wk: torch.Tensor, wv: torch.Tensor, wq: torch.Tensor,
+                bk: Union[torch.Tensor, None], bv: Union[torch.Tensor, None], bq: Union[torch.Tensor, None],
+                quad_weights: torch.Tensor, col_idx: torch.Tensor, row_off: torch.Tensor,
+                nh: int, nlon_in: int, nlat_out: int, nlon_out: int):
+        
+        ctx.save_for_backward(col_idx, row_off, quad_weights, k, v, q, wk, wv, wq, bk, bv, bq)
+        ctx.nh = nh
+        ctx.nlon_in = nlon_in
+        ctx.nlat_out = nlat_out
+        ctx.nlon_out = nlon_out
+
+        kw = F.conv2d(k, weight=wk, bias=bk)
+        vw = F.conv2d(v, weight=wv, bias=bv)
+        qw = F.conv2d(q, weight=wq, bias=bq)
+
+        # reshape, folding num heads into batch dim
+        B, _, H, W = kw.shape
+        kw = kw.reshape(B*nh, -1, H, W)
+        B, _, H, W = vw.shape
+        vw = vw.reshape(B*nh, -1, H, W)
+        B, _, H, W = qw.shape
+        qw = qw.reshape(B*nh, -1, H, W)
+
+        kw = kw.to(torch.float32)
+        vw = vw.to(torch.float32)
+        qw = qw.to(torch.float32)
+
+        output = _neighborhood_attention_s2_fwd_torch(kw, vw, qw, quad_weights,
+                                                      col_idx, row_off,
+                                                      nlon_in, nlat_out, nlon_out)
+
+        _, C, H, W = output.shape
+        output = output.reshape(B, -1, H, W)
+
+        return output
+
+    @staticmethod
+    @custom_bwd(device_type="cpu")
+    def backward(ctx, grad_output):
+        col_idx, row_off, quad_weights, k, v, q, wk, wv, wq, bk, bv, bq = ctx.saved_tensors
+        nh = ctx.nh
+        nlon_in = ctx.nlon_in
+        nlat_out = ctx.nlat_out
+        nlon_out = ctx.nlon_out
+
+        kw = F.conv2d(k, weight=wk, bias=bk)
+        vw = F.conv2d(v, weight=wv, bias=bv)
+        qw = F.conv2d(q, weight=wq, bias=bq)
+
+        # reshape, folding num heads into batch dim
+        B, _, H, W = kw.shape
+        kw = kw.reshape(B*nh, -1, H, W)
+        B, _, H, W = vw.shape
+        vw = vw.reshape(B*nh, -1, H, W)
+        B, _, H, W = qw.shape
+        qw = qw.reshape(B*nh, -1, H, W)
+        B, _, H, W  = grad_output.shape
+        grad_output = grad_output.reshape(B*nh, -1, H, W)
+
+        dvw = _neighborhood_attention_s2_bwd_dv_torch(kw, vw, qw, grad_output,
+                                                      quad_weights,
+                                                      col_idx, row_off,
+                                                      nlon_in, nlat_out, nlon_out)
+
+        dkw = _neighborhood_attention_s2_bwd_dk_torch(kw, vw, qw, grad_output,
+                                                      quad_weights,
+                                                      col_idx, row_off,
+                                                      nlon_in, nlat_out, nlon_out)
+
+        dqw = _neighborhood_attention_s2_bwd_dq_torch(kw, vw, qw, grad_output,
+                                                      quad_weights,
+                                                      col_idx, row_off,
+                                                      nlon_in, nlat_out, nlon_out)
+
+        # reshape again
+        _, C, H, W = dkw.shape
+        dkw = dkw.reshape(B, -1, H, W)
+        _, C, H, W = dvw.shape
+        dvw = dvw.reshape(B, -1, H, W)
+        _, C, H, W = dqw.shape
+        dqw = dqw.reshape(B, -1, H, W)
+
+        # input grads
+        dv = torch.nn.functional.conv2d(dvw, weight=wv.permute([1,0,2,3]), bias=None)
+        dk = torch.nn.functional.conv2d(dkw, weight=wk.permute([1,0,2,3]), bias=None)
+        dq = torch.nn.functional.conv2d(dqw, weight=wq.permute([1,0,2,3]), bias=None)
+
+        # weight grads
+        dwv = torch.einsum("bchw,bfhw->cf", dvw, v).reshape(*wv.shape).contiguous()
+        dwk = torch.einsum("bchw,bfhw->cf", dkw, k).reshape(*wk.shape).contiguous()
+        dwq = torch.einsum("bchw,bfhw->cf", dqw, q).reshape(*wq.shape).contiguous()
+
+        # bias grads:
+        if bv is not None:
+            dbv = torch.sum(dvw, dim=(0,2,3))
+        else:
+            dbv = None
+
+        if bk is not None:
+            dbk = torch.sum(dkw, dim=(0,2,3))
+        else:
+            dbk = None
+
+        if bq is not None:
+            dbq = torch.sum(dqw, dim=(0,2,3))
+        else:
+            dbq = None
+
+        return dk, dv, dq, dwk, dwv, dwq, dbk, dbv, dbq, \
+                None, None, None, None, None, None, None
+
+
+def _neighborhood_attention_s2_torch(k: torch.Tensor, v: torch.Tensor, q: torch.Tensor,
+                                     wk: torch.Tensor, wv: torch.Tensor, wq: torch.Tensor,
+                                     bk: Union[torch.Tensor, None], bv: Union[torch.Tensor, None],
+                                     bq: Union[torch.Tensor, None], quad_weights: torch.Tensor,
+                                     col_idx: torch.Tensor, row_off: torch.Tensor,
+                                     nh: int, nlon_in: int, nlat_out: int, nlon_out: int) -> torch.Tensor:
+   
+    return _NeighborhoodAttentionS2.apply(k, v, q, wk, wv, wq, bk, bv, bq,
+                                          quad_weights, col_idx, row_off,
+                                          nh, nlon_in, nlat_out, nlon_out)
+
+
+class _NeighborhoodAttentionS2Cuda(torch.autograd.Function):
+
+
+    @staticmethod
+    @custom_fwd(device_type="cuda")
+    def forward(ctx, k: torch.Tensor, v: torch.Tensor, q: torch.Tensor,
+                wk: torch.Tensor, wv: torch.Tensor, wq: torch.Tensor,
+                bk: Union[torch.Tensor, None], bv: Union[torch.Tensor, None], bq: Union[torch.Tensor, None],
+                quad_weights: torch.Tensor, col_idx: torch.Tensor, row_off: torch.Tensor,
+                max_psi_nnz: int, nh: int, nlon_in: int, nlat_out: int, nlon_out: int):
+        
+        ctx.save_for_backward(col_idx, row_off, quad_weights, k, v, q, wk, wv, wq, bk, bv, bq)
+        ctx.nh = nh
+        ctx.max_psi_nnz = max_psi_nnz
+        ctx.nlon_in = nlon_in
+        ctx.nlat_out = nlat_out
+        ctx.nlon_out = nlon_out
+
+        kw = F.conv2d(k, weight=wk, bias=bk)
+        vw = F.conv2d(v, weight=wv, bias=bv)
+        qw = F.conv2d(q, weight=wq, bias=bq)
+
+        # reshape, folding num heads into batch dim
+        B, _, H, W = kw.shape
+        kw = kw.reshape(B*nh, -1, H, W)
+        B, _, H, W = vw.shape
+        vw = vw.reshape(B*nh, -1, H, W)
+        B, _, H, W = qw.shape
+        qw = qw.reshape(B*nh, -1, H, W)
+
+        # convert to float32
+        inp_dtype = kw.dtype
+        kw = kw.to(torch.float32).contiguous()
+        vw = vw.to(torch.float32).contiguous()
+        qw = qw.to(torch.float32).contiguous()
+
+        output = attention_cuda_extension.forward(kw, vw, qw, quad_weights,
+                                                  col_idx, row_off,
+                                                  nlon_in, nlat_out, nlon_out)
+
+        _, C, H, W = output.shape
+        output = output.reshape(B, -1, H, W)
+
+        # convert back precision
+        output = output.to(dtype=inp_dtype)
+
+        return output
+
+    @staticmethod
+    @custom_bwd(device_type="cuda")
+    def backward(ctx, grad_output):
+        col_idx, row_off, quad_weights, k, v, q, wk, wv, wq, bk, bv, bq = ctx.saved_tensors
+        nh = ctx.nh
+        max_psi_nnz = ctx.max_psi_nnz
+        nlon_in = ctx.nlon_in
+        nlat_out = ctx.nlat_out
+        nlon_out = ctx.nlon_out
+
+        kw = F.conv2d(k, weight=wk, bias=bk)
+        vw = F.conv2d(v, weight=wv, bias=bv)
+        qw = F.conv2d(q, weight=wq, bias=bq)
+
+        # reshape, folding num heads into batch dim
+        B, _, H, W = kw.shape
+        kw = kw.reshape(B*nh, -1, H, W)
+        B, _, H, W = vw.shape
+        vw = vw.reshape(B*nh, -1, H, W)
+        B, _, H, W = qw.shape
+        qw = qw.reshape(B*nh, -1, H, W)
+        B, _, H, W  = grad_output.shape
+        grad_output = grad_output.reshape(B*nh, -1, H, W)
+
+        # save type and convert to float32
+        kw_dtype = kw.dtype
+        vw_dtype = vw.dtype
+        qw_dtype = qw.dtype
+
+        kw = kw.to(torch.float32).contiguous()
+        vw = vw.to(torch.float32).contiguous()
+        qw = qw.to(torch.float32).contiguous()
+        grad_output = grad_output.to(torch.float32).contiguous()
+
+        dkw,dvw,dqw = attention_cuda_extension.backward_dkvq(kw, vw, qw, grad_output,
+                                                             quad_weights,
+                                                             col_idx, row_off,
+                                                             nlon_in, nlat_out, nlon_out)
+
+        # reshape again
+        _, C, H, W = dkw.shape
+        dkw = dkw.reshape(B, -1, H, W)
+        _, C, H, W = dvw.shape
+        dvw = dvw.reshape(B, -1, H, W)
+        _, C, H, W = dqw.shape
+        dqw = dqw.reshape(B, -1, H, W)
+
+        # convert precision
+        dkw = dkw.to(dtype=kw_dtype)
+        dvw = dvw.to(dtype=vw_dtype)
+        dqw = dqw.to(dtype=qw_dtype)
+
+        # input grads
+        dv = torch.nn.functional.conv2d(dvw, weight=wv.permute([1,0,2,3]), bias=None)
+        dk = torch.nn.functional.conv2d(dkw, weight=wk.permute([1,0,2,3]), bias=None)
+        dq = torch.nn.functional.conv2d(dqw, weight=wq.permute([1,0,2,3]), bias=None)
+
+        # weight grads
+        dwv = torch.einsum("bchw,bfhw->cf", dvw, v).reshape(*wv.shape).contiguous()
+        dwk = torch.einsum("bchw,bfhw->cf", dkw, k).reshape(*wk.shape).contiguous()
+        dwq = torch.einsum("bchw,bfhw->cf", dqw, q).reshape(*wq.shape).contiguous()
+
+        # bias grads:
+        if bv is not None:
+            dbv = torch.sum(dvw, dim=(0,2,3))
+        else:
+            dbv = None
+
+        if bk is not None:
+            dbk = torch.sum(dkw, dim=(0,2,3))
+        else:
+            dbk = None
+
+        if bq is not None:
+            dbq = torch.sum(dqw, dim=(0,2,3))
+        else:
+            dbq = None
+
+        return dk, dv, dq, dwk, dwv, dwq, dbk, dbv, dbq, \
+                None, None, None, None, None, None, None, None
+
+
+def _neighborhood_attention_s2_cuda(k: torch.Tensor, v: torch.Tensor, q: torch.Tensor,
+                                    wk: torch.Tensor, wv: torch.Tensor, wq: torch.Tensor,
+                                    bk: Union[torch.Tensor, None], bv: Union[torch.Tensor, None],
+                                    bq: Union[torch.Tensor, None], quad_weights: torch.Tensor,
+                                    col_idx: torch.Tensor, row_off: torch.Tensor, max_psi_nnz: int,
+                                    nh: int, nlon_in: int, nlat_out: int, nlon_out: int) -> torch.Tensor:
+   
+    return _NeighborhoodAttentionS2Cuda.apply(k, v, q, wk, wv, wq, bk, bv, bq,
+                                              quad_weights, col_idx, row_off, max_psi_nnz,
+                                              nh, nlon_in, nlat_out, nlon_out)

.deps/torch_harmonics/attention.py

@@ -0,0 +1,382 @@
+# coding=utf-8
+
+# SPDX-FileCopyrightText: Copyright (c) 2025 The torch-harmonics Authors. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+#
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+#
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+#
+
+from typing import List, Tuple, Union, Optional
+from warnings import warn
+
+import math
+
+import torch
+import torch.nn as nn
+import numpy as np
+
+from torch_harmonics.quadrature import _precompute_latitudes
+from torch_harmonics.convolution import _precompute_convolution_tensor_s2
+from torch_harmonics._neighborhood_attention import _neighborhood_attention_s2_torch, _neighborhood_attention_s2_cuda
+from torch_harmonics.filter_basis import get_filter_basis
+
+# import custom C++/CUDA extensions
+try:
+    import attention_cuda_extension
+
+    _cuda_extension_available = True
+except ImportError as err:
+    attention_cuda_extension = None
+    _cuda_extension_available = False
+
+class AttentionS2(nn.Module):
+    """
+    (Global) attention on the 2-sphere.
+    Parameters
+    -----------
+    in_channels: int
+        number of channels of the input signal (corresponds to embed_dim in MHA in PyTorch)
+    num_heads: int
+        number of attention heads
+    in_shape: tuple
+        shape of the input grid
+    out_shape: tuple
+        shape of the output grid
+    grid_in: str, optional
+        input grid type, "equiangular" by default
+    grid_out: str, optional
+        output grid type, "equiangular" by default
+    bias: bool, optional
+        if specified, adds bias to input / output projection layers
+    k_channels: int
+        number of dimensions for interior inner product in the attention matrix (corresponds to kdim in MHA in PyTorch)
+    out_channels: int, optional
+        number of dimensions for interior inner product in the attention matrix (corresponds to vdim in MHA in PyTorch)
+    """
+
+    def __init__(
+            self,
+            in_channels: int,
+            num_heads: int,
+            in_shape: Tuple[int],
+            out_shape: Tuple[int],
+            grid_in: Optional[str] = "equiangular",
+            grid_out: Optional[str] = "equiangular",
+            scale: Optional[Union[torch.Tensor, float]] = None,
+            bias: Optional[bool] = True,
+            k_channels: Optional[int] = None,
+            out_channels: Optional[int] = None,
+            drop_rate: Optional[float]=0.0,
+    ):
+        super().__init__()
+
+        self.nlat_in, self.nlon_in = in_shape
+        self.nlat_out, self.nlon_out = out_shape
+
+        self.in_channels = in_channels
+        self.num_heads = num_heads
+        self.k_channels = in_channels if k_channels is None else k_channels
+        self.out_channels = in_channels if out_channels is None else out_channels
+        self.drop_rate = drop_rate
+        self.scale = scale
+
+        # integration weights
+        _, wgl = _precompute_latitudes(self.nlat_in, grid=grid_in)
+        quad_weights = 2.0 * torch.pi * wgl.to(dtype=torch.float32) / self.nlon_in
+        # we need to tile and flatten them accordingly
+        quad_weights = torch.tile(quad_weights.reshape(-1, 1), (1, self.nlon_in)).flatten()
+
+        # compute log because they are applied as an addition prior to the softmax ('attn_mask'), which includes an exponential.
+        # see https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html
+        # for info on how 'attn_mask' is applied to the attention weights
+        log_quad_weights = torch.log(quad_weights).reshape(1,1,-1)
+        self.register_buffer("log_quad_weights", log_quad_weights, persistent=False)
+
+        # learnable parameters
+        # TODO: double-check that this gives us the correct initialization magnitudes
+        # the standard MHA uses xavier uniform, NATTEN uses kaiming. Let's use that for now
+        if self.k_channels % self.num_heads != 0:
+            raise ValueError(f"Please make sure that number of heads {self.num_heads} divides k_channels {self.k_channels} evenly.")
+        if self.out_channels % self.num_heads != 0:
+            raise ValueError(f"Please make sure that number of heads {self.num_heads} divides out_channels {self.out_channels} evenly.")
+        scale_qkv = math.sqrt(3.0 / self.in_channels)
+        self.q_weights = nn.Parameter(scale_qkv * (2 * torch.rand(self.k_channels, self.in_channels, 1, 1) - 1))
+        self.k_weights = nn.Parameter(scale_qkv * (2 * torch.rand(self.k_channels, self.in_channels, 1, 1) - 1))
+        self.v_weights = nn.Parameter(scale_qkv * (2 * torch.rand(self.out_channels, self.in_channels, 1, 1) - 1))
+        scale_proj = math.sqrt(3.0 / self.out_channels)
+        self.proj_weights = nn.Parameter(scale_proj * (2 * torch.rand(self.out_channels, self.out_channels, 1, 1) - 1))
+
+        if bias:
+            self.q_bias = nn.Parameter(torch.zeros(self.k_channels))
+            self.k_bias = nn.Parameter(torch.zeros(self.k_channels))
+            self.v_bias = nn.Parameter(torch.zeros(self.out_channels))
+            self.proj_bias = nn.Parameter(torch.zeros(self.out_channels))
+        else:
+            self.q_bias = None
+            self.k_bias = None
+            self.v_bias = None
+            self.proj_bias = None
+
+
+    def extra_repr(self):
+        return f"in_shape={(self.nlat_in, self.nlon_in)}, out_shape={(self.nlat_out, self.nlon_out)}, in_channels={self.in_channels}, out_channels={self.out_channels}, k_channels={self.k_channels}"
+
+    def forward(self, query: torch.Tensor, key: Optional[torch.Tensor] = None, value: Optional[torch.Tensor] = None) -> torch.Tensor:
+
+        # self attention simplification
+        if key is None:
+            key = query
+
+        if value is None:
+            value = query
+
+        # change this later to allow arbitrary number of batch dims
+        assert (query.dim() == key.dim()) and (key.dim() == value.dim()) and (value.dim() == 4)
+
+        # perform MLP
+        query = nn.functional.conv2d(query, self.q_weights, bias=self.q_bias)
+        key = nn.functional.conv2d(key, self.k_weights, bias=self.k_bias)
+        value = nn.functional.conv2d(value, self.v_weights, bias=self.v_bias)
+
+        # reshape
+        B, _, H, W = query.shape
+        query = query.reshape(B, self.num_heads, -1, H, W)
+        B, _, H, W = key.shape
+        key = key.reshape(B, self.num_heads, -1, H, W)
+        B, _, H, W = value.shape
+        value = value.reshape(B, self.num_heads, -1, H, W)
+
+        # reshape to the right dimensions
+        B, _, C, H, W = query.shape
+        query = query.permute(0,1,3,4,2).reshape(B, self.num_heads, H*W, C)
+        B, _, C, H, W = key.shape
+        key = key.permute(0,1,3,4,2).reshape(B, self.num_heads, H*W, C)
+        B, _, C, H, W = value.shape
+        value = value.permute(0,1,3,4,2).reshape(B, self.num_heads, H*W, C)
+
+        # multiply the query, key and value tensors
+        out = nn.functional.scaled_dot_product_attention(query, key, value, attn_mask=self.log_quad_weights, dropout_p=self.drop_rate, scale=self.scale)
+
+        # reshape
+        B, _, _, C = out.shape
+        # (B, heads, H*W, C)
+        out = out.permute(0,1,3,2)
+        # (B, heads, C, H*W)
+        out = out.reshape(B, self.num_heads*C, self.nlat_out, self.nlon_out)
+        # (B, heads*C, H, W)
+        out = nn.functional.conv2d(out, self.proj_weights, bias=self.proj_bias)
+
+        return out
+
+
+class NeighborhoodAttentionS2(nn.Module):
+    """
+    Neighborhood attention on the 2-sphere.
+
+    Parameters
+    -----------
+    in_channels: int
+        number of channels of the input signal (corresponds to embed_dim in MHA in PyTorch)
+    in_shape: tuple
+        shape of the input grid
+    out_shape: tuple
+        shape of the output grid
+    grid_in: str, optional
+        input grid type, "equiangular" by default
+    grid_out: str, optional
+        output grid type, "equiangular" by default
+    bias: bool, optional
+        if specified, adds bias to input / output projection layers
+    theta_cutoff: float, optional
+        neighborhood size
+    k_channels: int
+        number of dimensions for interior inner product in the attention matrix (corresponds to kdim in MHA in PyTorch)
+    out_channels: int, optional
+        number of dimensions for interior inner product in the attention matrix (corresponds to vdim in MHA in PyTorch)
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        in_shape: Tuple[int],
+        out_shape: Tuple[int],
+        grid_in: Optional[str] = "equiangular",
+        grid_out: Optional[str] = "equiangular",
+        num_heads: Optional[int] = 1,
+        scale: Optional[Union[torch.Tensor, float]] = None,
+        bias: Optional[bool] = True,
+        theta_cutoff: Optional[float] = None,
+        k_channels: Optional[int] = None,
+        out_channels: Optional[int] = None,
+    ):
+        super().__init__()
+
+        self.nlat_in, self.nlon_in = in_shape
+        self.nlat_out, self.nlon_out = out_shape
+
+        self.in_channels = in_channels
+        self.num_heads = num_heads
+        self.k_channels = in_channels if k_channels is None else k_channels
+        self.out_channels = in_channels if out_channels is None else out_channels
+
+        # heuristic to compute theta cutoff based on the bandlimit of the input field and overlaps of the basis functions
+        if theta_cutoff is None:
+            theta_cutoff = torch.pi / float(self.nlat_out - 1)
+
+        if theta_cutoff <= 0.0:
+            raise ValueError("Error, theta_cutoff has to be positive.")
+
+        # integration weights
+        _, wgl = _precompute_latitudes(self.nlat_in, grid=grid_in)
+        quad_weights = 2.0 * torch.pi * wgl.to(dtype=torch.float32) / self.nlon_in
+        self.register_buffer("quad_weights", quad_weights, persistent=False)
+
+        # create a dummy filter basis to pass to the construction of the convolution tensor
+        # this is to avoid code duplication as the logic of pre-computing the sparsity pattern
+        # is identical to convolutions with a constant filter function
+        fb = get_filter_basis(kernel_shape=1, basis_type="zernike")
+
+        # precompute the neighborhood sparsity pattern
+        idx, _, roff = _precompute_convolution_tensor_s2(
+            in_shape,
+            out_shape,
+            fb,
+            grid_in=grid_in,
+            grid_out=grid_out,
+            theta_cutoff=theta_cutoff,
+            transpose_normalization=False,
+            basis_norm_mode="none",
+            merge_quadrature=True,
+        )
+
+        # this is kept for legacy resons in case we want to resuse sorting of these entries
+        row_idx = idx[1, ...].contiguous()
+        col_idx = idx[2, ...].contiguous()
+        roff_idx = roff.contiguous()
+
+        # store some metadata
+        self.max_psi_nnz = col_idx.max().item() + 1
+        self.register_buffer("psi_row_idx", row_idx, persistent=False)
+        self.register_buffer("psi_col_idx", col_idx, persistent=False)
+        self.register_buffer("psi_roff_idx", roff_idx, persistent=False)
+
+        # learnable parameters
+        # TODO: double-check that this gives us the correct initialization magnitudes
+        # the standard MHA uses xavier uniform, NATTEN uses kaiming. Let's use that for now
+        if self.k_channels % self.num_heads != 0:
+            raise ValueError(f"Please make sure that number of heads {self.num_heads} divides k_channels {self.k_channels} evenly.")
+        if self.out_channels % self.num_heads != 0:
+            raise ValueError(f"Please make sure that number of heads {self.num_heads} divides out_channels {self.out_channels} evenly.")
+        scale_qkv = math.sqrt(3.0 / self.in_channels)
+        self.q_weights = nn.Parameter(scale_qkv * (2 * torch.rand(self.k_channels, self.in_channels, 1, 1) - 1))
+        self.k_weights = nn.Parameter(scale_qkv * (2 * torch.rand(self.k_channels, self.in_channels, 1, 1) - 1))
+        self.v_weights = nn.Parameter(scale_qkv * (2 * torch.rand(self.out_channels, self.in_channels, 1, 1) - 1))
+        scale_proj = math.sqrt(3.0 / self.out_channels)
+        self.proj_weights = nn.Parameter(scale_proj * (2 * torch.rand(self.out_channels, self.out_channels, 1, 1) - 1))
+
+        if scale is not None:
+            self.scale = scale
+        else:
+            self.scale = 1 / math.sqrt(self.k_channels)
+
+        if bias:
+            self.q_bias = nn.Parameter(torch.zeros(self.k_channels))
+            self.k_bias = nn.Parameter(torch.zeros(self.k_channels))
+            self.v_bias = nn.Parameter(torch.zeros(self.out_channels))
+            self.proj_bias = nn.Parameter(torch.zeros(self.out_channels))
+        else:
+            self.q_bias = None
+            self.k_bias = None
+            self.v_bias = None
+            self.proj_bias = None
+
+    def extra_repr(self):
+        return f"in_shape={(self.nlat_in, self.nlon_in)}, out_shape={(self.nlat_out, self.nlon_out)}, in_channels={self.in_channels}, out_channels={self.out_channels}, k_channels={self.k_channels}"
+
+    def forward(self, query: torch.Tensor, key: Optional[torch.Tensor] = None, value: Optional[torch.Tensor] = None) -> torch.Tensor:
+
+        # self attention simplification
+        if key is None:
+            key = query
+
+        if value is None:
+            value = query
+
+        # change this later to allow arbitrary number of batch dims
+        assert (query.dim() == key.dim()) and (key.dim() == value.dim()) and (value.dim() == 4)
+
+        # do the scaling
+        query_scaled = query * self.scale
+
+        # TODO: insert dimension checks for input
+        if query.is_cuda and _cuda_extension_available:
+
+            out = _neighborhood_attention_s2_cuda(
+                key,
+                value,
+                query_scaled,
+                self.k_weights,
+                self.v_weights,
+                self.q_weights,
+                self.k_bias,
+                self.v_bias,
+                self.q_bias,
+                self.quad_weights,
+                self.psi_col_idx,
+                self.psi_roff_idx,
+                self.max_psi_nnz,
+                self.num_heads,
+                self.nlon_in,
+                self.nlat_out,
+                self.nlon_out,
+            )
+        else:
+            if query.is_cuda:
+                warn("couldn't find CUDA extension, falling back to slow PyTorch implementation")
+
+            # call attention
+            out = _neighborhood_attention_s2_torch(
+                key,
+                value,
+                query_scaled,
+                self.k_weights,
+                self.v_weights,
+                self.q_weights,
+                self.k_bias,
+                self.v_bias,
+                self.q_bias,
+                self.quad_weights,
+                self.psi_col_idx,
+                self.psi_roff_idx,
+                self.num_heads,
+                self.nlon_in,
+                self.nlat_out,
+                self.nlon_out,
+            )
+
+        out = nn.functional.conv2d(out, self.proj_weights, bias=self.proj_bias)
+
+        return out

.deps/torch_harmonics/cache.py

@@ -0,0 +1,74 @@
+# coding=utf-8
+
+# SPDX-FileCopyrightText: Copyright (c) 2025 The torch-harmonics Authors. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+#
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+#
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+#
+
+import functools
+from copy import deepcopy
+
+# copying LRU cache decorator a la:
+# https://stackoverflow.com/questions/54909357/how-to-get-functools-lru-cache-to-return-new-instances
+def lru_cache(maxsize=20, typed=False, copy=False):
+    """
+    Least Recently Used (LRU) cache decorator with optional deep copying.
+    
+    This is a wrapper around functools.lru_cache that adds the ability to return
+    deep copies of cached results to prevent unintended modifications to cached objects.
+    
+    Parameters
+    -----------
+    maxsize : int, optional
+        Maximum number of items to cache, by default 20
+    typed : bool, optional
+        Whether to cache different types separately, by default False
+    copy : bool, optional
+        Whether to return deep copies of cached results, by default False
+        
+    Returns
+    -------
+    function
+        Decorated function with LRU caching
+        
+    Example
+    -------
+    >>> @lru_cache(maxsize=10, copy=True)
+    ... def expensive_function(x):
+    ...     return [x, x*2, x*3]
+    """
+    def decorator(f):
+        cached_func = functools.lru_cache(maxsize=maxsize, typed=typed)(f)
+        def wrapper(*args, **kwargs):
+            res = cached_func(*args, **kwargs)
+            if copy:
+                return deepcopy(res)
+            else:
+                return res
+                
+        return wrapper
+    return decorator

.deps/torch_harmonics/convolution.py

@@ -0,0 +1,666 @@
+# coding=utf-8
+
+# SPDX-FileCopyrightText: Copyright (c) 2022 The torch-harmonics Authors. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+#
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+#
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+#
+
+import abc
+from typing import List, Tuple, Union, Optional
+from warnings import warn
+
+import math
+
+import torch
+import torch.nn as nn
+
+from functools import partial
+
+from torch_harmonics.cache import lru_cache
+from torch_harmonics.quadrature import _precompute_grid, _precompute_latitudes, _precompute_longitudes
+from torch_harmonics._disco_convolution import _get_psi, _disco_s2_contraction_torch, _disco_s2_transpose_contraction_torch
+from torch_harmonics._disco_convolution import _disco_s2_contraction_cuda, _disco_s2_transpose_contraction_cuda
+from torch_harmonics.filter_basis import FilterBasis, get_filter_basis
+
+# import custom C++/CUDA extensions if available
+try:
+    from disco_helpers import preprocess_psi
+    import disco_cuda_extension
+
+    _cuda_extension_available = True
+except ImportError as err:
+    disco_cuda_extension = None
+    _cuda_extension_available = False
+
+
+def _normalize_convolution_tensor_s2(
+    psi_idx, psi_vals, in_shape, out_shape, kernel_size, quad_weights, transpose_normalization=False, basis_norm_mode="mean", merge_quadrature=False, eps=1e-9
+):
+    """Normalizes convolution tensor values based on specified normalization mode.
+
+    This function applies different normalization strategies to the convolution tensor
+    values based on the basis_norm_mode parameter. It can normalize individual basis
+    functions, compute mean normalization across all basis functions, or use support
+    weights. The function also optionally merges quadrature weights into the tensor.
+
+    Parameters
+    -----------
+    psi_idx: torch.Tensor
+        Index tensor for the sparse convolution tensor.
+    psi_vals: torch.Tensor
+        Value tensor for the sparse convolution tensor.
+    in_shape: Tuple[int]
+        Tuple of (nlat_in, nlon_in) representing input grid dimensions.
+    out_shape: Tuple[int]
+        Tuple of (nlat_out, nlon_out) representing output grid dimensions.
+    kernel_size: int
+        Number of kernel basis functions.
+    quad_weights: torch.Tensor
+        Quadrature weights for numerical integration.
+    transpose_normalization: bool
+        If True, applies normalization in transpose direction.
+    basis_norm_mode: str
+        Normalization mode, one of ["none", "individual", "mean", "support"].
+    merge_quadrature: bool
+        If True, multiplies values by quadrature weights.
+    eps: float
+        Small epsilon value to prevent division by zero.
+
+    Returns
+    -------
+    torch.Tensor
+        Normalized convolution tensor values.
+
+    Raises
+    ------
+    ValueError
+        If basis_norm_mode is not one of the supported modes.
+    """
+
+    # exit here if no normalization is needed
+    if basis_norm_mode == "none":
+        return psi_vals
+
+    # reshape the indices implicitly to be ikernel, out_shape[0], in_shape[0], in_shape[1]
+    idx = torch.stack([psi_idx[0], psi_idx[1], psi_idx[2] // in_shape[1], psi_idx[2] % in_shape[1]], dim=0)
+
+    # getting indices for adressing kernels, input and output latitudes
+    ikernel = idx[0]
+
+    if transpose_normalization:
+        ilat_out = idx[2]
+        ilat_in = idx[1]
+        # here we are deliberately swapping input and output shapes to handle transpose normalization with the same code
+        nlat_out = in_shape[0]
+        correction_factor = out_shape[1] / in_shape[1]
+    else:
+        ilat_out = idx[1]
+        ilat_in = idx[2]
+        nlat_out = out_shape[0]
+
+    # get the quadrature weights
+    q = quad_weights[ilat_in].reshape(-1)
+
+    # buffer to store intermediate values
+    vnorm = torch.zeros(kernel_size, nlat_out, device=psi_vals.device)
+    support = torch.zeros(kernel_size, nlat_out, device=psi_vals.device)
+
+    # loop through dimensions to compute the norms
+    for ik in range(kernel_size):
+        for ilat in range(nlat_out):
+
+            # find indices corresponding to the given output latitude and kernel basis function
+            iidx = torch.argwhere((ikernel == ik) & (ilat_out == ilat))
+
+            # compute the 1-norm
+            # vnorm[ik, ilat] = torch.sqrt(torch.sum(psi_vals[iidx].abs().pow(2) * q[iidx]))
+            vnorm[ik, ilat] = torch.sum(psi_vals[iidx].abs() * q[iidx])
+
+            # compute the support
+            support[ik, ilat] = torch.sum(q[iidx])
+
+    # loop over values and renormalize
+    for ik in range(kernel_size):
+        for ilat in range(nlat_out):
+
+            iidx = torch.argwhere((ikernel == ik) & (ilat_out == ilat))
+
+            if basis_norm_mode == "individual":
+                val = vnorm[ik, ilat]
+            elif basis_norm_mode == "mean":
+                val = vnorm[ik, :].mean()
+            elif basis_norm_mode == "support":
+                val = support[ik, ilat]
+            elif basis_norm_mode == "none":
+                val = 1.0
+            else:
+                raise ValueError(f"Unknown basis normalization mode {basis_norm_mode}.")
+
+            psi_vals[iidx] = psi_vals[iidx] / (val + eps)
+
+            if merge_quadrature:
+                psi_vals[iidx] = psi_vals[iidx] * q[iidx]
+
+    if transpose_normalization and merge_quadrature:
+        psi_vals = psi_vals / correction_factor
+
+    return psi_vals
+
+
+@lru_cache(typed=True, copy=True)
+def _precompute_convolution_tensor_s2(
+    in_shape: Tuple[int],
+    out_shape: Tuple[int],
+    filter_basis: FilterBasis,
+    grid_in: Optional[str] = "equiangular",
+    grid_out: Optional[str] = "equiangular",
+    theta_cutoff: Optional[float] = 0.01 * math.pi,
+    theta_eps: Optional[float] = 1e-3,
+    transpose_normalization: Optional[bool] = False,
+    basis_norm_mode: Optional[str] = "mean",
+    merge_quadrature: Optional[bool] = False,
+):
+    """
+    Precomputes the rotated filters at positions $R^{-1}_j \omega_i = R^{-1}_j R_i \nu = Y(-\theta_j)Z(\phi_i - \phi_j)Y(\theta_j)\nu$.
+    Assumes a tensorized grid on the sphere with an equidistant sampling in longitude as described in Ocampo et al.
+    The output tensor has shape kernel_shape x nlat_out x (nlat_in * nlon_in).
+
+    The rotation of the Euler angles uses the YZY convention, which applied to the northpole $(0,0,1)^T$ yields
+    $$
+    Y(\alpha) Z(\beta) Y(\gamma) n =
+        {\begin{bmatrix}
+            \cos(\gamma)\sin(\alpha) + \cos(\alpha)\cos(\beta)\sin(\gamma) \\
+            \sin(\beta)\sin(\gamma) \\
+            \cos(\alpha)\cos(\gamma)-\cos(\beta)\sin(\alpha)\sin(\gamma)
+        \end{bmatrix}}
+    $$
+
+    Parameters
+    -----------
+    in_shape: Tuple[int]
+        Input shape of the convolution tensor
+    out_shape: Tuple[int]
+        Output shape of the convolution tensor
+    filter_basis: FilterBasis
+        Filter basis functions
+    grid_in: str
+        Input grid type
+    grid_out: str
+        Output grid type
+    theta_cutoff: float
+        Theta cutoff for the filter basis functions
+    theta_eps: float
+        Epsilon for the theta cutoff
+    transpose_normalization: bool
+        Whether to normalize the convolution tensor in the transpose direction
+    basis_norm_mode: str
+        Mode for basis normalization
+    merge_quadrature: bool
+        Whether to merge the quadrature weights into the convolution tensor
+
+    Returns
+    -------
+    out_idx: torch.Tensor
+        Index tensor of the convolution tensor
+    out_vals: torch.Tensor
+        Values tensor of the convolution tensor
+
+    """
+
+    assert len(in_shape) == 2
+    assert len(out_shape) == 2
+
+    kernel_size = filter_basis.kernel_size
+
+    nlat_in, nlon_in = in_shape
+    nlat_out, nlon_out = out_shape
+
+    # precompute input and output grids
+    lats_in, win = _precompute_latitudes(nlat_in, grid=grid_in)
+    lats_out, wout = _precompute_latitudes(nlat_out, grid=grid_out)
+
+    # compute the phi differences
+    # It's imporatant to not include the 2 pi point in the longitudes, as it is equivalent to lon=0
+    lons_in = _precompute_longitudes(nlon_in)
+
+    # compute quadrature weights and merge them into the convolution tensor.
+    # These quadrature integrate to 1 over the sphere.
+    if transpose_normalization:
+        quad_weights = wout.reshape(-1, 1) / nlon_in / 2.0
+    else:
+        quad_weights = win.reshape(-1, 1) / nlon_in / 2.0
+
+    # effective theta cutoff if multiplied with a fudge factor to avoid aliasing with grid width (especially near poles)
+    theta_cutoff_eff = (1.0 + theta_eps) * theta_cutoff
+
+    out_idx = []
+    out_vals = []
+
+    beta = lons_in
+    gamma = lats_in.reshape(-1, 1)
+
+    # compute trigs
+    cbeta = torch.cos(beta)
+    sbeta = torch.sin(beta)
+    cgamma = torch.cos(gamma)
+    sgamma = torch.sin(gamma)
+
+    # compute row offsets
+    out_roff = torch.zeros(nlat_out + 1, dtype=torch.int64, device=lons_in.device)
+    out_roff[0] = 0
+    for t in range(nlat_out):
+        # the last angle has a negative sign as it is a passive rotation, which rotates the filter around the y-axis
+        alpha = -lats_out[t]
+
+        # compute cartesian coordinates of the rotated position
+        # This uses the YZY convention of Euler angles, where the last angle (alpha) is a passive rotation,
+        # and therefore applied with a negative sign
+        x = torch.cos(alpha) * cbeta * sgamma + cgamma * torch.sin(alpha)
+        y = sbeta * sgamma
+        z = -cbeta * torch.sin(alpha) * sgamma + torch.cos(alpha) * cgamma
+
+        # normalization is important to avoid NaNs when arccos and atan are applied
+        # this can otherwise lead to spurious artifacts in the solution
+        norm = torch.sqrt(x * x + y * y + z * z)
+        x = x / norm
+        y = y / norm
+        z = z / norm
+
+        # compute spherical coordinates, where phi needs to fall into the [0, 2pi) range
+        theta = torch.arccos(z)
+        phi = torch.arctan2(y, x)
+        phi = torch.where(phi < 0.0, phi + 2 * torch.pi, phi)
+
+        # find the indices where the rotated position falls into the support of the kernel
+        iidx, vals = filter_basis.compute_support_vals(theta, phi, r_cutoff=theta_cutoff_eff)
+
+        # add the output latitude and reshape such that psi has dimensions kernel_shape x nlat_out x (nlat_in*nlon_in)
+        idx = torch.stack([iidx[:, 0], t * torch.ones_like(iidx[:, 0]), iidx[:, 1] * nlon_in + iidx[:, 2]], dim=0)
+
+        # append indices and values to the COO datastructure, compute row offsets
+        out_idx.append(idx)
+        out_vals.append(vals)
+        out_roff[t + 1] = out_roff[t] + iidx.shape[0]
+
+    # concatenate the indices and values
+    out_idx = torch.cat(out_idx, dim=-1)
+    out_vals = torch.cat(out_vals, dim=-1)
+
+    out_vals = _normalize_convolution_tensor_s2(
+        out_idx,
+        out_vals,
+        in_shape,
+        out_shape,
+        kernel_size,
+        quad_weights,
+        transpose_normalization=transpose_normalization,
+        basis_norm_mode=basis_norm_mode,
+        merge_quadrature=merge_quadrature,
+    )
+
+    out_idx = out_idx.contiguous()
+    out_vals = out_vals.to(dtype=torch.float32).contiguous()
+
+    return out_idx, out_vals, out_roff
+
+
+class DiscreteContinuousConv(nn.Module, metaclass=abc.ABCMeta):
+    """
+    Abstract base class for discrete-continuous convolutions
+
+    Parameters
+    -----------
+    in_channels: int
+        Number of input channels
+    out_channels: int
+        Number of output channels
+    kernel_shape: Union[int, Tuple[int], Tuple[int, int]]
+        Shape of the kernel
+    basis_type: Optional[str]
+        Type of the basis functions
+    groups: Optional[int]
+        Number of groups
+    bias: Optional[bool]
+        Whether to use bias
+
+    Returns
+    -------
+    out: torch.Tensor
+        Output tensor
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_shape: Union[int, Tuple[int], Tuple[int, int]],
+        basis_type: Optional[str] = "piecewise linear",
+        groups: Optional[int] = 1,
+        bias: Optional[bool] = True,
+    ):
+        super().__init__()
+
+        self.kernel_shape = kernel_shape
+
+        # get the filter basis functions
+        self.filter_basis = get_filter_basis(kernel_shape=kernel_shape, basis_type=basis_type)
+
+        # groups
+        self.groups = groups
+
+        # weight tensor
+        if in_channels % self.groups != 0:
+            raise ValueError("Error, the number of input channels has to be an integer multiple of the group size")
+        if out_channels % self.groups != 0:
+            raise ValueError("Error, the number of output channels has to be an integer multiple of the group size")
+        self.groupsize = in_channels // self.groups
+        scale = math.sqrt(1.0 / self.groupsize / self.kernel_size)
+        self.weight = nn.Parameter(scale * torch.randn(out_channels, self.groupsize, self.kernel_size))
+
+        if bias:
+            self.bias = nn.Parameter(torch.zeros(out_channels))
+        else:
+            self.bias = None
+
+    @property
+    def kernel_size(self):
+        return self.filter_basis.kernel_size
+
+    @abc.abstractmethod
+    def forward(self, x: torch.Tensor):
+        raise NotImplementedError
+
+
+class DiscreteContinuousConvS2(DiscreteContinuousConv):
+    """
+    Discrete-continuous (DISCO) convolutions on the 2-Sphere as described in [1].
+
+    Parameters
+    -----------
+    in_channels: int
+        Number of input channels
+    out_channels: int
+        Number of output channels
+    in_shape: Tuple[int]
+        Input shape of the convolution tensor
+    out_shape: Tuple[int]
+        Output shape of the convolution tensor
+    kernel_shape: Union[int, Tuple[int], Tuple[int, int]]
+        Shape of the kernel
+    basis_type: Optional[str]
+        Type of the basis functions
+    basis_norm_mode: Optional[str]
+        Mode for basis normalization
+    groups: Optional[int]
+        Number of groups
+    grid_in: Optional[str]
+        Input grid type
+    grid_out: Optional[str]
+        Output grid type
+    bias: Optional[bool]
+        Whether to use bias
+    theta_cutoff: Optional[float]
+        Theta cutoff for the filter basis functions
+
+    Returns
+    -------
+    out: torch.Tensor
+        Output tensor
+
+    References
+    ----------
+    [1] Ocampo, Price, McEwen, Scalable and equivariant spherical CNNs by discrete-continuous (DISCO) convolutions, ICLR (2023), arXiv:2209.13603
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        in_shape: Tuple[int],
+        out_shape: Tuple[int],
+        kernel_shape: Union[int, Tuple[int], Tuple[int, int]],
+        basis_type: Optional[str] = "piecewise linear",
+        basis_norm_mode: Optional[str] = "mean",
+        groups: Optional[int] = 1,
+        grid_in: Optional[str] = "equiangular",
+        grid_out: Optional[str] = "equiangular",
+        bias: Optional[bool] = True,
+        theta_cutoff: Optional[float] = None,
+    ):
+        super().__init__(in_channels, out_channels, kernel_shape, basis_type, groups, bias)
+
+        self.nlat_in, self.nlon_in = in_shape
+        self.nlat_out, self.nlon_out = out_shape
+
+        # make sure the p-shift works by checking that longitudes are divisible
+        assert self.nlon_in % self.nlon_out == 0
+
+        # heuristic to compute theta cutoff based on the bandlimit of the input field and overlaps of the basis functions
+        if theta_cutoff is None:
+            theta_cutoff = torch.pi / float(self.nlat_out - 1)
+
+        if theta_cutoff <= 0.0:
+            raise ValueError("Error, theta_cutoff has to be positive.")
+
+        idx, vals, _ = _precompute_convolution_tensor_s2(
+            in_shape,
+            out_shape,
+            self.filter_basis,
+            grid_in=grid_in,
+            grid_out=grid_out,
+            theta_cutoff=theta_cutoff,
+            transpose_normalization=False,
+            basis_norm_mode=basis_norm_mode,
+            merge_quadrature=True,
+        )
+
+        # sort the values
+        ker_idx = idx[0, ...].contiguous()
+        row_idx = idx[1, ...].contiguous()
+        col_idx = idx[2, ...].contiguous()
+        vals = vals.contiguous()
+
+        if _cuda_extension_available:
+            # preprocessed data-structure for GPU kernel
+            roff_idx = preprocess_psi(self.kernel_size, self.nlat_out, ker_idx, row_idx, col_idx, vals).contiguous()
+            self.register_buffer("psi_roff_idx", roff_idx, persistent=False)
+
+        # save all datastructures
+        self.register_buffer("psi_ker_idx", ker_idx, persistent=False)
+        self.register_buffer("psi_row_idx", row_idx, persistent=False)
+        self.register_buffer("psi_col_idx", col_idx, persistent=False)
+        self.register_buffer("psi_vals", vals, persistent=False)
+
+        # also store psi as COO matrix just in case for torch input
+        self.psi = _get_psi(self.kernel_size, self.psi_idx, self.psi_vals, self.nlat_in, self.nlon_in, self.nlat_out, self.nlon_out)
+
+    def extra_repr(self):
+        return f"in_shape={(self.nlat_in, self.nlon_in)}, out_shape={(self.nlat_out, self.nlon_out)}, in_chans={self.groupsize * self.groups}, out_chans={self.weight.shape[0]}, filter_basis={self.filter_basis}, kernel_shape={self.kernel_shape}, groups={self.groups}"
+
+    @property
+    def psi_idx(self):
+        return torch.stack([self.psi_ker_idx, self.psi_row_idx, self.psi_col_idx], dim=0).contiguous()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+
+        if x.is_cuda and _cuda_extension_available:
+            x = _disco_s2_contraction_cuda(
+                x, self.psi_roff_idx, self.psi_ker_idx, self.psi_row_idx, self.psi_col_idx, self.psi_vals, self.kernel_size, self.nlat_out, self.nlon_out
+            )
+        else:
+            if x.is_cuda:
+                warn("couldn't find CUDA extension, falling back to slow PyTorch implementation")
+            x = _disco_s2_contraction_torch(x, self.psi.to(x.device), self.nlon_out)
+
+        # extract shape
+        B, C, K, H, W = x.shape
+        x = x.reshape(B, self.groups, self.groupsize, K, H, W)
+
+        # do weight multiplication
+        out = torch.einsum("bgckxy,gock->bgoxy", x, self.weight.reshape(self.groups, -1, self.weight.shape[1], self.weight.shape[2])).contiguous()
+        out = out.reshape(B, -1, H, W)
+
+        if self.bias is not None:
+            out = out + self.bias.reshape(1, -1, 1, 1)
+
+        return out
+
+
+class DiscreteContinuousConvTransposeS2(DiscreteContinuousConv):
+    """
+    Discrete-continuous (DISCO) transpose convolutions on the 2-Sphere as described in [1].
+
+    Parameters
+    -----------
+    in_channels: int
+        Number of input channels
+    out_channels: int
+        Number of output channels
+    in_shape: Tuple[int]
+        Input shape of the convolution tensor
+    out_shape: Tuple[int]
+        Output shape of the convolution tensor
+    kernel_shape: Union[int, Tuple[int], Tuple[int, int]]
+        Shape of the kernel
+    basis_type: Optional[str]
+        Type of the basis functions
+    basis_norm_mode: Optional[str]
+        Mode for basis normalization
+    groups: Optional[int]
+        Number of groups
+    grid_in: Optional[str]
+        Input grid type
+    grid_out: Optional[str]
+        Output grid type
+    bias: Optional[bool]
+        Whether to use bias
+    theta_cutoff: Optional[float]
+        Theta cutoff for the filter basis functions
+
+    Returns
+    --------
+    out: torch.Tensor
+        Output tensor
+
+    References
+    ----------
+    [1] Ocampo, Price, McEwen, Scalable and equivariant spherical CNNs by discrete-continuous (DISCO) convolutions, ICLR (2023), arXiv:2209.13603
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        in_shape: Tuple[int],
+        out_shape: Tuple[int],
+        kernel_shape: Union[int, Tuple[int], Tuple[int, int]],
+        basis_type: Optional[str] = "piecewise linear",
+        basis_norm_mode: Optional[str] = "mean",
+        groups: Optional[int] = 1,
+        grid_in: Optional[str] = "equiangular",
+        grid_out: Optional[str] = "equiangular",
+        bias: Optional[bool] = True,
+        theta_cutoff: Optional[float] = None,
+    ):
+        super().__init__(in_channels, out_channels, kernel_shape, basis_type, groups, bias)
+
+        self.nlat_in, self.nlon_in = in_shape
+        self.nlat_out, self.nlon_out = out_shape
+
+        # make sure the p-shift works by checking that longitudes are divisible
+        assert self.nlon_out % self.nlon_in == 0
+
+        # bandlimit
+        if theta_cutoff is None:
+            theta_cutoff = torch.pi / float(self.nlat_in - 1)
+
+        if theta_cutoff <= 0.0:
+            raise ValueError("Error, theta_cutoff has to be positive.")
+
+        # switch in_shape and out_shape since we want the transpose convolution
+        idx, vals, _ = _precompute_convolution_tensor_s2(
+            out_shape,
+            in_shape,
+            self.filter_basis,
+            grid_in=grid_out,
+            grid_out=grid_in,
+            theta_cutoff=theta_cutoff,
+            transpose_normalization=True,
+            basis_norm_mode=basis_norm_mode,
+            merge_quadrature=True,
+        )
+
+        # sort the values
+        ker_idx = idx[0, ...].contiguous()
+        row_idx = idx[1, ...].contiguous()
+        col_idx = idx[2, ...].contiguous()
+        vals = vals.contiguous()
+
+        if _cuda_extension_available:
+            # preprocessed data-structure for GPU kernel
+            roff_idx = preprocess_psi(self.kernel_size, self.nlat_in, ker_idx, row_idx, col_idx, vals).contiguous()
+            self.register_buffer("psi_roff_idx", roff_idx, persistent=False)
+
+        # save all datastructures
+        self.register_buffer("psi_ker_idx", ker_idx, persistent=False)
+        self.register_buffer("psi_row_idx", row_idx, persistent=False)
+        self.register_buffer("psi_col_idx", col_idx, persistent=False)
+        self.register_buffer("psi_vals", vals, persistent=False)
+
+        # also store psi just in case
+        self.psi_st = _get_psi(self.kernel_size, self.psi_idx, self.psi_vals, self.nlat_in, self.nlon_in, self.nlat_out, self.nlon_out, semi_transposed=True)
+
+    def extra_repr(self):
+        return f"in_shape={(self.nlat_in, self.nlon_in)}, out_shape={(self.nlat_out, self.nlon_out)}, in_chans={self.groupsize * self.groups}, out_chans={self.weight.shape[0]}, filter_basis={self.filter_basis}, kernel_shape={self.kernel_shape}, groups={self.groups}"
+
+    @property
+    def psi_idx(self):
+        return torch.stack([self.psi_ker_idx, self.psi_row_idx, self.psi_col_idx], dim=0).contiguous()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+
+        # extract shape
+        B, C, H, W = x.shape
+        x = x.reshape(B, self.groups, self.groupsize, H, W)
+
+        # do weight multiplication
+        x = torch.einsum("bgcxy,gock->bgokxy", x, self.weight.reshape(self.groups, -1, self.weight.shape[1], self.weight.shape[2])).contiguous()
+        x = x.reshape(B, -1, x.shape[-3], H, W)
+
+        if x.is_cuda and _cuda_extension_available:
+            out = _disco_s2_transpose_contraction_cuda(
+                x, self.psi_roff_idx, self.psi_ker_idx, self.psi_row_idx, self.psi_col_idx, self.psi_vals, self.kernel_size, self.nlat_out, self.nlon_out
+            )
+        else:
+            if x.is_cuda:
+                warn("couldn't find CUDA extension, falling back to slow PyTorch implementation")
+            out = _disco_s2_transpose_contraction_torch(x, self.psi_st.to(x.device), self.nlon_out)
+
+        if self.bias is not None:
+            out = out + self.bias.reshape(1, -1, 1, 1)
+
+        return out

.deps/torch_harmonics/csrc/attention/attention.cuh

@@ -0,0 +1,50 @@
+// coding=utf-8
+//
+// SPDX-FileCopyrightText: Copyright (c) 2025 The torch-harmonics Authors. All rights reserved.
+// SPDX-License-Identifier: BSD-3-Clause
+//
+// Redistribution and use in source and binary forms, with or without
+// modification, are permitted provided that the following conditions are met:
+//
+// 1. Redistributions of source code must retain the above copyright notice, this
+// list of conditions and the following disclaimer.
+//
+// 2. Redistributions in binary form must reproduce the above copyright notice,
+// this list of conditions and the following disclaimer in the documentation
+// and/or other materials provided with the distribution.
+//
+// 3. Neither the name of the copyright holder nor the names of its
+// contributors may be used to endorse or promote products derived from
+// this software without specific prior written permission.
+//
+// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+// FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+// OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+#pragma once
+
+#include <cmath>
+#include <cstdint>
+#include <torch/torch.h>
+
+#define CHECK_CUDA_TENSOR(x) TORCH_INTERNAL_ASSERT(x.device().type() == torch::kCUDA)
+#define CHECK_CONTIGUOUS_TENSOR(x) TORCH_INTERNAL_ASSERT(x.is_contiguous() || x.is_contiguous(at::MemoryFormat::ChannelsLast))
+#define CHECK_CUDA_INPUT_TENSOR(x)                                                                                     \
+    CHECK_CUDA_TENSOR(x);                                                                                              \
+    CHECK_CONTIGUOUS_TENSOR(x)
+
+torch::Tensor s2_attention_fwd_cuda(at::Tensor kx, at::Tensor vx, at::Tensor qy, at::Tensor quad_weights,
+                                    at::Tensor psi_col_idx, at::Tensor psi_row_off, int nlon_in, int nlat_out,
+                                    int nlon_out);
+
+std::tuple<at::Tensor, at::Tensor, at::Tensor> s2_attention_bwd_dkvq_cuda(at::Tensor kx, at::Tensor vx, at::Tensor qy,
+                                                                          at::Tensor dy, at::Tensor quad_weights,
+                                                                          at::Tensor psi_col_idx, at::Tensor psi_row_off,
+                                                                          int nlon_in, int nlat_out, int nlon_out);

.deps/torch_harmonics/csrc/attention/attention_bwd_cuda.cu

@@ -0,0 +1,960 @@
+// coding=utf-8
+//
+// SPDX-FileCopyrightText: Copyright (c) 2025 The torch-harmonics Authors. All rights reserved.
+// SPDX-License-Identifier: BSD-3-Clause
+//
+// Redistribution and use in source and binary forms, with or without
+// modification, are permitted provided that the following conditions are met:
+//
+// 1. Redistributions of source code must retain the above copyright notice, this
+// list of conditions and the following disclaimer.
+//
+// 2. Redistributions in binary form must reproduce the above copyright notice,
+// this list of conditions and the following disclaimer in the documentation
+// and/or other materials provided with the distribution.
+//
+// 3. Neither the name of the copyright holder nor the names of its
+// contributors may be used to endorse or promote products derived from
+// this software without specific prior written permission.
+//
+// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+// FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+// OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+#include "attention.cuh"
+#include "c10/core/MemoryFormat.h"
+
+#include <ATen/core/TensorAccessor.h>
+#include <ATen/cuda/detail/TensorInfo.cuh>
+#include <ATen/cuda/detail/KernelUtils.h>
+#include <ATen/cuda/detail/IndexUtils.cuh>
+#include <ATen/cuda/CUDAUtils.h>
+
+#include <ctime>
+#include <cub/cub.cuh>
+#include <limits>
+
+#include "cudamacro.h"
+#include "attention_utils.cuh"
+
+#include <iostream>
+#include <chrono>
+#include <string>
+
+#define THREADS (64)
+
+#define MAX_LOCAL_ARR_LEN (16)
+
+#if 0
+class ScopeTimer
+{
+  public:
+    explicit ScopeTimer(const std::string &label = "") :
+        label_(label), start_(std::chrono::high_resolution_clock::now())
+    {
+    }
+
+    ~ScopeTimer()
+    {
+        auto end = std::chrono::high_resolution_clock::now();
+        auto elapsed = std::chrono::duration_cast<std::chrono::milliseconds>(end - start_);
+        std::cout << label_ << "Elapsed time: " << elapsed.count() << " ms" << std::endl;
+    }
+
+  private:
+    std::string label_;
+    std::chrono::high_resolution_clock::time_point start_;
+};
+
+// easier to understand version of manual shfl_xor_sync, performance appears similar
+static __device__ float __warp_sum_cub(float val)
+{
+    // use cub to reduce within a warp
+    __shared__ typename cub::WarpReduce<float>::TempStorage temp_storage;
+
+    // 1. Compute sum (initially only in lane 0)
+    float sum = cub::WarpReduce<float>(temp_storage).Sum(val);
+    // 2. Broadcast sum to all threads
+    sum = __shfl_sync(0xFFFFFFFF, sum, 0);
+    return sum;
+}
+
+// This kernel computes the backward pass for the S2 attention mechanism, using
+// shared memory as a cache and one warp per output point, warp-parallel over
+// channels, which should be layed out in the fastest dimension for coalesced
+// memory access.
+template <int BDIM_X>
+__global__ __launch_bounds__(BDIM_X) void s2_attention_bwd_dkvq_kernel(
+    int num_channels, int nlon_in, int nlat_out, int nlon_out,
+    const torch::PackedTensorAccessor32<float, 4, torch::RestrictPtrTraits> kx,
+    const torch::PackedTensorAccessor32<float, 4, torch::RestrictPtrTraits> vx,
+    const torch::PackedTensorAccessor32<float, 4, torch::RestrictPtrTraits> qy,
+    const torch::PackedTensorAccessor32<float, 4, torch::RestrictPtrTraits> dy,
+    torch::PackedTensorAccessor32<float, 4, torch::RestrictPtrTraits> dydk,
+    torch::PackedTensorAccessor32<float, 4, torch::RestrictPtrTraits> dydv,
+    torch::PackedTensorAccessor32<float, 4, torch::RestrictPtrTraits> dydq,
+    const torch::PackedTensorAccessor64<int64_t, 1, torch::RestrictPtrTraits> psi_col_idx,
+    const torch::PackedTensorAccessor64<int64_t, 1, torch::RestrictPtrTraits> psi_row_offset,
+    const torch::PackedTensorAccessor32<float, 1, torch::RestrictPtrTraits> quad_weights)
+{
+
+    extern __shared__ float sh[];
+    float *sh_alpha_k = sh + threadIdx.y * num_channels * 5;
+    float *sh_alpha_vw = sh_alpha_k + num_channels;
+    float *sh_alpha_kvw = sh_alpha_vw + num_channels;
+    float *sh_dy = sh_alpha_kvw + num_channels;
+    float *sh_qy = sh_dy + num_channels;
+    // (optionally, could use more shared memory for other intermediates)
+
+    const uint64_t batchId = blockIdx.y;
+    const uint64_t wid = uint64_t(blockIdx.x) * blockDim.y + threadIdx.y;
+    if (wid >= uint64_t(nlat_out) * nlon_in) return;
+    const int tidx = threadIdx.x;
+    const int ho = wid / nlon_out;
+    const int wo = wid - (ho * nlon_out);
+
+    // Zero shared memory
+    for (int chan = tidx; chan < num_channels; chan += WARP_SIZE) {
+        sh_alpha_k[chan] = 0.0f;
+        sh_alpha_vw[chan] = 0.0f;
+        sh_alpha_kvw[chan] = 0.0f;
+        sh_dy[chan] = dy[batchId][chan][ho][wo];
+        sh_qy[chan] = qy[batchId][chan][ho][wo];
+    }
+    float alpha_sum = 0.0f;
+    float qdotk_max = -FLT_MAX;
+    float integral = 0.0f;
+    __syncthreads();
+
+    const int64_t rbeg = psi_row_offset[ho];
+    const int64_t rend = psi_row_offset[ho + 1];
+    const int rlen = rend - rbeg;
+
+    // 1st pass: accumulate alpha_sum, integral, and shared stats, along with a progressively computed qdotk_max.
+    for (int off = 0; off < rlen; off++) {
+        const int64_t col = psi_col_idx[rbeg + off];
+        const int hi = col / nlon_in;
+        const int wi = col - (hi * nlon_in);
+        const int wip = (wi + wo) - ((wi + wo) / nlon_in) * nlon_in;
+        float qdotk = 0.0f, gdotv = 0.0f;
+        for (int chan = tidx; chan < num_channels; chan += WARP_SIZE) {
+            qdotk += sh_qy[chan] * kx[batchId][chan][hi][wip];
+            gdotv += sh_dy[chan] * vx[batchId][chan][hi][wip];
+        }
+        qdotk = __warp_sum_cub(qdotk);
+        gdotv = __warp_sum_cub(gdotv);
+        float qdotk_max_tmp = max(qdotk_max, qdotk);
+        float alpha_inz = expf(qdotk - qdotk_max_tmp) * quad_weights[hi];
+        float max_correction = expf(qdotk_max - qdotk_max_tmp);
+        alpha_sum = alpha_sum * max_correction + alpha_inz;
+        integral = integral * max_correction + alpha_inz * gdotv;
+        for (int chan = tidx; chan < num_channels; chan += WARP_SIZE) {
+            float kxval = kx[batchId][chan][hi][wip];
+            sh_alpha_k[chan] = sh_alpha_k[chan] * max_correction + alpha_inz * kxval;
+            sh_alpha_vw[chan] = sh_alpha_vw[chan] * max_correction + alpha_inz * gdotv;
+            sh_alpha_kvw[chan] = sh_alpha_kvw[chan] * max_correction + alpha_inz * kxval * gdotv;
+        }
+        qdotk_max = qdotk_max_tmp;
+    }
+
+    integral /= alpha_sum;
+
+    // Write dydq
+    for (int chan = tidx; chan < num_channels; chan += WARP_SIZE) {
+        dydq[batchId][chan][ho][wo]
+            = (sh_alpha_kvw[chan] * alpha_sum - sh_alpha_vw[chan] * sh_alpha_k[chan]) / (alpha_sum * alpha_sum);
+    }
+
+    // Third pass: accumulate gradients for k and v
+    for (int off = 0; off < rlen; off++) {
+        const int64_t col = psi_col_idx[rbeg + off];
+        const int hi = col / nlon_in;
+        const int wi = col - (hi * nlon_in);
+        const int wip = (wi + wo) - ((wi + wo) / nlon_in) * nlon_in;
+        float qdotk = 0.0f, gdotv = 0.0f;
+        for (int chan = tidx; chan < num_channels; chan += WARP_SIZE) {
+            qdotk += qy[batchId][chan][ho][wo] * kx[batchId][chan][hi][wip];
+            gdotv += sh_dy[chan] * vx[batchId][chan][hi][wip];
+        }
+        qdotk = __warp_sum_cub(qdotk);
+        gdotv = __warp_sum_cub(gdotv);
+        float alpha_inz = expf(qdotk - qdotk_max) * quad_weights[hi];
+        for (int chan = tidx; chan < num_channels; chan += WARP_SIZE) {
+            float qyval = qy[batchId][chan][ho][wo];
+            float dyval = sh_dy[chan];
+            atomicAdd(&dydk[batchId][chan][hi][wip], qyval * (alpha_inz / alpha_sum) * (gdotv - integral));
+            atomicAdd(&dydv[batchId][chan][hi][wip], (alpha_inz / alpha_sum) * dyval);
+        }
+    }
+}
+#endif
+
+// BEGIN backward kernels and functions
+
+// called with (blockDim.x=32 and blockDim.y>1, BDIM=blockDim.x*blockDim.y)
+template<int BDIM_X,
+         typename FLOATV_T> // either float or float4
+__global__
+__launch_bounds__(BDIM_X)
+void s2_attn_bwd_generic_vec_k(int nchan,  // no. of FLOATV_T elements along channel dim
+                               int nlat_in,
+                               int nlon_in,
+                               int nlat_out,
+                               int nlon_out,
+                               const FLOATV_T *__restrict__ kx,
+                               const FLOATV_T *__restrict__ vx,
+                               const FLOATV_T *__restrict__ qy,
+                               const FLOATV_T *__restrict__ dy,
+                               const int32_t *__restrict__ row_idx,
+                               const int64_t *__restrict__ row_off,
+                               const int64_t *__restrict__ col_idx,
+                               const   float *__restrict__ quad_weights,
+                                     FLOATV_T *__restrict__ dkx,
+                                     FLOATV_T *__restrict__ dvx,
+                                     FLOATV_T *__restrict__ dqy) {
+
+    extern __shared__ __align__(sizeof(float4)) float shext[];
+
+    // for dqy
+    FLOATV_T *sh_alpha_k__ = reinterpret_cast<FLOATV_T *>(shext) + threadIdx.y * nchan*5;
+    FLOATV_T *sh_alpha_vw_ = sh_alpha_k__ + nchan;
+    FLOATV_T *sh_alpha_kvw = sh_alpha_vw_ + nchan;
+
+    FLOATV_T *sh_dy        = sh_alpha_kvw + nchan;
+    FLOATV_T *sh_qy        = sh_dy        + nchan;
+
+    const int batch = blockIdx.y;
+
+    const uint64_t wid = uint64_t(blockIdx.x) * blockDim.y + threadIdx.y;
+    if (wid >= uint64_t(nlat_out)*nlon_in) {
+        return;
+    }
+
+    const int tidx = threadIdx.x;
+
+    // use permuted rows
+    const int h = wid / nlon_out;
+    const int wo = wid - (h*nlon_out);
+    const int ho = row_idx[h];
+
+    // offset input tensors
+    kx += int64_t(batch)*nlat_in*nlon_in*nchan;
+    vx += int64_t(batch)*nlat_in*nlon_in*nchan;
+    qy += int64_t(batch)*nlat_out*nlon_out*nchan + int64_t(ho)*nlon_out*nchan + int64_t(wo)*nchan;
+    dy += int64_t(batch)*nlat_out*nlon_out*nchan + int64_t(ho)*nlon_out*nchan + int64_t(wo)*nchan;
+
+    // offset output tensors
+    dkx += int64_t(batch)*nlat_in*nlon_in*nchan;
+    dvx += int64_t(batch)*nlat_in*nlon_in*nchan;
+    dqy += int64_t(batch)*nlat_out*nlon_out*nchan + int64_t(ho)*nlon_out*nchan + int64_t(wo)*nchan;
+
+    // zero/init shared memory
+    for (int chan = tidx; chan < nchan; chan += WARP_SIZE) {
+
+        sh_alpha_k__[chan] = __vset<FLOATV_T>(0.0f);
+        sh_alpha_vw_[chan] = __vset<FLOATV_T>(0.0f);
+        sh_alpha_kvw[chan] = __vset<FLOATV_T>(0.0f);
+
+        sh_dy[chan] = dy[chan];
+        sh_qy[chan] = qy[chan];
+    }
+
+#if __CUDA_ARCH__ < 900
+    // for architectures < 9.0, sh_dy and sh_qy will be read
+    // as individual floats at the end of the kernel, which
+    // breaks the assumption that each FLOATV_T location is
+    // written to and read by the same thread throughout the 
+    // kernel, in the case FLOATV_T==float4
+    if constexpr(std::is_same<FLOATV_T, float4>::value) { __syncwarp(); }
+#endif
+
+    // for dkx, dvx, dqy
+    float alpha_sum = 0.0f;
+    float qdotk_max = -FLT_MAX;
+    
+    // for dkx
+    float integral = 0.0f;
+
+    const int64_t rbeg = row_off[ho];
+    const int64_t rend = row_off[ho+1];
+
+    col_idx += rbeg;
+
+    const int rlen = rend - rbeg;
+
+    // accumulate alpha_sum, integral, and shared stats, 
+    // along with a progressively computed qdotk_max.
+    for (int off = 0; off < rlen; off++) {
+        
+        const int64_t col = col_idx[off];
+
+        const int hi = col / nlon_in;
+        const int wi = col - (hi * nlon_in);
+        const int wip = (wi + wo) - ((wi + wo) / nlon_in) * nlon_in;
+
+        const FLOATV_T *_kx = kx + int64_t(hi)*nlon_in*nchan + int64_t(wip)*nchan;
+        const FLOATV_T *_vx = vx + int64_t(hi)*nlon_in*nchan + int64_t(wip)*nchan;
+
+        FLOATV_T qdotk_v = __vset<FLOATV_T>(0.0f);
+        FLOATV_T gdotv_v = __vset<FLOATV_T>(0.0f);
+
+        for(int chan = tidx; chan < nchan; chan += WARP_SIZE) {
+
+            qdotk_v = __vadd(qdotk_v, __vmul(sh_qy[chan], _kx[chan]));
+            gdotv_v = __vadd(gdotv_v, __vmul(sh_dy[chan], _vx[chan]));
+        }
+
+        const float qdotk = __warp_sum(__vred(qdotk_v));
+        const float gdotv = __warp_sum(__vred(gdotv_v));
+
+        const float qdotk_max_tmp = max(qdotk_max, qdotk);
+        const float alpha_inz = expf(qdotk - qdotk_max_tmp) * quad_weights[hi];
+        const float max_correction = expf(qdotk_max - qdotk_max_tmp);
+        alpha_sum = alpha_sum * max_correction + alpha_inz;
+
+        integral = integral * max_correction + alpha_inz * gdotv;
+
+        const float ainz_gdotv = alpha_inz * gdotv;
+
+        for (int chan = tidx; chan < nchan; chan += WARP_SIZE) {
+
+            const FLOATV_T kxval = _kx[chan];
+
+            sh_alpha_k__[chan] = __vadd(__vscale(max_correction, sh_alpha_k__[chan]), __vscale(alpha_inz, kxval));
+            sh_alpha_vw_[chan] = __vadd(__vscale(max_correction, sh_alpha_vw_[chan]), __vset<FLOATV_T>(ainz_gdotv));
+            sh_alpha_kvw[chan] = __vadd(__vscale(max_correction, sh_alpha_kvw[chan]), __vscale(ainz_gdotv, kxval));
+        }
+        qdotk_max = qdotk_max_tmp;
+    }
+
+    const float alpha_sum_inv = 1.0f / alpha_sum;
+
+    integral *= alpha_sum_inv;
+
+    // Write dqy
+    for (int chan = tidx; chan < nchan; chan += WARP_SIZE) {
+
+        dqy[chan] = __vscale(alpha_sum_inv * alpha_sum_inv,
+                             __vsub(__vscale(alpha_sum, sh_alpha_kvw[chan]),
+                                    __vmul(sh_alpha_vw_[chan], sh_alpha_k__[chan])));
+    }
+
+    // accumulate gradients for k and v
+    for (int off = 0; off < rlen; off++) {
+
+        const int64_t col = col_idx[off];
+        const int hi = col / nlon_in;
+        const int wi = col - (hi * nlon_in);
+        const int wip = (wi + wo) - ((wi + wo) / nlon_in) * nlon_in;
+
+        const FLOATV_T *_kx = kx + int64_t(hi)*nlon_in*nchan + int64_t(wip)*nchan;
+        const FLOATV_T *_vx = vx + int64_t(hi)*nlon_in*nchan + int64_t(wip)*nchan;
+
+        FLOATV_T qdotk_v = __vset<FLOATV_T>(0.0f);
+        FLOATV_T gdotv_v = __vset<FLOATV_T>(0.0f);
+
+        for (int chan = tidx; chan < nchan; chan += WARP_SIZE) {
+
+            qdotk_v = __vadd(qdotk_v, __vmul(sh_qy[chan], _kx[chan]));
+            gdotv_v = __vadd(gdotv_v, __vmul(sh_dy[chan], _vx[chan]));
+        }
+
+        const float qdotk = __warp_sum(__vred(qdotk_v));
+        const float gdotv = __warp_sum(__vred(gdotv_v));
+
+        const float alpha_inz = expf(qdotk - qdotk_max) * quad_weights[hi];
+
+        FLOATV_T *_dkx = dkx + int64_t(hi)*nlon_in*nchan + int64_t(wip)*nchan;
+        FLOATV_T *_dvx = dvx + int64_t(hi)*nlon_in*nchan + int64_t(wip)*nchan;
+
+        const float alpha_mul = alpha_inz * alpha_sum_inv;
+
+        const float scale_fact_qy = (gdotv - integral)*alpha_mul;
+        const float scale_fact_dy =                    alpha_mul;
+
+        // float4, 128-bit atomics are only supported by devices of compute 
+        // capability 9.x+, so on older devices we resort to 32-bit atomics
+
+#if __CUDA_ARCH__ < 900
+        // to use 32-bit operations on consecutve addresses
+        float *sh_qy_scl = reinterpret_cast<float *>(sh_qy);
+        float *sh_dy_scl = reinterpret_cast<float *>(sh_dy);
+
+        float *_dkx_scl = reinterpret_cast<float *>(_dkx);
+        float *_dvx_scl = reinterpret_cast<float *>(_dvx);
+
+        constexpr int VEC_SIZE = sizeof(FLOATV_T)/sizeof(float);
+
+        // 32-bit, consecutive atomics to glmem;
+        // strided atomics results in a severe slowdown
+        for (int chan = tidx; chan < nchan*VEC_SIZE; chan += WARP_SIZE) {
+
+            atomicAdd(_dkx_scl + chan, scale_fact_qy * sh_qy_scl[chan]);
+            atomicAdd(_dvx_scl + chan, scale_fact_dy * sh_dy_scl[chan]);
+        }
+#else
+        // 128-bit, consecutive atomics to glmem
+        for (int chan = tidx; chan < nchan; chan += WARP_SIZE) {
+
+            atomicAdd(_dkx + chan, __vscale(scale_fact_qy, sh_qy[chan]));
+            atomicAdd(_dvx + chan, __vscale(scale_fact_dy, sh_dy[chan]));
+        }
+#endif
+    }
+
+    return;
+}
+
+// called with either (BDIM_X=32 and BDIM_Y>1) || (2^K=BDIM_X > 32 and BDIM_Y=1)
+template<int BDIM_X,
+         int BDIM_Y,
+         int NLOC,
+         typename FLOATV_T> // either float or float4
+__global__
+__launch_bounds__(BDIM_X*BDIM_Y)
+void s2_attn_bwd_special_vec_k(int nchan,  // no. of FLOATV_T elements along channel dim
+                               int nlat_in,
+                               int nlon_in,
+                               int nlat_out,
+                               int nlon_out,
+                               const FLOATV_T *__restrict__ kx,
+                               const FLOATV_T *__restrict__ vx,
+                               const FLOATV_T *__restrict__ qy,
+                               const FLOATV_T *__restrict__ dy,
+                               const int32_t *__restrict__ row_idx,
+                               const int64_t *__restrict__ row_off,
+                               const int64_t *__restrict__ col_idx,
+                               const   float *__restrict__ quad_weights,
+                                     FLOATV_T *__restrict__ dkx,
+                                     FLOATV_T *__restrict__ dvx,
+                                     FLOATV_T *__restrict__ dqy) {
+
+    static_assert(0 == (BDIM_X & (BDIM_X-1)));
+    static_assert(0 == (BDIM_Y & (BDIM_Y-1)));
+    static_assert((BDIM_X == 32 && BDIM_Y  > 1) ||
+                  (BDIM_X  > 32 && BDIM_Y == 1)) ;
+
+    constexpr int NLOC_M1 = NLOC-1;
+
+    const int tidx = threadIdx.x;
+    const int batch = blockIdx.y;
+    const uint64_t ctaid = uint64_t(blockIdx.x) * blockDim.y + threadIdx.y;
+    
+    if (ctaid >= uint64_t(nlat_out)*nlon_in) {
+        return;
+    }
+
+    extern __shared__ __align__(sizeof(float4)) float shext[];
+
+    FLOATV_T *sh_dy = reinterpret_cast<FLOATV_T *>(shext) + threadIdx.y*nchan*2 + tidx;
+    FLOATV_T *sh_qy = sh_dy + nchan;
+
+    // for dqy
+    FLOATV_T loc_k__[NLOC];
+    FLOATV_T loc_vw_[NLOC];
+    FLOATV_T loc_kvw[NLOC];
+
+    // use permuted rows
+    const int h = ctaid / nlon_out;
+    const int wo = ctaid - (h*nlon_out);
+    const int ho = row_idx[h];
+
+    // offset input tensors
+    kx += int64_t(batch)*nlat_in*nlon_in*nchan + tidx;
+    vx += int64_t(batch)*nlat_in*nlon_in*nchan + tidx;
+    qy += int64_t(batch)*nlat_out*nlon_out*nchan + int64_t(ho)*nlon_out*nchan + int64_t(wo)*nchan + tidx;
+    dy += int64_t(batch)*nlat_out*nlon_out*nchan + int64_t(ho)*nlon_out*nchan + int64_t(wo)*nchan + tidx;
+
+    // offset output tensors
+    dkx += int64_t(batch)*nlat_in*nlon_in*nchan + tidx;
+    dvx += int64_t(batch)*nlat_in*nlon_in*nchan + tidx;
+    dqy += int64_t(batch)*nlat_out*nlon_out*nchan + int64_t(ho)*nlon_out*nchan + int64_t(wo)*nchan + tidx;
+
+    #pragma unroll
+    for(int i = 0; i < NLOC; i++) {
+        loc_k__[i] = __vset<FLOATV_T>(0.0f);
+        loc_vw_[i] = __vset<FLOATV_T>(0.0f);
+        loc_kvw[i] = __vset<FLOATV_T>(0.0f);
+    }
+
+    #pragma unroll
+    for(int i = 0; i < NLOC_M1; i++) {
+        sh_dy[i*BDIM_X] = dy[i*BDIM_X];
+        sh_qy[i*BDIM_X] = qy[i*BDIM_X];
+    }
+    if (NLOC_M1*BDIM_X+tidx < nchan) {
+        sh_dy[NLOC_M1*BDIM_X] = dy[NLOC_M1*BDIM_X];
+        sh_qy[NLOC_M1*BDIM_X] = qy[NLOC_M1*BDIM_X];
+    }
+
+#if __CUDA_ARCH__ < 900
+    // for architectures < 9.0, sh_dy and sh_qy will be read
+    // as individual floats at the end of the kernel, which
+    // breaks the assumption that each FLOATV_T location is
+    // written to and read by the same thread throughout the 
+    // kernel, in the case FLOATV_T==float4
+    if constexpr(std::is_same<FLOATV_T, float4>::value) {
+        if constexpr(BDIM_X == 32) {    __syncwarp(); }
+        else                       { __syncthreads(); }
+    }
+#endif
+
+    // for dkx, dvx, dqy
+    float alpha_sum = 0.0f;
+    float qdotk_max = -FLT_MAX;
+    
+    // for dkx
+    float integral = 0.0f;
+
+    const int64_t rbeg = row_off[ho];
+    const int64_t rend = row_off[ho+1];
+
+    col_idx += rbeg;
+
+    const int rlen = rend - rbeg;
+
+    // accumulate alpha_sum, integral, and shared stats, 
+    // along with a progressively computed qdotk_max.
+    for (int off = 0; off < rlen; off++) {
+        
+        const int64_t col = col_idx[off];
+
+        const int hi = col / nlon_in;
+        const int wi = col - (hi * nlon_in);
+        const int wip = (wi + wo) - ((wi + wo) / nlon_in) * nlon_in;
+
+        const FLOATV_T *_kx = kx + int64_t(hi)*nlon_in*nchan + int64_t(wip)*nchan;
+        const FLOATV_T *_vx = vx + int64_t(hi)*nlon_in*nchan + int64_t(wip)*nchan;
+
+        FLOATV_T qdotk_v = __vset<FLOATV_T>(0.0f);
+        FLOATV_T gdotv_v = __vset<FLOATV_T>(0.0f);
+
+        #pragma unroll
+        for(int i = 0; i < NLOC_M1; i++) {
+            qdotk_v = __vadd(qdotk_v, __vmul(sh_qy[i*BDIM_X], _kx[i*BDIM_X]));
+            gdotv_v = __vadd(gdotv_v, __vmul(sh_dy[i*BDIM_X], _vx[i*BDIM_X]));
+        }
+        if (NLOC_M1*BDIM_X+tidx < nchan) {
+            qdotk_v = __vadd(qdotk_v, __vmul(sh_qy[NLOC_M1*BDIM_X], _kx[NLOC_M1*BDIM_X]));
+            gdotv_v = __vadd(gdotv_v, __vmul(sh_dy[NLOC_M1*BDIM_X], _vx[NLOC_M1*BDIM_X]));
+        }
+
+        float qdotk = __vred(qdotk_v);
+        float gdotv = __vred(gdotv_v);
+
+        if constexpr(BDIM_X == 32) { 
+            qdotk = __warp_sum(qdotk); 
+            gdotv = __warp_sum(gdotv); 
+        } else {
+            qdotk = __block_sum<BDIM_X>(qdotk);
+            gdotv = __block_sum<BDIM_X>(gdotv);
+        }
+
+        const float qdotk_max_tmp = max(qdotk_max, qdotk);
+        const float alpha_inz = expf(qdotk - qdotk_max_tmp) * quad_weights[hi];
+        const float max_correction = expf(qdotk_max - qdotk_max_tmp);
+
+        alpha_sum = alpha_sum * max_correction + alpha_inz;
+        integral  = integral * max_correction + alpha_inz * gdotv;
+
+        const float ainz_gdotv = alpha_inz * gdotv;
+
+        #pragma unroll
+        for(int i = 0; i < NLOC_M1; i++) {
+            const FLOATV_T kxval = _kx[i*BDIM_X];
+            loc_k__[i] = __vadd(__vscale(max_correction, loc_k__[i]), __vscale(alpha_inz, kxval));
+            loc_vw_[i] = __vadd(__vscale(max_correction, loc_vw_[i]), __vset<FLOATV_T>(ainz_gdotv));
+            loc_kvw[i] = __vadd(__vscale(max_correction, loc_kvw[i]), __vscale(ainz_gdotv, kxval));
+        }
+        if (NLOC_M1*BDIM_X+tidx < nchan) {
+            const FLOATV_T kxval = _kx[NLOC_M1*BDIM_X];
+            loc_k__[NLOC_M1] = __vadd(__vscale(max_correction, loc_k__[NLOC_M1]), __vscale(alpha_inz, kxval));
+            loc_vw_[NLOC_M1] = __vadd(__vscale(max_correction, loc_vw_[NLOC_M1]), __vset<FLOATV_T>(ainz_gdotv));
+            loc_kvw[NLOC_M1] = __vadd(__vscale(max_correction, loc_kvw[NLOC_M1]), __vscale(ainz_gdotv, kxval));
+        }
+
+        qdotk_max = qdotk_max_tmp;
+    }
+
+    const float alpha_sum_inv = 1.0f / alpha_sum;
+
+    integral *= alpha_sum_inv;
+
+    // Write dqy
+    const float alpha_sum_inv_sq = alpha_sum_inv*alpha_sum_inv;
+
+    #pragma unroll
+    for(int i = 0; i < NLOC_M1; i++) {
+        dqy[i*BDIM_X] = __vscale(alpha_sum_inv_sq,
+                                 __vsub(__vscale(alpha_sum, loc_kvw[i]),
+                                        __vmul(loc_vw_[i], loc_k__[i])));
+    }
+    if (NLOC_M1*BDIM_X+tidx < nchan) {
+        dqy[NLOC_M1*BDIM_X] = __vscale(alpha_sum_inv_sq,
+                                       __vsub(__vscale(alpha_sum, loc_kvw[NLOC_M1]),
+                                              __vmul(loc_vw_[NLOC_M1], loc_k__[NLOC_M1])));
+    }
+
+    // accumulate gradients for k and v
+    for (int off = 0; off < rlen; off++) {
+
+        const int64_t col = col_idx[off];
+
+        const int hi = col / nlon_in;
+        const int wi = col - (hi * nlon_in);
+        const int wip = (wi + wo) - ((wi + wo) / nlon_in) * nlon_in;
+
+        const FLOATV_T *_kx = kx + int64_t(hi)*nlon_in*nchan + int64_t(wip)*nchan;
+        const FLOATV_T *_vx = vx + int64_t(hi)*nlon_in*nchan + int64_t(wip)*nchan;
+
+        FLOATV_T qdotk_v = __vset<FLOATV_T>(0.0f);
+        FLOATV_T gdotv_v = __vset<FLOATV_T>(0.0f);
+
+        #pragma unroll
+        for(int i = 0; i < NLOC_M1; i++) {
+            qdotk_v = __vadd(qdotk_v, __vmul(sh_qy[i*BDIM_X], _kx[i*BDIM_X]));
+            gdotv_v = __vadd(gdotv_v, __vmul(sh_dy[i*BDIM_X], _vx[i*BDIM_X]));
+        }
+        if (NLOC_M1*BDIM_X+tidx < nchan) {
+            qdotk_v = __vadd(qdotk_v, __vmul(sh_qy[NLOC_M1*BDIM_X], _kx[NLOC_M1*BDIM_X]));
+            gdotv_v = __vadd(gdotv_v, __vmul(sh_dy[NLOC_M1*BDIM_X], _vx[NLOC_M1*BDIM_X]));
+        }
+
+        float qdotk = __vred(qdotk_v);
+        float gdotv = __vred(gdotv_v);
+
+        if constexpr(BDIM_X == 32) { 
+            qdotk = __warp_sum(qdotk); 
+            gdotv = __warp_sum(gdotv); 
+        } else {
+            qdotk = __block_sum<BDIM_X>(qdotk);
+            gdotv = __block_sum<BDIM_X>(gdotv);
+        }
+
+        const float alpha_inz = expf(qdotk - qdotk_max) * quad_weights[hi];
+
+        FLOATV_T *_dkx = dkx + int64_t(hi)*nlon_in*nchan + int64_t(wip)*nchan;
+        FLOATV_T *_dvx = dvx + int64_t(hi)*nlon_in*nchan + int64_t(wip)*nchan;
+
+        const float alpha_mul = alpha_inz * alpha_sum_inv;
+
+        const float scale_fact_qy = (gdotv - integral)*alpha_mul;
+        const float scale_fact_dy =                    alpha_mul;
+
+        // float4, 128-bit atomics are only supported by devices of compute 
+        // capability 9.x+, so on older devices we resort to 32-bit atomics
+
+#if __CUDA_ARCH__ < 900
+        // making the loop count known at compile time doesn't seem
+        // to make any difference here so let's keep this (much) 
+        // simpler version
+        float *sh_qy_scl = reinterpret_cast<float *>(sh_qy - tidx);
+        float *sh_dy_scl = reinterpret_cast<float *>(sh_dy - tidx);
+
+        float *_dkx_scl = reinterpret_cast<float *>(_dkx - tidx);
+        float *_dvx_scl = reinterpret_cast<float *>(_dvx - tidx);
+
+        constexpr int VEC_SIZE = sizeof(FLOATV_T)/sizeof(float);
+        
+        // 32-bit, consecutive atomics to glmem
+        // strided atomics results in a severe slowdown
+        for (int chan = tidx; chan < nchan*VEC_SIZE; chan += BDIM_X) {
+
+            atomicAdd(_dkx_scl + chan, scale_fact_qy * sh_qy_scl[chan]);
+            atomicAdd(_dvx_scl + chan, scale_fact_dy * sh_dy_scl[chan]);
+        }
+#else
+        #pragma unroll
+        for(int i = 0; i < NLOC_M1; i++) {
+            atomicAdd(_dkx + i*BDIM_X, __vscale(scale_fact_qy, sh_qy[i*BDIM_X]));
+            atomicAdd(_dvx + i*BDIM_X, __vscale(scale_fact_dy, sh_dy[i*BDIM_X]));
+        }
+        if (NLOC_M1*BDIM_X+tidx < nchan) {
+            atomicAdd(_dkx + NLOC_M1*BDIM_X, __vscale(scale_fact_qy, sh_qy[NLOC_M1*BDIM_X]));
+            atomicAdd(_dvx + NLOC_M1*BDIM_X, __vscale(scale_fact_dy, sh_dy[NLOC_M1*BDIM_X]));
+        }
+#endif
+    }
+
+    return;
+}
+
+template<typename FLOATV_T>
+void launch_gen_attn_bwd(int batch_size,
+                         int nchans,
+                         int nlat_in,
+                         int nlon_in,
+                         int nlat_out,
+                         int nlon_out,
+                         FLOATV_T *_kxp,
+                         FLOATV_T *_vxp,
+                         FLOATV_T *_qyp,
+                         FLOATV_T *_dyp,
+                         int32_t *_row_idx,
+                         int64_t *_row_off,
+                         int64_t *_col_idx,
+                         float *_quad_weights,
+                         FLOATV_T *_dkxp,
+                         FLOATV_T *_dvxp,
+                         FLOATV_T *_dqyp,
+                         cudaStream_t stream) {
+
+    dim3 block(WARP_SIZE, THREADS / WARP_SIZE);
+    dim3 grid(DIV_UP(nlat_out*nlon_out, block.y), batch_size);
+
+    size_t shsize = sizeof(FLOATV_T)*nchans*5 * block.y; // 5 arrays per warp
+
+    s2_attn_bwd_generic_vec_k<THREADS>
+                             <<<grid, block, shsize, stream>>>(nchans, nlat_in, nlon_in, nlat_out, nlon_out,
+                                                               _kxp, _vxp, _qyp, _dyp, _row_idx, _row_off, _col_idx,
+                                                               _quad_weights, _dkxp, _dvxp, _dqyp);
+    CHECK_ERROR("s2_attn_bwd_generic_vec_k");
+
+    return;
+}
+
+template<int BDIM_X,
+         int BDIM_Y,
+         int CUR_LOC_SIZE,
+         int MAX_LOC_SIZE, // max size of FLOATV_T[] local array
+         typename FLOATV_T>
+void launch_spc_attn_bwd(int batch_size,
+                         int nloc,      // "BDIM_X*nloc" >= nchans
+                         int nchans,
+                         int nlat_in,
+                         int nlon_in,
+                         int nlat_out,
+                         int nlon_out,
+                         FLOATV_T *_kxp,
+                         FLOATV_T *_vxp,
+                         FLOATV_T *_qyp,
+                         FLOATV_T *_dyp,
+                         int32_t *_row_idx,
+                         int64_t *_row_off,
+                         int64_t *_col_idx,
+                         float *_quad_weights,
+                         FLOATV_T *_dkxp,
+                         FLOATV_T *_dvxp,
+                         FLOATV_T *_dqyp,
+                         cudaStream_t stream) {
+
+    if (CUR_LOC_SIZE == nloc) {
+
+        dim3 block(BDIM_X, BDIM_Y);
+        dim3 grid(DIV_UP(nlat_out*nlon_out, block.y), batch_size);
+
+        size_t shsize = sizeof(FLOATV_T)*nchans*2 * block.y; // 2 arrays per cta, block.y > 1 iif block.x==32
+
+        s2_attn_bwd_special_vec_k<BDIM_X, BDIM_Y, CUR_LOC_SIZE>
+                                 <<<grid, block, shsize, stream>>>(nchans, nlat_in, nlon_in, nlat_out, nlon_out,
+                                                                   _kxp, _vxp, _qyp, _dyp, _row_idx, _row_off, _col_idx,
+                                                                   _quad_weights, _dkxp, _dvxp, _dqyp);
+        CHECK_ERROR("s2_attn_bwd_special_vec_k");
+
+        return;
+    }
+    if constexpr(CUR_LOC_SIZE < MAX_LOC_SIZE) {
+        launch_spc_attn_bwd<BDIM_X,
+                            BDIM_Y,
+                            CUR_LOC_SIZE+1,
+                            MAX_LOC_SIZE>(batch_size, nloc, nchans, nlat_in, nlon_in, nlat_out, nlon_out,
+                                          _kxp, _vxp, _qyp, _dyp, _row_idx, _row_off, _col_idx, _quad_weights,
+                                          _dkxp, _dvxp, _dqyp, stream);
+    }
+    return;
+}
+
+static void s2_attn_bwd_dispatch(int batch_size,
+                                 int nchans,
+                                 int nlon_in,
+                                 int nlat_out,
+                                 int nlon_out,
+                                 at::Tensor kxP,
+                                 at::Tensor vxP,
+                                 at::Tensor qyP,
+                                 at::Tensor dyP,
+                                 at::Tensor row_off,
+                                 at::Tensor col_idx,
+                                 at::Tensor quad_weights,
+                                 at::Tensor dkxP,
+                                 at::Tensor dvxP,
+                                 at::Tensor dqyP) {
+
+    static_assert(0 == (MAX_LOCAL_ARR_LEN & (MAX_LOCAL_ARR_LEN-1)));
+
+    // get stream
+    auto stream = at::cuda::getCurrentCUDAStream().stream();
+
+    // sort row indices (ho-s) in descending order
+    // based on (row_off[ho+1]-row_off[ho])
+    at::Tensor row_idx = sortRows(nlat_out, row_off, stream);
+
+    const int nlat_in = kxP.size(1);
+
+    // smallest power of two "bdimx" (>=32) s.t. bdimx*MAX_LOCAL_ARR_LEN >= nchans
+    int bdimx;
+    bdimx = DIV_UP(nchans, MAX_LOCAL_ARR_LEN);
+    bdimx = max(bdimx, WARP_SIZE);
+    bdimx = next_pow2(bdimx);
+
+    float *_kxp = reinterpret_cast<float *>(kxP.data_ptr());
+    float *_vxp = reinterpret_cast<float *>(vxP.data_ptr());
+    float *_qyp = reinterpret_cast<float *>(qyP.data_ptr());
+    float *_dyp = reinterpret_cast<float *>(dyP.data_ptr());
+
+    float *_dkxp = reinterpret_cast<float *>(dkxP.data_ptr());
+    float *_dvxp = reinterpret_cast<float *>(dvxP.data_ptr());
+    float *_dqyp = reinterpret_cast<float *>(dqyP.data_ptr());
+
+    int32_t *_row_idx = reinterpret_cast<int32_t *>(row_idx.data_ptr());
+    int64_t *_row_off = reinterpret_cast<int64_t *>(row_off.data_ptr());
+    int64_t *_col_idx = reinterpret_cast<int64_t *>(col_idx.data_ptr());
+    float *_quad_weights = reinterpret_cast<float *>(quad_weights.data_ptr());
+
+    constexpr int VEC_SIZE = sizeof(float4) / sizeof(float);
+
+    if (!is_aligned<sizeof(float4)>(_kxp) ||
+        !is_aligned<sizeof(float4)>(_vxp) ||
+        !is_aligned<sizeof(float4)>(_qyp) ||
+        !is_aligned<sizeof(float4)>(_dyp) ||
+        !is_aligned<sizeof(float4)>(_dkxp) ||
+        !is_aligned<sizeof(float4)>(_dvxp) ||
+        !is_aligned<sizeof(float4)>(_dqyp) ||
+        (nchans % VEC_SIZE) != 0) {
+
+        const int nloc = DIV_UP(nchans, bdimx);
+
+        // to avoid the compilation of unused template instances;
+        // we use a block size BDIM_X that is the smallest power of 2
+        // such that BDIM_X*MAX_LOCAL_ARR_LEN >= nchans, so
+        // BDIM_X > 32 are used only for:
+        //
+        //  (BDIM_X-1)*MAX_LOCAL_ARR_LEN < nchans <= BDIM_X*MAX_LOCAL_ARR_LEN
+        constexpr int MIN_LOC_ARR_LEN = MAX_LOCAL_ARR_LEN/2+1;
+
+        // use 2D blocks only if 32 threads are enough; w.r.t fowrard,
+        // we use the special kernel only up to BDIM_X=512 as with 1024
+        // each thread cannot use more than 64 registers, resulting in
+        // large amounts of registers spills
+        switch(bdimx) {
+            case   32: launch_spc_attn_bwd< 32, 2,               1, MAX_LOCAL_ARR_LEN>(batch_size, nloc, nchans, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, _vxp, _qyp, _dyp, _row_idx, _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp, stream); break;
+            case   64: launch_spc_attn_bwd< 64, 1, MIN_LOC_ARR_LEN, MAX_LOCAL_ARR_LEN>(batch_size, nloc, nchans, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, _vxp, _qyp, _dyp, _row_idx, _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp, stream); break;
+            case  128: launch_spc_attn_bwd<128, 1, MIN_LOC_ARR_LEN, MAX_LOCAL_ARR_LEN>(batch_size, nloc, nchans, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, _vxp, _qyp, _dyp, _row_idx, _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp, stream); break;
+            case  256: launch_spc_attn_bwd<256, 1, MIN_LOC_ARR_LEN, MAX_LOCAL_ARR_LEN>(batch_size, nloc, nchans, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, _vxp, _qyp, _dyp, _row_idx, _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp, stream); break;
+            case  512: launch_spc_attn_bwd<512, 1, MIN_LOC_ARR_LEN, MAX_LOCAL_ARR_LEN>(batch_size, nloc, nchans, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, _vxp, _qyp, _dyp, _row_idx, _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp, stream); break;
+            default:   launch_gen_attn_bwd                                            (batch_size,       nchans, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, _vxp, _qyp, _dyp, _row_idx, _row_off, _col_idx, _quad_weights, _dkxp, _dvxp, _dqyp, stream); break;
+        }
+
+    } else {
+
+        float4 *_kxp4 = reinterpret_cast<float4 *>(kxP.data_ptr());
+        float4 *_vxp4 = reinterpret_cast<float4 *>(vxP.data_ptr());
+        float4 *_qyp4 = reinterpret_cast<float4 *>(qyP.data_ptr());
+        float4 *_dyp4 = reinterpret_cast<float4 *>(dyP.data_ptr());
+
+        float4 *_dkxp4 = reinterpret_cast<float4 *>(dkxP.data_ptr());
+        float4 *_dvxp4 = reinterpret_cast<float4 *>(dvxP.data_ptr());
+        float4 *_dqyp4 = reinterpret_cast<float4 *>(dqyP.data_ptr());
+
+        nchans /= VEC_SIZE;
+        const int nloc = DIV_UP(nchans, bdimx);
+
+        constexpr int MAX_LOCAL_VEC_LEN = MAX_LOCAL_ARR_LEN / VEC_SIZE;
+        
+        constexpr int MIN_LOC_VEC_LEN = MAX_LOCAL_VEC_LEN/2+1;
+
+        // use 2D blocks only if 32 threads are enough
+        switch(bdimx) {
+            case   32: launch_spc_attn_bwd< 32, 2,               1, MAX_LOCAL_VEC_LEN>(batch_size, nloc, nchans, nlat_in, nlon_in, nlat_out, nlon_out, _kxp4, _vxp4, _qyp4, _dyp4, _row_idx, _row_off, _col_idx, _quad_weights, _dkxp4, _dvxp4, _dqyp4, stream); break;
+            case   64: launch_spc_attn_bwd< 64, 1, MIN_LOC_VEC_LEN, MAX_LOCAL_VEC_LEN>(batch_size, nloc, nchans, nlat_in, nlon_in, nlat_out, nlon_out, _kxp4, _vxp4, _qyp4, _dyp4, _row_idx, _row_off, _col_idx, _quad_weights, _dkxp4, _dvxp4, _dqyp4, stream); break;
+            case  128: launch_spc_attn_bwd<128, 1, MIN_LOC_VEC_LEN, MAX_LOCAL_VEC_LEN>(batch_size, nloc, nchans, nlat_in, nlon_in, nlat_out, nlon_out, _kxp4, _vxp4, _qyp4, _dyp4, _row_idx, _row_off, _col_idx, _quad_weights, _dkxp4, _dvxp4, _dqyp4, stream); break;
+            case  256: launch_spc_attn_bwd<256, 1, MIN_LOC_VEC_LEN, MAX_LOCAL_VEC_LEN>(batch_size, nloc, nchans, nlat_in, nlon_in, nlat_out, nlon_out, _kxp4, _vxp4, _qyp4, _dyp4, _row_idx, _row_off, _col_idx, _quad_weights, _dkxp4, _dvxp4, _dqyp4, stream); break;
+            case  512: launch_spc_attn_bwd<512, 1, MIN_LOC_VEC_LEN, MAX_LOCAL_VEC_LEN>(batch_size, nloc, nchans, nlat_in, nlon_in, nlat_out, nlon_out, _kxp4, _vxp4, _qyp4, _dyp4, _row_idx, _row_off, _col_idx, _quad_weights, _dkxp4, _dvxp4, _dqyp4, stream); break;
+            default:   launch_gen_attn_bwd                                            (batch_size,       nchans, nlat_in, nlon_in, nlat_out, nlon_out, _kxp4, _vxp4, _qyp4, _dyp4, _row_idx, _row_off, _col_idx, _quad_weights, _dkxp4, _dvxp4, _dqyp4, stream); break;
+        }
+    }
+
+    return;
+}
+
+// END backward kernels and functions
+
+std::tuple<at::Tensor, at::Tensor, at::Tensor> s2_attention_bwd_dkvq_cuda(at::Tensor kx, at::Tensor vx, at::Tensor qy,
+                                                                          at::Tensor dy, at::Tensor quad_weights,
+                                                                          at::Tensor psi_col_idx, at::Tensor psi_row_off,
+                                                                          int nlon_in, int nlat_out, int nlon_out)
+{
+
+    CHECK_CUDA_INPUT_TENSOR(kx);
+    CHECK_CUDA_INPUT_TENSOR(vx);
+    CHECK_CUDA_INPUT_TENSOR(qy);
+    CHECK_CUDA_INPUT_TENSOR(dy);
+    CHECK_CUDA_TENSOR(quad_weights);
+    CHECK_CUDA_TENSOR(psi_col_idx);
+    CHECK_CUDA_TENSOR(psi_row_off);
+    
+    const size_t uo_num_channels = kx.size(1);
+    const int batch_size = kx.size(0);
+
+    // extract dtype
+    auto kx_type = kx.dtype();
+    auto vx_type = vx.dtype();
+    auto qy_type = qy.dtype();
+    auto dy_type = dy.dtype();
+
+    torch::Tensor kxP = kx.to(torch::kFloat32);
+    torch::Tensor vxP = vx.to(torch::kFloat32);
+    torch::Tensor qyP = qy.to(torch::kFloat32);
+    torch::Tensor dyP = dy.to(torch::kFloat32);
+
+    // exract memory format: this is much safer than checking is_contiguous(at::MemoryFormat::ChannelsLast)
+    // the former fails for num_channels == 1
+    bool kx_is_channels_last = kxP.strides()[1] == 1;
+    bool vx_is_channels_last = vxP.strides()[1] == 1;
+    bool qy_is_channels_last = qyP.strides()[1] == 1;
+    bool dy_is_channels_last = dyP.strides()[1] == 1;
+
+    // transpose if required
+    if (!kx_is_channels_last) { kxP = permute_4D_to0231(kxP); }
+    if (!vx_is_channels_last) { vxP = permute_4D_to0231(vxP); }
+    if (!qy_is_channels_last) { qyP = permute_4D_to0231(qyP); }
+    if (!dy_is_channels_last) { dyP = permute_4D_to0231(dyP); }
+
+    torch::Tensor dkxP = torch::zeros_like(kxP);
+    torch::Tensor dvxP = torch::zeros_like(vxP);
+    torch::Tensor dqyP = torch::zeros_like(qyP);
+
+    s2_attn_bwd_dispatch(batch_size,
+                         uo_num_channels,
+                         nlon_in,
+                         nlat_out,
+                         nlon_out,
+                         kxP, vxP, qyP, dyP,
+                         psi_row_off,
+                         psi_col_idx,
+                         quad_weights,
+                         dkxP, dvxP, dqyP);
+
+    torch::Tensor dkx = dkxP;
+    torch::Tensor dvx = dvxP;
+    torch::Tensor dqy = dqyP;
+
+    if (!kx_is_channels_last) { dkx = permute_4D_to0312(dkx); }
+    if (!vx_is_channels_last) { dvx = permute_4D_to0312(dvx); }
+    if (!qy_is_channels_last) { dqy = permute_4D_to0312(dqy); }
+
+    // convert precision back to starting
+    dkx = dkx.to(kx_type);
+    dvx = dvx.to(vx_type);
+    dqy = dqy.to(qy_type);
+
+    return std::make_tuple(dkx, dvx, dqy);
+// #endif
+}

.deps/torch_harmonics/csrc/attention/attention_fwd_cuda.cu

@@ -0,0 +1,525 @@
+// coding=utf-8
+//
+// SPDX-FileCopyrightText: Copyright (c) 2025 The torch-harmonics Authors. All rights reserved.
+// SPDX-License-Identifier: BSD-3-Clause
+// 
+// Redistribution and use in source and binary forms, with or without
+// modification, are permitted provided that the following conditions are met:
+//
+// 1. Redistributions of source code must retain the above copyright notice, this
+// list of conditions and the following disclaimer.
+//
+// 2. Redistributions in binary form must reproduce the above copyright notice,
+// this list of conditions and the following disclaimer in the documentation
+// and/or other materials provided with the distribution.
+//
+// 3. Neither the name of the copyright holder nor the names of its
+// contributors may be used to endorse or promote products derived from
+// this software without specific prior written permission.
+//
+// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+// FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+// OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+#include "attention.cuh"
+#include <ATen/cuda/detail/TensorInfo.cuh>
+#include <ATen/cuda/detail/KernelUtils.h>
+#include <ATen/cuda/detail/IndexUtils.cuh>
+#include <ATen/cuda/CUDAUtils.h>
+
+#include <cuda_runtime.h>
+
+#include <cub/cub.cuh>
+#include <limits>
+
+#include "cudamacro.h"
+#include "attention_utils.cuh"
+
+#define THREADS (64)
+
+#define MAX_LOCAL_ARR_LEN (16)
+
+// BEGIN - forward kernels and functions
+
+// called with (blockDim.x=32 and blockDim.y>1, BDIM_X=blockDim.x*blockDim.y)
+template<int BDIM_X,
+         typename FLOATV_T> // either float or float4
+__global__
+__launch_bounds__(BDIM_X)
+void s2_attn_fwd_generic_vec_k(int nchan,  // no. of FLOATV_T elements along channel dim
+                               int nlat_in,
+                               int nlon_in,
+                               int nlat_out,
+                               int nlon_out,
+                               const FLOATV_T *__restrict__ kx,
+                               const FLOATV_T *__restrict__ vx,
+                               const FLOATV_T *__restrict__ qy,
+                               const int32_t *__restrict__ row_idx,
+                               const int64_t *__restrict__ row_off,
+                               const int64_t *__restrict__ col_idx,
+                               const   float *__restrict__ quad_weights,
+                                     FLOATV_T *__restrict__ y) {
+
+    extern __shared__ __align__(sizeof(float4)) float shext[];
+    FLOATV_T *shy = reinterpret_cast<FLOATV_T *>(shext) + threadIdx.y*nchan;
+
+    const int batch = blockIdx.y;
+    const int wid = blockIdx.x*blockDim.y + threadIdx.y;
+
+    if (wid >= nlat_out*nlon_out) {
+        return;
+    }
+
+    const int tidx = threadIdx.x;
+
+    const int h = wid / nlon_out;
+    const int wo = wid - (h*nlon_out);
+    const int ho = row_idx[h];
+
+    for(int chan = tidx; chan < nchan; chan += WARP_SIZE) {
+        shy[chan] = __vset<FLOATV_T>(0.f);
+    }
+
+    kx += int64_t(batch)*nlat_in*nlon_in*nchan;
+    vx += int64_t(batch)*nlat_in*nlon_in*nchan;
+    qy += int64_t(batch)*nlat_out*nlon_out*nchan + int64_t(ho)*nchan*nlon_out + int64_t(wo)*nchan;
+    y  += int64_t(batch)*nlat_out*nlon_out*nchan + int64_t(ho)*nchan*nlon_out + int64_t(wo)*nchan;
+
+    float alpha_sum = 0.0f;
+    float qdotk_max = -FLT_MAX;
+
+    const int64_t rbeg = row_off[ho];
+    const int64_t rend = row_off[ho+1];
+
+    col_idx += rbeg;
+
+    const int rlen = rend-rbeg;
+
+    for(int off = 0; off < rlen; off++) {
+
+        const int64_t col = col_idx[off];
+
+        const int hi = col / nlon_in;
+        const int wi = col - (hi*nlon_in);
+        const int wip = (wi+wo) - ((wi+wo) / nlon_in) * nlon_in;
+
+        const FLOATV_T *_kx = kx + int64_t(hi)*nlon_in*nchan + int64_t(wip)*nchan;
+        const FLOATV_T *_vx = vx + int64_t(hi)*nlon_in*nchan + int64_t(wip)*nchan;
+
+        FLOATV_T qdotkv = __vset<FLOATV_T>(0.f);
+
+        for(int chan = tidx; chan < nchan; chan += WARP_SIZE) {
+            qdotkv = __vadd(qdotkv,
+                            __vmul( qy[chan],
+                                   _kx[chan]));
+        }
+
+        float qdotk = __warp_sum(__vred(qdotkv));
+
+        float qdotk_max_tmp;
+        float alpha;
+        float exp_save;
+
+        qdotk_max_tmp = max(qdotk_max, qdotk);
+        alpha = expf(qdotk - qdotk_max_tmp) * quad_weights[hi];
+        exp_save = expf(qdotk_max - qdotk_max_tmp);
+
+        alpha_sum = alpha + alpha_sum*exp_save;
+
+        for(int chan = tidx; chan < nchan; chan += WARP_SIZE) {
+            shy[chan] = __vadd(__vscale(exp_save, shy[chan]),
+                               __vscale(   alpha, _vx[chan]));
+        }
+        qdotk_max = qdotk_max_tmp;
+    }
+
+    alpha_sum = 1.0f / alpha_sum;
+    for(int chan = tidx; chan < nchan; chan += WARP_SIZE) {
+        y[chan] = __vscale(alpha_sum, shy[chan]);
+    }
+
+    return;
+}
+
+// called with either (BDIM_X=32 and BDIM_Y>1) || (2^K=BDIM_X > 32 and BDIM_Y=1)
+template<int BDIM_X,
+         int BDIM_Y,
+         int NLOC,
+         typename FLOATV_T> // either float or float4
+__global__
+__launch_bounds__(BDIM_X*BDIM_Y)
+void s2_attn_fwd_special_vec_k(int nchan, // no. of FLOATV_T elements along channel dim
+                               int nlat_in,
+                               int nlon_in,
+                               int nlat_out,
+                               int nlon_out,
+                               const FLOATV_T *__restrict__ kx,
+                               const FLOATV_T *__restrict__ vx,
+                               const FLOATV_T *__restrict__ qy,
+                               const int32_t *__restrict__ row_idx,
+                               const int64_t *__restrict__ row_off,
+                               const int64_t *__restrict__ col_idx,
+                               const   float *__restrict__ quad_weights,
+                                     FLOATV_T *__restrict__ y) {
+
+    static_assert(0 == (BDIM_X & (BDIM_X-1)));
+    static_assert(0 == (BDIM_Y & (BDIM_Y-1)));
+    static_assert((BDIM_X == 32 && BDIM_Y  > 1) ||
+                  (BDIM_X  > 32 && BDIM_Y == 1)) ;
+
+    constexpr int NLOC_M1 = NLOC-1;
+
+    const int tidx = threadIdx.x;
+    const int batch = blockIdx.y;
+    const int ctaid = blockIdx.x*blockDim.y + threadIdx.y;
+
+    if (ctaid >= nlat_out*nlon_out) {
+        return;
+    }
+
+    FLOATV_T locy[NLOC];
+
+    extern __shared__ __align__(sizeof(float4)) float shext[];
+    FLOATV_T *shq = reinterpret_cast<FLOATV_T *>(shext) + threadIdx.y*nchan + tidx;
+
+    const int h = ctaid / nlon_out;
+    const int wo = ctaid - (h*nlon_out);
+    const int ho = row_idx[h];
+
+    kx += int64_t(batch)*nlat_in*nlon_in*nchan + tidx;
+    vx += int64_t(batch)*nlat_in*nlon_in*nchan + tidx;
+    qy += int64_t(batch)*nlat_out*nlon_out*nchan + int64_t(ho)*nlon_out*nchan + int64_t(wo)*nchan + tidx;
+    y  += int64_t(batch)*nlat_out*nlon_out*nchan + int64_t(ho)*nlon_out*nchan + int64_t(wo)*nchan + tidx;
+
+    #pragma unroll
+    for(int i = 0; i < NLOC; i++) {
+        locy[i] = __vset<FLOATV_T>(0.f);
+    }
+
+    #pragma unroll
+    for(int i = 0; i < NLOC_M1; i++) {
+        shq[i*BDIM_X] = qy[i*BDIM_X];
+    }
+    if (NLOC_M1*BDIM_X+tidx < nchan) {
+        shq[NLOC_M1*BDIM_X] = qy[NLOC_M1*BDIM_X];
+    }
+
+    float alpha_sum = 0.0f;
+    float qdotk_max = -FLT_MAX;
+
+    const int64_t rbeg = row_off[ho];
+    const int64_t rend = row_off[ho+1];
+
+    col_idx += rbeg;
+
+    const int rlen = rend-rbeg;
+
+    for(int off = 0; off < rlen; off++) {
+
+        const int64_t col = col_idx[off];
+
+        const int hi = col / nlon_in;
+        const int wi = col - (hi*nlon_in);
+        const int wip = (wi+wo) - ((wi+wo) / nlon_in) * nlon_in;
+
+        const FLOATV_T *_kx = kx + int64_t(hi)*nlon_in*nchan + int64_t(wip)*nchan;
+        const FLOATV_T *_vx = vx + int64_t(hi)*nlon_in*nchan + int64_t(wip)*nchan;
+
+        FLOATV_T qdotkv = __vset<FLOATV_T>(0.f);
+
+        #pragma unroll
+        for(int i = 0; i < NLOC_M1; i++) {
+            qdotkv = __vadd(qdotkv,
+                            __vmul(shq[i*BDIM_X],
+                                   _kx[i*BDIM_X]));
+        }
+        if (NLOC_M1*BDIM_X+tidx < nchan) {
+            qdotkv = __vadd(qdotkv,
+                            __vmul(shq[NLOC_M1*BDIM_X],
+                                   _kx[NLOC_M1*BDIM_X]));
+        }
+
+        float qdotk = __vred(qdotkv);
+        if constexpr(BDIM_X == 32) { qdotk =          __warp_sum(qdotk); }
+        else                       { qdotk = __block_sum<BDIM_X>(qdotk); }
+
+        float qdotk_max_tmp;
+        float alpha;
+        float exp_save;
+
+        qdotk_max_tmp = max(qdotk_max, qdotk);
+        alpha = expf(qdotk - qdotk_max_tmp) * quad_weights[hi];
+        exp_save = expf(qdotk_max - qdotk_max_tmp);
+
+        alpha_sum = alpha + alpha_sum*exp_save;
+
+        #pragma unroll
+        for(int i = 0; i < NLOC_M1; i++) {
+            locy[i] = __vadd(__vscale(exp_save, locy[i]),
+                             __vscale(alpha, _vx[i*BDIM_X]));
+        }
+        if (NLOC_M1*BDIM_X+tidx < nchan) {
+            locy[NLOC_M1] = __vadd(__vscale(exp_save, locy[NLOC_M1]),
+                                   __vscale(alpha, _vx[NLOC_M1*BDIM_X]));
+        }
+
+        qdotk_max = qdotk_max_tmp;
+    }
+
+    alpha_sum = 1.0f / alpha_sum;
+
+    #pragma unroll
+    for(int i = 0; i < NLOC_M1; i++) {
+        y[i*BDIM_X] = __vscale(alpha_sum, locy[i]);
+    }
+    if (NLOC_M1*BDIM_X+tidx < nchan) {
+        y[NLOC_M1*BDIM_X] = __vscale(alpha_sum, locy[NLOC_M1]);
+    }
+
+    return;
+}
+
+template<typename FLOATV_T>
+void launch_gen_attn_fwd(int batch_size,
+                         int nchans,
+                         int nlat_in,
+                         int nlon_in,
+                         int nlat_out,
+                         int nlon_out,
+                         FLOATV_T *__restrict__ _kxp,
+                         FLOATV_T *__restrict__ _vxp,
+                         FLOATV_T *__restrict__ _qyp,
+                         int32_t *_row_idx,
+                         int64_t *_row_off,
+                         int64_t *_col_idx,
+                         float *_quad_weights,
+                         FLOATV_T *__restrict__ _yp,
+                         cudaStream_t stream) {
+
+    dim3 block(WARP_SIZE, THREADS/WARP_SIZE);
+    dim3 grid(DIV_UP(nlat_out*nlon_out, block.y), batch_size);
+
+    size_t shsize = sizeof(FLOATV_T)*nchans * block.y;
+
+    s2_attn_fwd_generic_vec_k<THREADS>
+                             <<<grid, block, shsize, stream>>>(nchans, nlat_in, nlon_in, nlat_out, nlon_out,
+                                                               _kxp, _vxp, _qyp, _row_idx, _row_off, _col_idx, _quad_weights, _yp);
+    CHECK_ERROR("s2_attn_fwd_generic_vec_k");
+
+    return;
+}
+
+template<int BDIM_X,
+         int BDIM_Y,
+         int CUR_LOC_SIZE,
+         int MAX_LOC_SIZE, // max size of FLOATV_T[] local array
+         typename FLOATV_T>
+void launch_spc_attn_fwd(int batch_size,
+                         int nloc,      // "BDIM_X*nloc" >= nchans
+                         int nchans,
+                         int nlat_in,
+                         int nlon_in,
+                         int nlat_out,
+                         int nlon_out,
+                         FLOATV_T *__restrict__ _kxp,
+                         FLOATV_T *__restrict__ _vxp,
+                         FLOATV_T *__restrict__ _qyp,
+                         int32_t *_row_idx,
+                         int64_t *_row_off,
+                         int64_t *_col_idx,
+                         float *_quad_weights,
+                         FLOATV_T *__restrict__ _yp,
+                         cudaStream_t stream) {
+
+    if (CUR_LOC_SIZE == nloc) {
+
+        dim3 block(BDIM_X, BDIM_Y);
+        dim3 grid(DIV_UP(nlat_out*nlon_out, block.y), batch_size);
+
+        size_t shsize = sizeof(FLOATV_T)*nchans * block.y; // block.y > 1 iif block.x==32
+
+        s2_attn_fwd_special_vec_k<BDIM_X, BDIM_Y, CUR_LOC_SIZE>
+                                 <<<grid, block, shsize, stream>>>(nchans, nlat_in, nlon_in, nlat_out, nlon_out,
+                                                                   _kxp, _vxp, _qyp, _row_idx, _row_off, _col_idx, _quad_weights, _yp);
+        CHECK_ERROR("s2_attn_fwd_special_vec_k");
+
+        return;
+    }
+    if constexpr(CUR_LOC_SIZE < MAX_LOC_SIZE) {
+        launch_spc_attn_fwd<BDIM_X,
+                             BDIM_Y,
+                             CUR_LOC_SIZE+1,
+                             MAX_LOC_SIZE>(batch_size, nloc, nchans, nlat_in, nlon_in, nlat_out, nlon_out,
+                                           _kxp, _vxp, _qyp, _row_idx, _row_off, _col_idx, _quad_weights, _yp,
+                                           stream);
+    }
+    return;
+}
+
+static void s2_attn_fwd_dispatch(int batch_size,
+                                 int nchans,
+                                 int nlon_in,
+                                 int nlat_out,
+                                 int nlon_out,
+                                 at::Tensor kxP,
+                                 at::Tensor vxP,
+                                 at::Tensor qyP,
+                                 at::Tensor row_off,
+                                 at::Tensor col_idx,
+                                 at::Tensor quad_weights,
+                                 at::Tensor yP) {
+
+    static_assert(0 == (MAX_LOCAL_ARR_LEN & (MAX_LOCAL_ARR_LEN-1)));
+
+    // get stream
+    auto stream = at::cuda::getCurrentCUDAStream().stream();
+
+    // sort row indices (ho-s) in descending order
+    // based on (row_off[ho+1]-row_off[ho])
+    at::Tensor row_idx = sortRows(nlat_out, row_off, stream);
+
+    const int nlat_in = kxP.size(1);
+
+    // smallest power of two "bdimx" (>=32) s.t. bdimx*MAX_LOCAL_ARR_LEN >= nchans
+    int bdimx;
+    bdimx = DIV_UP(nchans, MAX_LOCAL_ARR_LEN);
+    bdimx = max(bdimx, WARP_SIZE);
+    bdimx = next_pow2(bdimx);
+
+    float *_kxp = reinterpret_cast<float *>(kxP.data_ptr());
+    float *_vxp = reinterpret_cast<float *>(vxP.data_ptr());
+    float *_qyp = reinterpret_cast<float *>(qyP.data_ptr());
+    float *_yp  = reinterpret_cast<float *>(yP.data_ptr());
+
+    int32_t *_row_idx = reinterpret_cast<int32_t *>(row_idx.data_ptr());
+    int64_t *_row_off = reinterpret_cast<int64_t *>(row_off.data_ptr());
+    int64_t *_col_idx = reinterpret_cast<int64_t *>(col_idx.data_ptr());
+    float *_quad_weights = reinterpret_cast<float *>(quad_weights.data_ptr());
+
+    constexpr int VEC_SIZE = sizeof(float4) / sizeof(float);
+
+    if (!is_aligned<sizeof(float4)>(_kxp) ||
+        !is_aligned<sizeof(float4)>(_vxp) ||
+        !is_aligned<sizeof(float4)>(_qyp) ||
+        !is_aligned<sizeof(float4)>(_yp)  ||
+        (nchans % VEC_SIZE) != 0) {
+
+        const int nloc = DIV_UP(nchans, bdimx);
+        
+        // to avoid the compilation of unused template instances;
+        // we use a block size BDIM_X that is the smallest power of 2
+        // such that BDIM_X*MAX_LOCAL_ARR_LEN >= nchans, so
+        // BDIM_X > 32 are used only for:
+        //
+        //  (BDIM_X-1)*MAX_LOCAL_ARR_LEN < nchans <= BDIM_X*MAX_LOCAL_ARR_LEN
+        constexpr int MIN_LOC_ARR_LEN = MAX_LOCAL_ARR_LEN/2+1;
+
+        // use 2D blocks only if 32 threads are enough
+        switch(bdimx) {
+            case   32: launch_spc_attn_fwd<  32, 2,               1, MAX_LOCAL_ARR_LEN>(batch_size, nloc, nchans, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, _vxp, _qyp, _row_idx, _row_off, _col_idx, _quad_weights, _yp, stream); break;
+            case   64: launch_spc_attn_fwd<  64, 1, MIN_LOC_ARR_LEN, MAX_LOCAL_ARR_LEN>(batch_size, nloc, nchans, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, _vxp, _qyp, _row_idx, _row_off, _col_idx, _quad_weights, _yp, stream); break;
+            case  128: launch_spc_attn_fwd< 128, 1, MIN_LOC_ARR_LEN, MAX_LOCAL_ARR_LEN>(batch_size, nloc, nchans, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, _vxp, _qyp, _row_idx, _row_off, _col_idx, _quad_weights, _yp, stream); break;
+            case  256: launch_spc_attn_fwd< 256, 1, MIN_LOC_ARR_LEN, MAX_LOCAL_ARR_LEN>(batch_size, nloc, nchans, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, _vxp, _qyp, _row_idx, _row_off, _col_idx, _quad_weights, _yp, stream); break;
+            case  512: launch_spc_attn_fwd< 512, 1, MIN_LOC_ARR_LEN, MAX_LOCAL_ARR_LEN>(batch_size, nloc, nchans, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, _vxp, _qyp, _row_idx, _row_off, _col_idx, _quad_weights, _yp, stream); break;
+            case 1024: launch_spc_attn_fwd<1024, 1, MIN_LOC_ARR_LEN, MAX_LOCAL_ARR_LEN>(batch_size, nloc, nchans, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, _vxp, _qyp, _row_idx, _row_off, _col_idx, _quad_weights, _yp, stream); break;
+            default:   launch_gen_attn_fwd                                             (batch_size,       nchans, nlat_in, nlon_in, nlat_out, nlon_out, _kxp, _vxp, _qyp, _row_idx, _row_off, _col_idx, _quad_weights, _yp, stream); break;
+        }
+
+    } else {
+
+        float4 *_kxp4 = reinterpret_cast<float4 *>(_kxp);
+        float4 *_vxp4 = reinterpret_cast<float4 *>(_vxp);
+        float4 *_qyp4 = reinterpret_cast<float4 *>(_qyp);
+        float4 *_yp4  = reinterpret_cast<float4 *>(_yp);
+
+        nchans /= VEC_SIZE;
+        const int nloc = DIV_UP(nchans, bdimx);
+
+        constexpr int MAX_LOCAL_VEC_LEN = MAX_LOCAL_ARR_LEN / VEC_SIZE;
+        
+        constexpr int MIN_LOC_VEC_LEN = MAX_LOCAL_VEC_LEN/2+1;
+
+        // use 2D blocks only if 32 threads are enough
+        switch(bdimx) {
+            case   32: launch_spc_attn_fwd<  32, 2,               1, MAX_LOCAL_VEC_LEN>(batch_size, nloc, nchans, nlat_in, nlon_in, nlat_out, nlon_out, _kxp4, _vxp4, _qyp4, _row_idx, _row_off, _col_idx, _quad_weights, _yp4, stream); break;
+            case   64: launch_spc_attn_fwd<  64, 1, MIN_LOC_VEC_LEN, MAX_LOCAL_VEC_LEN>(batch_size, nloc, nchans, nlat_in, nlon_in, nlat_out, nlon_out, _kxp4, _vxp4, _qyp4, _row_idx, _row_off, _col_idx, _quad_weights, _yp4, stream); break;
+            case  128: launch_spc_attn_fwd< 128, 1, MIN_LOC_VEC_LEN, MAX_LOCAL_VEC_LEN>(batch_size, nloc, nchans, nlat_in, nlon_in, nlat_out, nlon_out, _kxp4, _vxp4, _qyp4, _row_idx, _row_off, _col_idx, _quad_weights, _yp4, stream); break;
+            case  256: launch_spc_attn_fwd< 256, 1, MIN_LOC_VEC_LEN, MAX_LOCAL_VEC_LEN>(batch_size, nloc, nchans, nlat_in, nlon_in, nlat_out, nlon_out, _kxp4, _vxp4, _qyp4, _row_idx, _row_off, _col_idx, _quad_weights, _yp4, stream); break;
+            case  512: launch_spc_attn_fwd< 512, 1, MIN_LOC_VEC_LEN, MAX_LOCAL_VEC_LEN>(batch_size, nloc, nchans, nlat_in, nlon_in, nlat_out, nlon_out, _kxp4, _vxp4, _qyp4, _row_idx, _row_off, _col_idx, _quad_weights, _yp4, stream); break;
+            case 1024: launch_spc_attn_fwd<1024, 1, MIN_LOC_VEC_LEN, MAX_LOCAL_VEC_LEN>(batch_size, nloc, nchans, nlat_in, nlon_in, nlat_out, nlon_out, _kxp4, _vxp4, _qyp4, _row_idx, _row_off, _col_idx, _quad_weights, _yp4, stream); break;
+            default:   launch_gen_attn_fwd                                             (batch_size,       nchans, nlat_in, nlon_in, nlat_out, nlon_out, _kxp4, _vxp4, _qyp4, _row_idx, _row_off, _col_idx, _quad_weights, _yp4, stream); break;
+        }
+    }
+
+    return;
+}
+
+// END - forward kernels and functions
+
+torch::Tensor s2_attention_fwd_cuda(at::Tensor kx,
+                                    at::Tensor vx,
+                                    at::Tensor qy,
+                                    at::Tensor quad_weights,
+                                    at::Tensor psi_col_idx,
+                                    at::Tensor psi_row_off,
+                                    int nlon_in,
+                                    int nlat_out,
+                                    int nlon_out) {
+    CHECK_CUDA_INPUT_TENSOR(kx);
+    CHECK_CUDA_INPUT_TENSOR(vx);
+    CHECK_CUDA_INPUT_TENSOR(qy);
+    CHECK_CUDA_TENSOR(quad_weights);
+    CHECK_CUDA_TENSOR(psi_col_idx);
+    CHECK_CUDA_TENSOR(psi_row_off);
+
+    size_t uo_num_channels = kx.size(1);
+
+    const int batch_size = kx.size(0);
+
+     // extract dtype
+     auto qy_type = qy.dtype();
+
+    torch::Tensor kxP = kx.to(torch::kFloat32);
+    torch::Tensor vxP = vx.to(torch::kFloat32);
+    torch::Tensor qyP = qy.to(torch::kFloat32);
+
+    // these are much safer than checking is_contiguous(at::MemoryFormat::ChannelsLast)
+    // the former fails for num_channels == 1
+    bool kx_is_channels_last = kxP.strides()[1] == 1;
+    bool vx_is_channels_last = vxP.strides()[1] == 1;
+    bool qy_is_channels_last = qyP.strides()[1] == 1;
+
+    if (!kx_is_channels_last) { kxP = permute_4D_to0231(kxP); }
+    if (!vx_is_channels_last) { vxP = permute_4D_to0231(vxP); }
+    if (!qy_is_channels_last) { qyP = permute_4D_to0231(qyP); }
+
+    torch::Tensor yP = torch::empty_like(qyP);
+
+    s2_attn_fwd_dispatch(batch_size,
+                         uo_num_channels,
+                         nlon_in,
+                         nlat_out,
+                         nlon_out,
+                         kxP, vxP, qyP,
+                         psi_row_off,
+                         psi_col_idx,
+                         quad_weights,
+                         yP);
+
+    torch::Tensor y = yP;
+    if (!qy_is_channels_last) { y = permute_4D_to0312(y); }
+
+    // convert precision back to starting
+    y = y.to(qy_type);
+
+    C10_CUDA_KERNEL_LAUNCH_CHECK();
+
+    return y;
+}

.deps/torch_harmonics/csrc/attention/attention_interface.cu

@@ -0,0 +1,38 @@
+// coding=utf-8
+//
+// SPDX-FileCopyrightText: Copyright (c) 2025 The torch-harmonics Authors. All rights reserved.
+// SPDX-License-Identifier: BSD-3-Clause
+//
+// Redistribution and use in source and binary forms, with or without
+// modification, are permitted provided that the following conditions are met:
+//
+// 1. Redistributions of source code must retain the above copyright notice, this
+// list of conditions and the following disclaimer.
+//
+// 2. Redistributions in binary form must reproduce the above copyright notice,
+// this list of conditions and the following disclaimer in the documentation
+// and/or other materials provided with the distribution.
+//
+// 3. Neither the name of the copyright holder nor the names of its
+// contributors may be used to endorse or promote products derived from
+// this software without specific prior written permission.
+//
+// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+// FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+// OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+#include "attention.cuh"
+#include <torch/extension.h>
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
+{
+    m.def("forward", &s2_attention_fwd_cuda, "(Local) Attention on S2");
+    m.def("backward_dkvq", &s2_attention_bwd_dkvq_cuda, "(Local) Attention gradient on S2 (gradient for k,v,&q)");
+}

.deps/torch_harmonics/csrc/attention/attention_utils.cu

@@ -0,0 +1,180 @@
+// coding=utf-8
+//
+// SPDX-FileCopyrightText: Copyright (c) 2025 The torch-harmonics Authors. All rights reserved.
+// SPDX-License-Identifier: BSD-3-Clause
+// 
+// Redistribution and use in source and binary forms, with or without
+// modification, are permitted provided that the following conditions are met:
+//
+// 1. Redistributions of source code must retain the above copyright notice, this
+// list of conditions and the following disclaimer.
+//
+// 2. Redistributions in binary form must reproduce the above copyright notice,
+// this list of conditions and the following disclaimer in the documentation
+// and/or other materials provided with the distribution.
+//
+// 3. Neither the name of the copyright holder nor the names of its
+// contributors may be used to endorse or promote products derived from
+// this software without specific prior written permission.
+//
+// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+// FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+// OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+#include "attention.cuh"
+#include <ATen/cuda/detail/TensorInfo.cuh>
+#include <ATen/cuda/detail/KernelUtils.h>
+#include <ATen/cuda/detail/IndexUtils.cuh>
+#include <ATen/cuda/CUDAUtils.h>
+
+#include <cuda_runtime.h>
+
+#include <cub/cub.cuh>
+#include <limits>
+
+#include "cudamacro.h"
+#include "attention_utils.cuh"
+
+#define THREADS (64)
+
+#define TRANSP_WARPS_X_TILE_GENERIC (32)
+#define TRANSP_WARPS_X_TILE_SM100    (4)
+
+// BEGIN - CSR rows sorting kernels and functions
+__global__ void set_rlen_rids_k(const int n,
+                                const int64_t *__restrict__ offs,
+                                      int *__restrict__ rids,
+                                      int *__restrict__ rlen) {
+
+    const int nth = gridDim.x*blockDim.x;
+    const int tid = blockIdx.x*blockDim.x + threadIdx.x;
+
+    for(int i = tid; i < n; i += nth) {
+        rids[i] = i;
+        rlen[i] = offs[i+1]-offs[i];
+    }
+
+    return;
+}   
+
+at::Tensor sortRows(int nlat_out, at::Tensor row_off, cudaStream_t stream) {
+
+    int64_t *_row_off_d = reinterpret_cast<int64_t *>(row_off.data_ptr());
+
+    auto options = torch::TensorOptions().dtype(torch::kInt32).device(row_off.device());
+
+    torch::Tensor rids_d = torch::empty({nlat_out}, options);
+    torch::Tensor rlen_d = torch::empty({nlat_out}, options);
+
+    int *_rids_d = reinterpret_cast<int *>(rids_d.data_ptr());
+    int *_rlen_d = reinterpret_cast<int *>(rlen_d.data_ptr());
+
+    const int grid = DIV_UP(nlat_out, THREADS);
+    const int block = THREADS;
+
+    set_rlen_rids_k<<<grid, block, 0, stream>>>(nlat_out,
+                                                _row_off_d,
+                                                _rids_d,
+                                                _rlen_d);
+
+    torch::Tensor rids_sort_d = torch::empty({nlat_out}, options);
+    torch::Tensor rlen_sort_d = torch::empty({nlat_out}, options);
+
+    int *_rids_sort_d = reinterpret_cast<int *>(rids_sort_d.data_ptr());
+    int *_rlen_sort_d = reinterpret_cast<int *>(rlen_sort_d.data_ptr());
+
+    size_t temp_storage_bytes = 0;
+    CHECK_CUDA(cub::DeviceRadixSort::SortPairsDescending(NULL, temp_storage_bytes,
+                                                         _rlen_d, _rlen_sort_d, 
+                                                         _rids_d, _rids_sort_d,
+                                                         nlat_out, 0, sizeof(*_rlen_d)*8, stream));
+
+    options = torch::TensorOptions().dtype(torch::kByte).device(row_off.device());
+    torch::Tensor temp_storage_d = torch::empty({int64_t(temp_storage_bytes)}, options);
+
+    void *_temp_storage_d = reinterpret_cast<void *>(temp_storage_d.data_ptr());
+
+    CHECK_CUDA(cub::DeviceRadixSort::SortPairsDescending(_temp_storage_d, temp_storage_bytes,
+                                                         _rlen_d, _rlen_sort_d, 
+                                                         _rids_d, _rids_sort_d,
+                                                         nlat_out, 0, sizeof(*_rlen_d)*8, stream));
+    return rids_sort_d;
+}
+// END - CSR rows sorting kernels and functions
+
+
+// BEGIN - 4D tensor permutation kernels and functions
+__global__ void empty_k() {}
+
+static int getPtxver() {
+    cudaFuncAttributes attrs;
+    CHECK_CUDA(cudaFuncGetAttributes(&attrs, empty_k));
+    return attrs.ptxVersion*10;
+}
+
+at::Tensor permute_4D_to0231(at::Tensor src) {
+
+    auto options = torch::TensorOptions().dtype(src.dtype()).device(src.device());
+    torch::Tensor dst = torch::empty({src.size(0), src.size(2), src.size(3), src.size(1)}, options);
+
+    const int ptxv = getPtxver();
+
+    // to be further specialized for additional archs, if necessary
+    if (ptxv < 100) {
+        AT_DISPATCH_FLOATING_TYPES(src.scalar_type(), "permute_to0231_k_tile_generic", ([&] {
+            launch_permute_to0231<TRANSP_WARPS_X_TILE_GENERIC, scalar_t>(src, dst);
+        }));
+        CHECK_ERROR("permute_to0231_k_tile_generic");
+    } else {
+        AT_DISPATCH_FLOATING_TYPES(src.scalar_type(), "permute_to0231_k_tile_sm100", ([&] {
+            launch_permute_to0231<TRANSP_WARPS_X_TILE_SM100, scalar_t>(src, dst);
+        }));
+        CHECK_ERROR("permute_to0231_k_tile_sm100");
+    }
+
+    return dst;
+}
+
+at::Tensor permute_4D_to0312(at::Tensor src) {
+
+    auto options = torch::TensorOptions().dtype(src.dtype()).device(src.device());
+    torch::Tensor dst = torch::empty({src.size(0), src.size(3), src.size(1), src.size(2)}, options);
+
+    const int ptxv = getPtxver();
+
+    // to be further specialized for additional archs, if necessary
+    if (ptxv < 100) {
+        AT_DISPATCH_FLOATING_TYPES(src.scalar_type(), "permute_to0312_k_tile_generic", ([&] {
+            launch_permute_to0312<TRANSP_WARPS_X_TILE_GENERIC, scalar_t>(src, dst);
+        }));
+        CHECK_ERROR("permute_to0312_k_tile_generic");
+    } else {
+        AT_DISPATCH_FLOATING_TYPES(src.scalar_type(), "permute_to0312_k_tile_sm100", ([&] {
+            launch_permute_to0312<TRANSP_WARPS_X_TILE_SM100, scalar_t>(src, dst);
+        }));
+        CHECK_ERROR("permute_to0312_k_tile_sm100");
+    }
+
+    return dst;
+}
+// END - tensor permutation kernels and functions
+
+// BEGIN - general host-side functions
+unsigned int next_pow2(unsigned int x) { 
+
+    x -= 1;
+
+    #pragma unroll
+    for(int i = 1; i <= sizeof(x)*8 / 2; i *= 2) {
+        x |= x >> i;    
+    }
+    return x+1;
+}
+// END - general host-side functions

.deps/torch_harmonics/csrc/attention/attention_utils.cuh

@@ -0,0 +1,373 @@
+// coding=utf-8
+//
+// SPDX-FileCopyrightText: Copyright (c) 2025 The torch-harmonics Authors. All rights reserved.
+// SPDX-License-Identifier: BSD-3-Clause
+//
+// Redistribution and use in source and binary forms, with or without
+// modification, are permitted provided that the following conditions are met:
+//
+// 1. Redistributions of source code must retain the above copyright notice, this
+// list of conditions and the following disclaimer.
+//
+// 2. Redistributions in binary form must reproduce the above copyright notice,
+// this list of conditions and the following disclaimer in the documentation
+// and/or other materials provided with the distribution.
+//
+// 3. Neither the name of the copyright holder nor the names of its
+// contributors may be used to endorse or promote products derived from
+// this software without specific prior written permission.
+//
+// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+// FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+// OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+#pragma once
+
+#include <ATen/ATen.h>
+
+#define WARP_SIZE (32)
+#define FULL_MASK (0xFFFFFFFF)
+#define DIV_UP(a,b) (((a)+((b)-1))/(b))
+
+// CSR rows sorting kernels and functions
+at::Tensor sortRows(int nlat_out, at::Tensor row_off, cudaStream_t stream);
+
+// 4D tensor permutation kernels and functions
+at::Tensor permute_4D_to0231(at::Tensor src);
+at::Tensor permute_4D_to0312(at::Tensor src);
+
+// Host tensor dump and CSR manipulation functions
+void dump_tensor(const char *fname, at::Tensor t);
+void dump_csr(const char *fname, at::Tensor roff, at::Tensor cols);
+
+int part_csr_rows(int *row_perm,
+                  const at::Tensor roff,
+                  const at::Tensor cols,
+                  int **part_off,
+                  int **part_val);
+
+int verify_part(const int npart,
+                const int *part_off,
+                const int *part_val,
+                const at::Tensor roff,
+                const at::Tensor cols);
+
+void verify_part_new(const int nlon_out,
+                const int nlat_in,
+                const int nlon_in,
+                const int npart,      // partitioning data
+                const int *part_off,
+                const int *part_val,
+                const at::Tensor roff, 
+                const at::Tensor cols);
+
+unsigned int next_pow2(unsigned int x);
+
+
+// utility host functions and templates
+
+template<unsigned int ALIGN>
+int is_aligned(const void *ptr) {
+
+    static_assert(0 == (ALIGN & (ALIGN-1)));
+    return (0 == (uintptr_t(ptr) & (ALIGN-1)));
+}
+
+
+// utility device functions and templates
+
+template<typename FLOATV_T>
+__device__ FLOATV_T __vset(float x) {
+    static_assert(sizeof(FLOATV_T) == 0, "Unsupported type for __vset");
+    return FLOATV_T{};
+}
+
+template<>
+__device__ float __forceinline__ __vset<float>(float x) {
+    return x;
+}
+
+__device__ float __forceinline__ __vmul(float a, float b) {
+    return a*b;
+}
+
+__device__ float __forceinline__ __vadd(float a, float b) {
+    return a+b;
+}
+
+__device__ float __forceinline__ __vsub(float a, float b) {
+    return a-b;
+}
+
+__device__ float __forceinline__ __vred(float a) {
+    return a;
+}
+
+__device__ float __forceinline__ __vscale(float s, float v) {
+    return v*s;
+}
+
+__device__ float __forceinline__ __vdiv(float s, float v) {
+    return v/s;
+}
+
+template<>
+__device__ float4 __forceinline__ __vset<float4>(float x) {
+    return make_float4(x, x, x, x);
+}
+
+__device__ float4 __forceinline__ __vmul(float4 a, float4 b) {
+    return make_float4(a.x*b.x, a.y*b.y, a.z*b.z, a.w*b.w);
+}
+
+__device__ float4 __forceinline__ __vadd(float4 a, float4 b) {
+    return make_float4(a.x+b.x, a.y+b.y, a.z+b.z, a.w+b.w);
+}
+
+__device__ float4 __forceinline__ __vsub(float4 a, float4 b) {
+    return make_float4(a.x-b.x, a.y-b.y, a.z-b.z, a.w-b.w);
+}
+
+__device__ float __forceinline__ __vred(float4 a) {
+    return a.x + a.y + a.z + a.w;
+}
+
+__device__ float4 __forceinline__ __vscale(float s, float4 v) {
+    return make_float4(s*v.x, s*v.y, s*v.z, s*v.w);
+}
+
+__device__ float4 __forceinline__ __vdiv(float s, float4 v) {
+    return make_float4(s/v.x, s/v.y, s/v.z, s/v.w);;
+}
+
+template<typename VAL_T>
+__device__ VAL_T __warp_sum(VAL_T val) {
+
+    #pragma unroll
+    for(int i = WARP_SIZE/2; i; i /= 2) {
+        val += __shfl_xor_sync(FULL_MASK, val, i, WARP_SIZE);
+    }
+    return val;
+}
+
+template<int BDIM_X,
+         int BDIM_Y=1,
+         int BDIM_Z=1,
+         typename VAL_T>
+__device__ VAL_T __block_sum(VAL_T val) {
+
+    const int NWARP = (BDIM_X*BDIM_Y*BDIM_Z) / WARP_SIZE;
+
+    val = __warp_sum(val);
+
+    if constexpr(NWARP > 1) {
+
+        int tid = threadIdx.x;
+        if constexpr(BDIM_Y > 1) { tid += threadIdx.y*BDIM_X;        }
+        if constexpr(BDIM_Z > 1) { tid += threadIdx.z*BDIM_X*BDIM_Y; }
+
+        const int lid = tid%WARP_SIZE;
+        const int wid = tid/WARP_SIZE;
+
+        __shared__ VAL_T sh[NWARP];
+
+        if (lid == 0) {
+            sh[wid] = val;
+        }
+        __syncthreads();
+
+        if (wid == 0) {
+            val = (lid < NWARP) ? sh[lid] : 0;
+
+            val = __warp_sum(val);
+            __syncwarp();
+
+            if (!lid) {
+                sh[0] = val;
+            }
+        }
+        __syncthreads();
+
+        val = sh[0];
+        __syncthreads();
+    }
+    return val;
+}
+
+// transpose utils
+template<int BDIM_X,
+         int BDIM_Y,
+         typename VAL_T>
+__global__
+__launch_bounds__(BDIM_X*BDIM_Y)
+void  permute_to0231_k(const int nchn,
+                       const int nlat,
+                       const int nlon,
+                       const at::PackedTensorAccessor32<VAL_T, 4, at::RestrictPtrTraits> src,
+                             at::PackedTensorAccessor32<VAL_T, 4, at::RestrictPtrTraits> dst) {
+
+    static_assert(!(BDIM_X & (BDIM_X-1)));
+    static_assert(!(BDIM_Y & (BDIM_Y-1)));
+    static_assert(BDIM_X >= BDIM_Y);
+
+    __shared__ VAL_T sh[BDIM_X][BDIM_X+1];
+
+    const int tidx = threadIdx.x;
+    const int tidy = threadIdx.y;
+
+    const int coff  = blockIdx.x*BDIM_X;           // channel offset
+    const int woff  = blockIdx.y*BDIM_X;           // width offset
+    const int batch = blockIdx.z / nlat;           // batch (same for all block)
+    const int h     = blockIdx.z - (batch * nlat); // height (same for all block)
+
+    const int nchn_full = (nchn-coff) >= BDIM_X;
+    const int nlon_full = (nlon-woff) >= BDIM_X;
+
+    if (nchn_full && nlon_full) {
+        #pragma unroll
+        for(int j = 0; j < BDIM_X; j += BDIM_Y) {
+            sh[j+tidy][tidx] = src[batch][coff + j+tidy][h][woff+tidx];
+        }
+        __syncthreads();
+
+        #pragma unroll
+        for(int j = 0; j < BDIM_X; j += BDIM_Y) {
+            dst[batch][h][woff + j+tidy][coff+tidx] = sh[tidx][j+tidy];
+        }
+    } else {
+        if (woff+tidx < nlon) {
+            #pragma unroll
+            for(int j = 0; j < BDIM_X; j += BDIM_Y) {
+                sh[j+tidy][tidx] = (coff + j+tidy < nchn) ? src[batch][coff + j+tidy][h][woff+tidx] : VAL_T(0);
+            }
+        }
+        __syncthreads();
+
+        if (coff+tidx < nchn) {
+            #pragma unroll
+            for(int j = 0; j < BDIM_X; j += BDIM_Y) {
+                if (woff + j+tidy < nlon) {
+                    dst[batch][h][woff + j+tidy][coff+tidx] = sh[tidx][j+tidy];
+                }
+            }
+        }
+    }
+    return;
+}
+
+template<int WARPS_X_TILE, typename VAL_T>
+void launch_permute_to0231(at::Tensor src, at::Tensor dst){
+    dim3 block;
+    dim3 grid;
+
+    block.x = WARP_SIZE;
+    block.y = WARPS_X_TILE;
+    grid.x = DIV_UP(src.size(1), block.x);
+    grid.y = DIV_UP(src.size(3), block.x);
+    grid.z = src.size(2)*src.size(0);
+
+    assert(grid.y < 65536);
+    assert(grid.z < 65536);
+
+    // get stream
+    auto stream = at::cuda::getCurrentCUDAStream().stream();
+
+    permute_to0231_k<WARP_SIZE, WARPS_X_TILE>
+                        <<<grid, block, 0, stream>>>(src.size(1),
+                                                     src.size(2),
+                                                     src.size(3),
+                                                     src.packed_accessor32<VAL_T, 4, at::RestrictPtrTraits>(),
+                                                     dst.packed_accessor32<VAL_T, 4, at::RestrictPtrTraits>());
+}
+
+template<int BDIM_X,
+         int BDIM_Y,
+         typename VAL_T>
+__global__
+__launch_bounds__(BDIM_X*BDIM_Y)
+void  permute_to0312_k(const int nchn,
+                       const int nlat,
+                       const int nlon,
+                       const at::PackedTensorAccessor32<VAL_T, 4, at::RestrictPtrTraits> src,
+                             at::PackedTensorAccessor32<VAL_T, 4, at::RestrictPtrTraits> dst) {
+
+    static_assert(!(BDIM_X & (BDIM_X-1)));
+    static_assert(!(BDIM_Y & (BDIM_Y-1)));
+    static_assert(BDIM_X >= BDIM_Y);
+
+    __shared__ VAL_T sh[BDIM_X][BDIM_X+1];
+
+    const int tidx = threadIdx.x;
+    const int tidy = threadIdx.y;
+
+    const int woff  = blockIdx.x*BDIM_X;           // width offset
+    const int coff  = blockIdx.y*BDIM_X;           // channel offset
+    const int batch = blockIdx.z / nlat;           // batch (same for all block)
+    const int h     = blockIdx.z - (batch * nlat); // height (same for all block)
+
+    const int nchn_full = (nchn-coff) >= BDIM_X;
+    const int nlon_full = (nlon-woff) >= BDIM_X;
+
+    if (nchn_full && nlon_full) {
+        #pragma unroll
+        for(int j = 0; j < BDIM_X; j += BDIM_Y) {
+            sh[j+tidy][tidx] = src[batch][h][woff + j+tidy][coff+tidx];
+        }
+        __syncthreads();
+
+        #pragma unroll
+        for(int j = 0; j < BDIM_X; j += BDIM_Y) {
+            dst[batch][coff + j+tidy][h][woff+tidx] = sh[tidx][j+tidy];
+        }
+    } else {
+        if (coff+tidx < nchn) {
+            #pragma unroll
+            for(int j = 0; j < BDIM_X; j += BDIM_Y) {
+                sh[j+tidy][tidx] = (woff + j+tidy < nlon) ? src[batch][h][woff + j+tidy][coff+tidx] : VAL_T(0);
+            }
+        }
+        __syncthreads();
+
+        if (woff+tidx < nlon) {
+            #pragma unroll
+            for(int j = 0; j < BDIM_X; j += BDIM_Y) {
+                if (coff + j+tidy < nchn) {
+                    dst[batch][coff + j+tidy][h][woff+tidx] = sh[tidx][j+tidy];;
+                }
+            }
+        }
+    }
+    return;
+}
+
+template<int WARPS_X_TILE, typename VAL_T>
+void launch_permute_to0312(at::Tensor src, at::Tensor dst){
+    dim3 block;
+    dim3 grid;
+
+    block.x = WARP_SIZE;
+    block.y = WARPS_X_TILE;
+    grid.x = DIV_UP(src.size(2), block.x);
+    grid.y = DIV_UP(src.size(3), block.x);
+    grid.z = src.size(1)*src.size(0);
+
+    assert(grid.y < 65536);
+    assert(grid.z < 65536);
+
+    // get stream
+    auto stream = at::cuda::getCurrentCUDAStream().stream();
+
+    permute_to0312_k<WARP_SIZE, WARPS_X_TILE>
+                        <<<grid, block, 0, stream>>>(src.size(3),
+                                                     src.size(1),
+                                                     src.size(2),
+                                                     src.packed_accessor32<VAL_T, 4, at::RestrictPtrTraits>(),
+                                                     dst.packed_accessor32<VAL_T, 4, at::RestrictPtrTraits>());
+}

.deps/torch_harmonics/csrc/attention/cudamacro.h

@@ -0,0 +1,47 @@
+// coding=utf-8
+//
+// SPDX-FileCopyrightText: Copyright (c) 2025 The torch-harmonics Authors. All rights reserved.
+// SPDX-License-Identifier: BSD-3-Clause
+//
+// Redistribution and use in source and binary forms, with or without
+// modification, are permitted provided that the following conditions are met:
+//
+// 1. Redistributions of source code must retain the above copyright notice, this
+// list of conditions and the following disclaimer.
+//
+// 2. Redistributions in binary form must reproduce the above copyright notice,
+// this list of conditions and the following disclaimer in the documentation
+// and/or other materials provided with the distribution.
+//
+// 3. Neither the name of the copyright holder nor the names of its
+// contributors may be used to endorse or promote products derived from
+// this software without specific prior written permission.
+//
+// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+// FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+// OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+#pragma once
+
+#define CHECK_CUDA(call) {                                                   \
+    cudaError_t err = call;                                                  \
+    if( cudaSuccess != err) {                                                \
+        fprintf(stderr, "Cuda error in file '%s' in line %i : %s.\n",        \
+                __FILE__, __LINE__, cudaGetErrorString( err) );              \
+        exit(EXIT_FAILURE);                                                  \
+    }}
+
+#define CHECK_ERROR(errorMessage) {                                          \
+    cudaError_t err = cudaGetLastError();                                    \
+    if( cudaSuccess != err) {                                                \
+        fprintf(stderr, "Cuda error: %s in file '%s' in line %i : %s.\n",    \
+                errorMessage, __FILE__, __LINE__, cudaGetErrorString( err) );\
+        exit(EXIT_FAILURE);                                                  \
+    }}

.deps/torch_harmonics/csrc/disco/disco.h

@@ -0,0 +1,39 @@
+// coding=utf-8
+//
+// SPDX-FileCopyrightText: Copyright (c) 2024 The torch-harmonics Authors. All rights reserved.
+// SPDX-License-Identifier: BSD-3-Clause
+//
+// Redistribution and use in source and binary forms, with or without
+// modification, are permitted provided that the following conditions are met:
+//
+// 1. Redistributions of source code must retain the above copyright notice, this
+// list of conditions and the following disclaimer.
+//
+// 2. Redistributions in binary form must reproduce the above copyright notice,
+// this list of conditions and the following disclaimer in the documentation
+// and/or other materials provided with the distribution.
+//
+// 3. Neither the name of the copyright holder nor the names of its
+// contributors may be used to endorse or promote products derived from
+// this software without specific prior written permission.
+//
+// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+// FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+// OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+#pragma once
+
+#include <stdio.h>
+#include <stdlib.h>
+#include <torch/extension.h>
+#include <cassert>
+
+#define CHECK_CONTIGUOUS_TENSOR(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
+#define CHECK_INPUT_TENSOR(x) CHECK_CONTIGUOUS_TENSOR(x)

.deps/torch_harmonics/csrc/disco/disco_cuda.cuh

@@ -0,0 +1,54 @@
+// coding=utf-8
+//
+// SPDX-FileCopyrightText: Copyright (c) 2024 The torch-harmonics Authors. All rights reserved.
+// SPDX-License-Identifier: BSD-3-Clause
+//
+// Redistribution and use in source and binary forms, with or without
+// modification, are permitted provided that the following conditions are met:
+//
+// 1. Redistributions of source code must retain the above copyright notice, this
+// list of conditions and the following disclaimer.
+//
+// 2. Redistributions in binary form must reproduce the above copyright notice,
+// this list of conditions and the following disclaimer in the documentation
+// and/or other materials provided with the distribution.
+//
+// 3. Neither the name of the copyright holder nor the names of its
+// contributors may be used to endorse or promote products derived from
+// this software without specific prior written permission.
+//
+// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+// FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+// OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+#pragma once
+
+#include "disco.h"
+
+#include <cuda_runtime.h>
+#include <c10/cuda/CUDAStream.h>
+
+#define CHECK_CUDA_TENSOR(x) TORCH_CHECK(x.device().is_cuda(), #x " must be a CUDA tensor")
+#define CHECK_CUDA_INPUT_TENSOR(x)                                                                                     \
+    CHECK_CUDA_TENSOR(x);                                                                                              \
+    CHECK_CONTIGUOUS_TENSOR(x)
+
+#define DIV_UP(a, b) (((a) + ((b)-1)) / (b))
+
+#define MIN_THREADS (64)
+#define ELXTH_MAX (32)
+
+// forward kernel
+torch::Tensor disco_cuda_fwd(torch::Tensor inp, torch::Tensor roff_idx, torch::Tensor ker_idx, torch::Tensor row_idx,
+                             torch::Tensor col_idx, torch::Tensor val, int64_t K, int64_t Ho, int64_t Wo);
+
+// backward kernel
+torch::Tensor disco_cuda_bwd(torch::Tensor inp, torch::Tensor roff_idx, torch::Tensor ker_idx, torch::Tensor row_idx,
+                             torch::Tensor col_idx, torch::Tensor val, int64_t K, int64_t Ho, int64_t Wo);

.deps/torch_harmonics/csrc/disco/disco_cuda_bwd.cu

@@ -0,0 +1,281 @@
+// coding=utf-8
+//
+// SPDX-FileCopyrightText: Copyright (c) 2024 The torch-harmonics Authors. All rights reserved.
+// SPDX-License-Identifier: BSD-3-Clause
+//
+// Redistribution and use in source and binary forms, with or without
+// modification, are permitted provided that the following conditions are met:
+//
+// 1. Redistributions of source code must retain the above copyright notice, this
+// list of conditions and the following disclaimer.
+//
+// 2. Redistributions in binary form must reproduce the above copyright notice,
+// this list of conditions and the following disclaimer in the documentation
+// and/or other materials provided with the distribution.
+//
+// 3. Neither the name of the copyright holder nor the names of its
+// contributors may be used to endorse or promote products derived from
+// this software without specific prior written permission.
+//
+// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+// FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+// OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+#include "disco.h"
+#include "disco_cuda.cuh"
+
+template <int BDIM_X, int ELXTH, typename REAL_T>
+__device__ void disco_bwd_d(const int Hi, const int Wi, const int K, const int Ho, const int Wo, const int pscale,
+                            const int64_t *__restrict__ roff, const int64_t *__restrict__ kers,
+                            const int64_t *__restrict__ rows, const int64_t *__restrict__ cols,
+                            const REAL_T *__restrict__ vals, const REAL_T *__restrict__ inp, REAL_T *__restrict__ out)
+{
+
+    const int tid = threadIdx.x;
+
+    const int64_t bidx = blockIdx.x; // gloabl row
+    const int64_t bidy = blockIdx.y; // bc
+
+    int64_t soff = roff[bidx];
+    int64_t eoff = roff[bidx + 1];
+
+    const int64_t ker = kers[soff];
+    const int64_t row = rows[soff];
+
+    inp += bidy * K * Hi * Wi + ker * Hi * Wi + row * Wi;
+    out += bidy * Ho * Wo;
+
+    // align to larger supported fp type
+    extern __shared__ __align__(sizeof(double)) unsigned char __sh_ptr[]; // REAL_T __sh[2*(BDIM_X*ELXTH)*pscale]
+
+    REAL_T(*__sh)[BDIM_X * ELXTH * 2] = reinterpret_cast<REAL_T(*)[BDIM_X * ELXTH * 2]>(__sh_ptr);
+
+    // copy current inp row in regs
+    REAL_T __reg[ELXTH];
+
+#pragma unroll
+    for (int i = 0; i < ELXTH; i++) { __reg[i] = (i * BDIM_X + tid < Wi) ? inp[i * BDIM_X + tid] : REAL_T(0); }
+
+    // reset shared row up to Wo+2, remaining
+    // ppscale*(BDIM_X*ELXTH - Wo) locations
+    // will be written to but never copied to
+    // global mem
+    for (int i = 0; i < pscale; i++) {
+#pragma unroll
+        for (int j = 0; j < 2 * BDIM_X * ELXTH; j += BDIM_X) { __sh[i][j + tid] = 0; }
+    }
+    __syncthreads();
+
+    int col_prev = cols[soff];
+
+    int h_prev = col_prev / Wo;
+    int w_prev = col_prev % Wo;
+
+    // loops along the colums of CTA's row
+    for (int64_t nz = soff; nz < eoff; nz++) {
+
+        const int col = cols[nz];
+        const REAL_T val = vals[nz];
+
+        // if we are processing a nz with a col value
+        // leading to a new row of inp then copy it
+        // to shmem;
+        // we read a col that points to a new output
+        // row if (col / Wo) > (col_prev / Wo)
+        if (col >= col_prev - w_prev + Wo) {
+            __syncthreads();
+            for (int i = 0; i < pscale; i++) {
+                for (int j = tid; j < Wi; j += BDIM_X) {
+
+                    const REAL_T v = __sh[i][j] + __sh[i][Wi + j];
+
+                    atomicAdd(&out[h_prev * Wo + j * pscale + i], v);
+
+                    __sh[i][j] = 0;
+                    __sh[i][Wi + j] = 0;
+                }
+            }
+            __syncthreads();
+
+            col_prev = col;
+            h_prev = col / Wo;
+            w_prev = col % Wo;
+        }
+
+        const int w = w_prev + (col - col_prev);
+        const int w_mod_ps = w % pscale;
+        const int w_div_ps = w / pscale;
+
+#pragma unroll
+        for (int i = 0; i < ELXTH; i++) {
+
+            const int pp = i * BDIM_X + tid;
+            __sh[w_mod_ps][w_div_ps + pp] += val * __reg[i];
+        }
+
+        // to avoid race conditions on __sh[]
+        // among consecutive iterations along nz
+        __syncthreads();
+    }
+    __syncthreads();
+
+    // write last row
+    for (int i = 0; i < pscale; i++) {
+
+        for (int j = tid; j < Wi; j += BDIM_X) {
+
+            const REAL_T v = __sh[i][j] + __sh[i][Wi + j];
+            atomicAdd(&out[h_prev * Wo + j * pscale + i], v);
+        }
+    }
+    return;
+}
+
+template <int BDIM_X, int ELXTH, int PSCALE, typename REAL_T>
+__global__
+    __launch_bounds__(BDIM_X) void disco_bwd_blk_k(const int Hi, const int Wi, const int K, const int Ho, const int Wo,
+                                                   const int pscale, const int64_t *__restrict__ roff,
+                                                   const int64_t *__restrict__ kers, const int64_t *__restrict__ rows,
+                                                   const int64_t *__restrict__ cols, const REAL_T *__restrict__ vals,
+                                                   const REAL_T *__restrict__ inp, REAL_T *__restrict__ out)
+{
+
+    if constexpr (PSCALE != 0) {
+        disco_bwd_d<BDIM_X, ELXTH>(Hi, Wi, K, Ho, Wo, PSCALE, roff, kers, rows, cols, vals, inp, out);
+    } else {
+        disco_bwd_d<BDIM_X, ELXTH>(Hi, Wi, K, Ho, Wo, pscale, roff, kers, rows, cols, vals, inp, out);
+    }
+
+    return;
+}
+
+template <int NTH, int ELXTH, typename REAL_T>
+static void launch_kernel(int BC, int Hi, int Wi, int K, int Ho, int Wo, int64_t nrows, int64_t *roff_d, int64_t *ker_d,
+                          int64_t *row_d, int64_t *col_d, REAL_T *val_d, REAL_T *inp_d, REAL_T *out_d,
+                          cudaStream_t stream)
+{
+
+    static_assert(sizeof(REAL_T) == 2 || sizeof(REAL_T) == 4 || sizeof(REAL_T) == 8);
+
+    if constexpr (ELXTH <= ELXTH_MAX) {
+        if (NTH * ELXTH >= Wi) {
+            dim3 grid(nrows, BC);
+
+            const int pscale = Wo / Wi;
+            size_t shmem = sizeof(*out_d) * (2 * (NTH * ELXTH) * pscale);
+
+            switch (pscale) {
+            case 1:
+                disco_bwd_blk_k<NTH, ELXTH, 1><<<grid, NTH, shmem, stream>>>(Hi, Wi, K, Ho, Wo, pscale, roff_d, ker_d,
+                                                                             row_d, col_d, val_d, inp_d, out_d);
+                break;
+            case 2:
+                disco_bwd_blk_k<NTH, ELXTH, 2><<<grid, NTH, shmem, stream>>>(Hi, Wi, K, Ho, Wo, pscale, roff_d, ker_d,
+                                                                             row_d, col_d, val_d, inp_d, out_d);
+                break;
+            case 3:
+                disco_bwd_blk_k<NTH, ELXTH, 3><<<grid, NTH, shmem, stream>>>(Hi, Wi, K, Ho, Wo, pscale, roff_d, ker_d,
+                                                                             row_d, col_d, val_d, inp_d, out_d);
+                break;
+            default:
+                disco_bwd_blk_k<NTH, ELXTH, 0><<<grid, NTH, shmem, stream>>>(Hi, Wi, K, Ho, Wo, pscale, roff_d, ker_d,
+                                                                             row_d, col_d, val_d, inp_d, out_d);
+            }
+        } else {
+            launch_kernel<NTH, ELXTH + 1>(BC, Hi, Wi, K, Ho, Wo, nrows, roff_d, ker_d, row_d, col_d, val_d, inp_d,
+                                          out_d, stream);
+        }
+    }
+    return;
+}
+
+torch::Tensor disco_cuda_bwd(torch::Tensor inp, torch::Tensor roff_idx, torch::Tensor ker_idx, torch::Tensor row_idx,
+                             torch::Tensor col_idx, torch::Tensor val, int64_t K, int64_t Ho, int64_t Wo)
+{
+
+    // some sanity checks
+    CHECK_CUDA_INPUT_TENSOR(inp);
+    CHECK_CUDA_INPUT_TENSOR(roff_idx);
+    CHECK_CUDA_INPUT_TENSOR(ker_idx);
+    CHECK_CUDA_INPUT_TENSOR(row_idx);
+    CHECK_CUDA_INPUT_TENSOR(col_idx);
+    CHECK_CUDA_INPUT_TENSOR(val);
+
+    // extract some shapes
+    int64_t B = inp.size(0);
+    int64_t C = inp.size(1);
+    int64_t BC = B * C;
+    int64_t Hi = inp.size(3);
+    int64_t Wi = inp.size(4);
+    int64_t nrows = roff_idx.size(0) - 1;
+
+    // allocate output
+    int64_t out_dims[] = {B, C, Ho, Wo};
+    auto options = torch::TensorOptions().device(inp.device()).dtype(inp.dtype());
+    torch::Tensor out = torch::zeros(out_dims, options);
+
+    // get stream
+    auto stream = at::cuda::getCurrentCUDAStream().stream();
+
+    // assert
+    static_assert(0 == (ELXTH_MAX % 2));
+
+    if (Wo <= 64 * ELXTH_MAX) {
+        AT_DISPATCH_FLOATING_TYPES(inp.scalar_type(), "disco_backward_cuda", ([&] {
+                                       launch_kernel<64, 1, scalar_t>(
+                                           BC, Hi, Wi, K, Ho, Wo, nrows, roff_idx.data_ptr<int64_t>(),
+                                           ker_idx.data_ptr<int64_t>(), row_idx.data_ptr<int64_t>(),
+                                           col_idx.data_ptr<int64_t>(), val.data_ptr<scalar_t>(),
+                                           inp.data_ptr<scalar_t>(), out.data_ptr<scalar_t>(), stream);
+                                   }));
+    } else if (Wo <= 128 * ELXTH_MAX) {
+        AT_DISPATCH_FLOATING_TYPES(inp.scalar_type(), "disco_backward_cuda", ([&] {
+                                       launch_kernel<128, (ELXTH_MAX / 2) + 1, scalar_t>(
+                                           BC, Hi, Wi, K, Ho, Wo, nrows, roff_idx.data_ptr<int64_t>(),
+                                           ker_idx.data_ptr<int64_t>(), row_idx.data_ptr<int64_t>(),
+                                           col_idx.data_ptr<int64_t>(), val.data_ptr<scalar_t>(),
+                                           inp.data_ptr<scalar_t>(), out.data_ptr<scalar_t>(), stream);
+                                   }));
+    } else if (Wo <= 256 * ELXTH_MAX) {
+        AT_DISPATCH_FLOATING_TYPES(inp.scalar_type(), "disco_backward_cuda", ([&] {
+                                       launch_kernel<256, (ELXTH_MAX / 2) + 1, scalar_t>(
+                                           BC, Hi, Wi, K, Ho, Wo, nrows, roff_idx.data_ptr<int64_t>(),
+                                           ker_idx.data_ptr<int64_t>(), row_idx.data_ptr<int64_t>(),
+                                           col_idx.data_ptr<int64_t>(), val.data_ptr<scalar_t>(),
+                                           inp.data_ptr<scalar_t>(), out.data_ptr<scalar_t>(), stream);
+                                   }));
+    } else if (Wo <= 512 * ELXTH_MAX) {
+        AT_DISPATCH_FLOATING_TYPES(inp.scalar_type(), "disco_backward_cuda", ([&] {
+                                       launch_kernel<512, (ELXTH_MAX / 2) + 1, scalar_t>(
+                                           BC, Hi, Wi, K, Ho, Wo, nrows, roff_idx.data_ptr<int64_t>(),
+                                           ker_idx.data_ptr<int64_t>(), row_idx.data_ptr<int64_t>(),
+                                           col_idx.data_ptr<int64_t>(), val.data_ptr<scalar_t>(),
+                                           inp.data_ptr<scalar_t>(), out.data_ptr<scalar_t>(), stream);
+                                   }));
+    } else if (Wo <= 1024 * ELXTH_MAX) {
+        AT_DISPATCH_FLOATING_TYPES(inp.scalar_type(), "disco_backward_cuda", ([&] {
+                                       launch_kernel<1024, (ELXTH_MAX / 2) + 1, scalar_t>(
+                                           BC, Hi, Wi, K, Ho, Wo, nrows, roff_idx.data_ptr<int64_t>(),
+                                           ker_idx.data_ptr<int64_t>(), row_idx.data_ptr<int64_t>(),
+                                           col_idx.data_ptr<int64_t>(), val.data_ptr<scalar_t>(),
+                                           inp.data_ptr<scalar_t>(), out.data_ptr<scalar_t>(), stream);
+                                   }));
+    } else {
+        fprintf(stderr, "%s:%d: error, unsupported Wo value (%ld), max supported is %d\n", __FILE__, __LINE__, Wo,
+                1024 * ELXTH_MAX);
+        exit(EXIT_FAILURE);
+    }
+
+    return out;
+}
+
+// PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+//  m.def("backward", &disco_cuda_bwd, "DISCO backward (CUDA)");
+//}

.deps/torch_harmonics/csrc/disco/disco_cuda_fwd.cu

@@ -0,0 +1,265 @@
+// coding=utf-8
+//
+// SPDX-FileCopyrightText: Copyright (c) 2024 The torch-harmonics Authors. All rights reserved.
+// SPDX-License-Identifier: BSD-3-Clause
+//
+// Redistribution and use in source and binary forms, with or without
+// modification, are permitted provided that the following conditions are met:
+//
+// 1. Redistributions of source code must retain the above copyright notice, this
+// list of conditions and the following disclaimer.
+//
+// 2. Redistributions in binary form must reproduce the above copyright notice,
+// this list of conditions and the following disclaimer in the documentation
+// and/or other materials provided with the distribution.
+//
+// 3. Neither the name of the copyright holder nor the names of its
+// contributors may be used to endorse or promote products derived from
+// this software without specific prior written permission.
+//
+// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+// FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+// OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+#include "disco.h"
+#include "disco_cuda.cuh"
+
+template <int BDIM_X, int ELXTH, typename REAL_T>
+__device__ void disco_fwd_d(const int Hi, const int Wi, const int K, const int Ho, const int Wo, const int pscale,
+                            const int64_t *__restrict__ roff, const int64_t *__restrict__ kers,
+                            const int64_t *__restrict__ rows, const int64_t *__restrict__ cols,
+                            const REAL_T *__restrict__ vals, const REAL_T *__restrict__ inp, REAL_T *__restrict__ out)
+{
+
+    const int tid = threadIdx.x;
+
+    const int64_t bidx = blockIdx.x; // gloabl row
+    const int64_t bidy = blockIdx.y; // bc
+
+    int64_t soff = roff[bidx];
+    int64_t eoff = roff[bidx + 1];
+
+    const int64_t ker = kers[soff];
+    const int64_t row = rows[soff];
+
+    inp += bidy * Hi * Wi;
+    out += bidy * K * Ho * Wo + ker * Ho * Wo + row * Wo;
+
+    REAL_T __reg[ELXTH] = {0};
+
+    // align to larger supported fp type
+    extern __shared__ __align__(
+        sizeof(double)) unsigned char __sh_ptr[]; // REAL_T __sh[2*Wi + ppscale*(BDIM_X*ELXTH - Wo)]
+    REAL_T *__sh = reinterpret_cast<REAL_T *>(__sh_ptr);
+
+    int col_prev = cols[soff];
+
+    int h_prev = col_prev / Wi;
+    int w_prev = col_prev % Wi;
+
+    // copy current inp row in shmem
+    for (int i = tid; i < Wi; i += BDIM_X) {
+        const REAL_T v = inp[h_prev * Wi + i];
+        __sh[i] = v;
+        __sh[Wi + i] = v;
+    }
+    // locations __sh[2*Wi : ppscale*(BDIM_X*ELXTH-Wo)] are not used
+    __syncthreads();
+
+    // loops along the colums of CTA's row
+    for (int64_t nz = soff; nz < eoff; nz++) {
+
+        const int col = cols[nz];
+        const REAL_T val = vals[nz];
+
+        // if we are processing a nz with a col value
+        // leading to a new row of inp then copy it
+        // to shmem;
+        // checks whether (h_prev < h) with:
+        //  (col >= col_prev - (col_prev % Wi) + Wi)
+        if (col >= col_prev - w_prev + Wi) {
+
+            col_prev = col;
+            h_prev = col / Wi;
+            w_prev = col % Wi;
+
+            __syncthreads();
+            for (int i = tid; i < Wi; i += BDIM_X) {
+                const REAL_T v = inp[h_prev * Wi + i];
+                __sh[i] = v;
+                __sh[Wi + i] = v;
+            }
+            __syncthreads();
+        }
+
+        const int w = w_prev + (col - col_prev);
+
+#pragma unroll
+        for (int i = 0; i < ELXTH; i++) {
+
+            const int pp = i * BDIM_X + tid;
+
+            // original lines:
+            //
+            //   if (pp >= Wo) break;
+            //   const int wpp = (w + pscale*pp) % Wi;
+            //
+            // value of (w + pscale*pp) < (Wi + (Wi/Wo)*Wo) = 2*Wi
+            // so we can allocate twice the amount of shmem,
+            // replicate the current inp row and avoid the costly mod
+            //
+            // also, to avoid the conditional, sh can be extended to
+            // cover the maximum location accessed during this loop
+            //
+            // REAL_T __sh[2*Wi + ppscale*NUM_REM]
+            //
+            //   Wi + (Wi/Wo)*BDIM_X*ELXTH = (since BDIM_X*ELXTH >= Wo) =
+            // = Wi + (Wi/Wo)*(Wo + (BDIM_X*ELXTH - Wo)) =
+            // = 2*Wi + ppscale*NUM_REM
+            //
+            // with NUM_REM = BDIM_X*ELXTH - Wo
+
+            const int wpp = w + pscale * pp;
+
+            __reg[i] += val * __sh[wpp];
+        }
+    }
+
+#pragma unroll
+    for (int i = 0; i < ELXTH; i++) {
+
+        const int pp = i * BDIM_X + tid;
+        if (pp >= Wo) break;
+
+        out[pp] = __reg[i];
+    }
+
+    return;
+}
+
+template <int BDIM_X, int ELXTH, typename REAL_T>
+__global__
+    __launch_bounds__(BDIM_X) void disco_fwd_blk_k(const int Hi, const int Wi, const int K, const int Ho, const int Wo,
+                                                   const int pscale, const int64_t *__restrict__ roff,
+                                                   const int64_t *__restrict__ kers, const int64_t *__restrict__ rows,
+                                                   const int64_t *__restrict__ cols, const REAL_T *__restrict__ vals,
+                                                   const REAL_T *__restrict__ inp, REAL_T *__restrict__ out)
+{
+
+    disco_fwd_d<BDIM_X, ELXTH>(Hi, Wi, K, Ho, Wo, pscale, roff, kers, rows, cols, vals, inp, out);
+
+    return;
+}
+
+template <int NTH, int ELXTH, typename REAL_T>
+static void launch_kernel(int BC, int Hi, int Wi, int K, int Ho, int Wo, int64_t nrows, int64_t *roff_d, int64_t *ker_d,
+                          int64_t *row_d, int64_t *col_d, REAL_T *val_d, REAL_T *inp_d, REAL_T *out_d,
+                          cudaStream_t stream)
+{
+
+    static_assert(sizeof(REAL_T) == 2 || sizeof(REAL_T) == 4 || sizeof(REAL_T) == 8);
+
+    if constexpr (ELXTH <= ELXTH_MAX) {
+        if (NTH * ELXTH >= Wo) {
+            dim3 grid(nrows, BC);
+
+            const int pscale = Wi / Wo;
+            size_t shmem = sizeof(*out_d) * (Wi * 2 + pscale * (NTH * ELXTH - Wo));
+
+            disco_fwd_blk_k<NTH, ELXTH><<<grid, NTH, shmem, stream>>>(Hi, Wi, K, Ho, Wo, pscale, roff_d, ker_d, row_d,
+                                                                      col_d, val_d, inp_d, out_d);
+        } else {
+            launch_kernel<NTH, ELXTH + 1>(BC, Hi, Wi, K, Ho, Wo, nrows, roff_d, ker_d, row_d, col_d, val_d, inp_d,
+                                          out_d, stream);
+        }
+    }
+    return;
+}
+
+torch::Tensor disco_cuda_fwd(torch::Tensor inp, torch::Tensor roff_idx, torch::Tensor ker_idx, torch::Tensor row_idx,
+                             torch::Tensor col_idx, torch::Tensor val, int64_t K, int64_t Ho, int64_t Wo)
+{
+
+    // some sanity checks
+    CHECK_CUDA_INPUT_TENSOR(inp);
+    CHECK_CUDA_INPUT_TENSOR(roff_idx);
+    CHECK_CUDA_INPUT_TENSOR(ker_idx);
+    CHECK_CUDA_INPUT_TENSOR(row_idx);
+    CHECK_CUDA_INPUT_TENSOR(col_idx);
+    CHECK_CUDA_INPUT_TENSOR(val);
+
+    // extract some shapes
+    int64_t B = inp.size(0);
+    int64_t C = inp.size(1);
+    int64_t BC = B * C;
+    int64_t Hi = inp.size(2);
+    int64_t Wi = inp.size(3);
+    int64_t nrows = roff_idx.size(0) - 1;
+
+    // allocate output
+    int64_t out_dims[] = {B, C, K, Ho, Wo};
+    auto options = torch::TensorOptions().device(inp.device()).dtype(inp.dtype());
+    torch::Tensor out = torch::zeros(out_dims, options);
+
+    // get stream
+    auto stream = at::cuda::getCurrentCUDAStream().stream();
+
+    // assert
+    static_assert(0 == (ELXTH_MAX % 2));
+
+    // pick the correct launch config
+    if (Wo <= 64 * ELXTH_MAX) {
+        AT_DISPATCH_FLOATING_TYPES(inp.scalar_type(), "disco_forward_cuda", ([&] {
+                                       launch_kernel<64, 1, scalar_t>(
+                                           BC, Hi, Wi, K, Ho, Wo, nrows, roff_idx.data_ptr<int64_t>(),
+                                           ker_idx.data_ptr<int64_t>(), row_idx.data_ptr<int64_t>(),
+                                           col_idx.data_ptr<int64_t>(), val.data_ptr<scalar_t>(),
+                                           inp.data_ptr<scalar_t>(), out.data_ptr<scalar_t>(), stream);
+                                   }));
+    } else if (Wo <= 128 * ELXTH_MAX) {
+        AT_DISPATCH_FLOATING_TYPES(inp.scalar_type(), "disco_forward_cuda", ([&] {
+                                       launch_kernel<128, (ELXTH_MAX / 2) + 1, scalar_t>(
+                                           BC, Hi, Wi, K, Ho, Wo, nrows, roff_idx.data_ptr<int64_t>(),
+                                           ker_idx.data_ptr<int64_t>(), row_idx.data_ptr<int64_t>(),
+                                           col_idx.data_ptr<int64_t>(), val.data_ptr<scalar_t>(),
+                                           inp.data_ptr<scalar_t>(), out.data_ptr<scalar_t>(), stream);
+                                   }));
+    } else if (Wo <= 256 * ELXTH_MAX) {
+        AT_DISPATCH_FLOATING_TYPES(inp.scalar_type(), "disco_forward_cuda", ([&] {
+                                       launch_kernel<256, (ELXTH_MAX / 2) + 1, scalar_t>(
+                                           BC, Hi, Wi, K, Ho, Wo, nrows, roff_idx.data_ptr<int64_t>(),
+                                           ker_idx.data_ptr<int64_t>(), row_idx.data_ptr<int64_t>(),
+                                           col_idx.data_ptr<int64_t>(), val.data_ptr<scalar_t>(),
+                                           inp.data_ptr<scalar_t>(), out.data_ptr<scalar_t>(), stream);
+                                   }));
+    } else if (Wo <= 512 * ELXTH_MAX) {
+        AT_DISPATCH_FLOATING_TYPES(inp.scalar_type(), "disco_forward_cuda", ([&] {
+                                       launch_kernel<512, (ELXTH_MAX / 2) + 1, scalar_t>(
+                                           BC, Hi, Wi, K, Ho, Wo, nrows, roff_idx.data_ptr<int64_t>(),
+                                           ker_idx.data_ptr<int64_t>(), row_idx.data_ptr<int64_t>(),
+                                           col_idx.data_ptr<int64_t>(), val.data_ptr<scalar_t>(),
+                                           inp.data_ptr<scalar_t>(), out.data_ptr<scalar_t>(), stream);
+                                   }));
+    } else if (Wo <= 1024 * ELXTH_MAX) {
+        AT_DISPATCH_FLOATING_TYPES(inp.scalar_type(), "disco_forward_cuda", ([&] {
+                                       launch_kernel<1024, (ELXTH_MAX / 2) + 1, scalar_t>(
+                                           BC, Hi, Wi, K, Ho, Wo, nrows, roff_idx.data_ptr<int64_t>(),
+                                           ker_idx.data_ptr<int64_t>(), row_idx.data_ptr<int64_t>(),
+                                           col_idx.data_ptr<int64_t>(), val.data_ptr<scalar_t>(),
+                                           inp.data_ptr<scalar_t>(), out.data_ptr<scalar_t>(), stream);
+                                   }));
+    } else {
+        fprintf(stderr, "%s:%d: error, unsupported Wo value (%ld), max supported is %d\n", __FILE__, __LINE__, Wo,
+                1024 * ELXTH_MAX);
+        exit(EXIT_FAILURE);
+    }
+
+    return out;
+}

.deps/torch_harmonics/csrc/disco/disco_helpers.cpp

@@ -0,0 +1,149 @@
+// coding=utf-8
+//
+// SPDX-FileCopyrightText: Copyright (c) 2024 The torch-harmonics Authors. All rights reserved.
+// SPDX-License-Identifier: BSD-3-Clause
+//
+// Redistribution and use in source and binary forms, with or without
+// modification, are permitted provided that the following conditions are met:
+//
+// 1. Redistributions of source code must retain the above copyright notice, this
+// list of conditions and the following disclaimer.
+//
+// 2. Redistributions in binary form must reproduce the above copyright notice,
+// this list of conditions and the following disclaimer in the documentation
+// and/or other materials provided with the distribution.
+//
+// 3. Neither the name of the copyright holder nor the names of its
+// contributors may be used to endorse or promote products derived from
+// this software without specific prior written permission.
+//
+// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+// FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+// OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+#include "disco.h"
+
+template <typename REAL_T>
+void preprocess_psi_kernel(int64_t nnz, int64_t K, int64_t Ho, int64_t *ker_h, int64_t *row_h, int64_t *col_h,
+                           int64_t *roff_h, REAL_T *val_h, int64_t &nrows)
+{
+
+    int64_t *Koff = new int64_t[K];
+    for (int i = 0; i < K; i++) { Koff[i] = 0; }
+
+    for (int64_t i = 0; i < nnz; i++) { Koff[ker_h[i]]++; }
+
+    int64_t prev = Koff[0];
+    Koff[0] = 0;
+    for (int i = 1; i < K; i++) {
+        int64_t save = Koff[i];
+        Koff[i] = prev + Koff[i - 1];
+        prev = save;
+    }
+
+    int64_t *ker_sort = new int64_t[nnz];
+    int64_t *row_sort = new int64_t[nnz];
+    int64_t *col_sort = new int64_t[nnz];
+    float *val_sort = new float[nnz];
+
+    for (int64_t i = 0; i < nnz; i++) {
+
+        const int64_t ker = ker_h[i];
+        const int64_t off = Koff[ker]++;
+
+        ker_sort[off] = ker;
+        row_sort[off] = row_h[i];
+        col_sort[off] = col_h[i];
+        val_sort[off] = val_h[i];
+    }
+    for (int64_t i = 0; i < nnz; i++) {
+        ker_h[i] = ker_sort[i];
+        row_h[i] = row_sort[i];
+        col_h[i] = col_sort[i];
+        val_h[i] = val_sort[i];
+    }
+
+    delete[] Koff;
+    delete[] ker_sort;
+    delete[] row_sort;
+    delete[] col_sort;
+    delete[] val_sort;
+
+    // compute rows offsets
+    nrows = 1;
+    roff_h[0] = 0;
+    for (int64_t i = 1; i < nnz; i++) {
+
+        if (row_h[i - 1] == row_h[i]) continue;
+        roff_h[nrows++] = i;
+
+        if (nrows > Ho * K) {
+            fprintf(stderr, "%s:%d: error, found more rows in the K COOs than Ho*K (%ld)\n", __FILE__, __LINE__,
+                    int64_t(Ho) * K);
+            exit(EXIT_FAILURE);
+        }
+    }
+    roff_h[nrows] = nnz;
+
+    return;
+}
+
+torch::Tensor preprocess_psi(const int64_t K, const int64_t Ho, torch::Tensor ker_idx, torch::Tensor row_idx,
+                             torch::Tensor col_idx, torch::Tensor val)
+{
+
+    CHECK_INPUT_TENSOR(ker_idx);
+    CHECK_INPUT_TENSOR(row_idx);
+    CHECK_INPUT_TENSOR(col_idx);
+    CHECK_INPUT_TENSOR(val);
+
+    // get the input device and make sure all tensors are on the same device
+    auto device = ker_idx.device();
+    TORCH_INTERNAL_ASSERT(device.type() == row_idx.device().type() && (device.type() == col_idx.device().type()) && (device.type() == val.device().type()));
+
+    // move to cpu
+    ker_idx = ker_idx.to(torch::kCPU);
+    row_idx = row_idx.to(torch::kCPU);
+    col_idx = col_idx.to(torch::kCPU);
+    val = val.to(torch::kCPU);
+
+    int64_t nnz = val.size(0);
+    int64_t *ker_h = ker_idx.data_ptr<int64_t>();
+    int64_t *row_h = row_idx.data_ptr<int64_t>();
+    int64_t *col_h = col_idx.data_ptr<int64_t>();
+    int64_t *roff_h = new int64_t[Ho * K + 1];
+    int64_t nrows;
+
+    AT_DISPATCH_FLOATING_TYPES(val.scalar_type(), "preprocess_psi", ([&] {
+                                   preprocess_psi_kernel<scalar_t>(nnz, K, Ho, ker_h, row_h, col_h, roff_h,
+                                                                   val.data_ptr<scalar_t>(), nrows);
+                               }));
+
+    // create output tensor
+    auto roff_idx = torch::empty({nrows + 1}, row_idx.options());
+    int64_t *roff_out_h = roff_idx.data_ptr<int64_t>();
+
+    for (int64_t i = 0; i < (nrows + 1); i++) { roff_out_h[i] = roff_h[i]; }
+    delete[] roff_h;
+
+    // move to original device
+    ker_idx = ker_idx.to(device);
+    row_idx = row_idx.to(device);
+    col_idx = col_idx.to(device);
+    val = val.to(device);
+    roff_idx = roff_idx.to(device);
+
+    return roff_idx;
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
+{
+    m.def("preprocess_psi", &preprocess_psi, "Sort psi matrix, required for using disco_cuda.");
+}

.deps/torch_harmonics/csrc/disco/disco_interface.cu

@@ -0,0 +1,38 @@
+// coding=utf-8
+//
+// SPDX-FileCopyrightText: Copyright (c) 2024 The torch-harmonics Authors. All rights reserved.
+// SPDX-License-Identifier: BSD-3-Clause
+//
+// Redistribution and use in source and binary forms, with or without
+// modification, are permitted provided that the following conditions are met:
+//
+// 1. Redistributions of source code must retain the above copyright notice, this
+// list of conditions and the following disclaimer.
+//
+// 2. Redistributions in binary form must reproduce the above copyright notice,
+// this list of conditions and the following disclaimer in the documentation
+// and/or other materials provided with the distribution.
+//
+// 3. Neither the name of the copyright holder nor the names of its
+// contributors may be used to endorse or promote products derived from
+// this software without specific prior written permission.
+//
+// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+// FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+// OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+#include "disco.h"
+#include "disco_cuda.cuh"
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
+{
+    m.def("forward", &disco_cuda_fwd, "DISCO forward (CUDA)");
+    m.def("backward", &disco_cuda_bwd, "DISCO backward (CUDA)");
+}

.deps/torch_harmonics/distributed/__init__.py

@@ -0,0 +1,58 @@
+# coding=utf-8
+
+# SPDX-FileCopyrightText: Copyright (c) 2022 The torch-harmonics Authors. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+# 
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+#
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+#
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+#
+
+# we need this in order to enable distributed
+from .utils import init, finalize, is_initialized, polar_group, azimuth_group
+from .utils import polar_group_size, azimuth_group_size, polar_group_rank, azimuth_group_rank
+from .primitives import compute_split_shapes, split_tensor_along_dim
+from .primitives import (
+    distributed_transpose_azimuth,
+    distributed_transpose_polar,
+    reduce_from_polar_region,
+    reduce_from_azimuth_region,
+    scatter_to_polar_region,
+    gather_from_polar_region,
+    copy_to_polar_region,
+    copy_to_azimuth_region,
+    reduce_from_scatter_to_polar_region,
+    gather_from_copy_to_polar_region
+)
+
+# import the sht
+from .distributed_sht import DistributedRealSHT, DistributedInverseRealSHT
+from .distributed_sht import DistributedRealVectorSHT, DistributedInverseRealVectorSHT
+
+# import DISCO
+from .distributed_convolution import DistributedDiscreteContinuousConvS2
+from .distributed_convolution import DistributedDiscreteContinuousConvTransposeS2
+
+# import resampling
+from .distributed_resample import DistributedResampleS2

.deps/torch_harmonics/distributed/distributed_convolution.py

@@ -0,0 +1,474 @@
+# coding=utf-8
+
+# SPDX-FileCopyrightText: Copyright (c) 2022 The torch-harmonics Authors. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+#
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+#
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+#
+
+import abc
+from typing import List, Tuple, Union, Optional
+from itertools import accumulate
+from warnings import warn
+
+import math
+
+import torch
+import torch.nn as nn
+
+from functools import partial
+
+from torch_harmonics.quadrature import _precompute_grid, _precompute_latitudes, _precompute_longitudes
+from torch_harmonics._disco_convolution import _get_psi, _disco_s2_contraction_torch, _disco_s2_transpose_contraction_torch
+from torch_harmonics._disco_convolution import _disco_s2_contraction_cuda, _disco_s2_transpose_contraction_cuda
+from torch_harmonics.filter_basis import get_filter_basis
+from torch_harmonics.convolution import (
+    _precompute_convolution_tensor_s2,
+    DiscreteContinuousConv,
+)
+
+
+from torch_harmonics.distributed import polar_group_size, azimuth_group_size
+from torch_harmonics.distributed import distributed_transpose_azimuth, distributed_transpose_polar
+from torch_harmonics.distributed import reduce_from_polar_region, scatter_to_polar_region, gather_from_polar_region, copy_to_polar_region
+from torch_harmonics.distributed import polar_group_rank, azimuth_group_rank
+from torch_harmonics.distributed import compute_split_shapes, split_tensor_along_dim
+
+# import custom C++/CUDA extensions if available
+try:
+    from disco_helpers import preprocess_psi
+    import disco_cuda_extension
+
+    _cuda_extension_available = True
+except ImportError as err:
+    disco_cuda_extension = None
+    _cuda_extension_available = False
+
+
+def _split_distributed_convolution_tensor_s2(
+    idx: torch.Tensor,
+    vals: torch.Tensor,
+    in_shape: Tuple[int],
+    out_shape: Tuple[int],
+):
+    """
+    Splits a pre-computed convolution tensor along the latitude dimension for distributed processing.
+
+    This function takes a convolution tensor that was generated by the serial routine and filters
+    it to only include entries corresponding to the local latitude slice assigned to this process.
+    The filtering is done based on the polar group rank and the computed split shapes.
+
+    Parameters
+    ----------
+    idx: torch.Tensor
+        Indices of the pre-computed convolution tensor
+    vals: torch.Tensor
+        Values of the pre-computed convolution tensor
+    in_shape: Tuple[int]
+        Shape of the input tensor (nlat_in, nlon_in)
+    out_shape: Tuple[int]
+        Shape of the output tensor (nlat_out, nlon_out)
+
+    Returns
+    -------
+    idx: torch.Tensor
+        Filtered indices corresponding to the local latitude slice
+    vals: torch.Tensor
+        Filtered values corresponding to the local latitude slice
+    """
+
+    nlat_in, nlon_in = in_shape
+    nlat_out, nlon_out = out_shape
+
+    comm_size_polar = polar_group_size()
+    comm_rank_polar = polar_group_rank()
+    split_shapes = compute_split_shapes(nlat_in, num_chunks=comm_size_polar)
+    offsets = [0] + list(accumulate(split_shapes))
+    start_idx = offsets[comm_rank_polar]
+    end_idx = offsets[comm_rank_polar + 1]
+
+    # once normalization is done we can throw away the entries which correspond to input latitudes we do not care about
+    lats = idx[2] // nlon_in
+    lons = idx[2] % nlon_in
+    ilats = torch.argwhere((lats < end_idx) & (lats >= start_idx)).squeeze()
+    vals = vals[ilats]
+    # for the indices we need to recompute them to refer to local indices of the input tenor
+    idx = torch.stack([idx[0, ilats], idx[1, ilats], (lats[ilats] - start_idx) * nlon_in + lons[ilats]], dim=0)
+
+    # make results contiguous
+    idx = idx.contiguous()
+    vals = vals.to(dtype=torch.float32).contiguous()
+
+    return idx, vals
+
+
+class DistributedDiscreteContinuousConvS2(DiscreteContinuousConv):
+    """
+    Distributed version of Discrete-continuous convolutions (DISCO) on the 2-Sphere as described in [1].
+    We assume the data can be splitted in polar and azimuthal directions.
+
+    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
+
+    Returns
+    -------
+    out: torch.Tensor
+        Output tensor
+
+    References
+    ----------
+    [1] Ocampo, Price, McEwen, Scalable and equivariant spherical CNNs by discrete-continuous (DISCO) convolutions, ICLR (2023), arXiv:2209.13603
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        in_shape: Tuple[int],
+        out_shape: Tuple[int],
+        kernel_shape: Union[int, Tuple[int], Tuple[int, int]],
+        basis_type: Optional[str] = "piecewise linear",
+        basis_norm_mode: Optional[str] = "mean",
+        groups: Optional[int] = 1,
+        grid_in: Optional[str] = "equiangular",
+        grid_out: Optional[str] = "equiangular",
+        bias: Optional[bool] = True,
+        theta_cutoff: Optional[float] = None,
+    ):
+        super().__init__(in_channels, out_channels, kernel_shape, basis_type, groups, bias)
+
+        self.nlat_in, self.nlon_in = in_shape
+        self.nlat_out, self.nlon_out = out_shape
+
+        # get the comms grid:
+        self.comm_size_polar = polar_group_size()
+        self.comm_rank_polar = polar_group_rank()
+        self.comm_size_azimuth = azimuth_group_size()
+        self.comm_rank_azimuth = azimuth_group_rank()
+
+        # we need those shapes:
+        self.lat_in_shapes = compute_split_shapes(self.nlat_in, self.comm_size_polar)
+        self.lon_in_shapes = compute_split_shapes(self.nlon_in, self.comm_size_azimuth)
+        self.lat_out_shapes = compute_split_shapes(self.nlat_out, self.comm_size_polar)
+        self.lon_out_shapes = compute_split_shapes(self.nlon_out, self.comm_size_azimuth)
+
+        # compute theta cutoff based on the bandlimit of the input field
+        if theta_cutoff is None:
+            theta_cutoff = torch.pi / float(self.nlat_out - 1)
+
+        if theta_cutoff <= 0.0:
+            raise ValueError("Error, theta_cutoff has to be positive.")
+
+        # Note that the psi matrix is of shape nlat_out x nlat_in * nlon_in. Since the contraction in nlon direction is a convolution,
+        # we will keep local to all nodes and split the computation up along nlat. We further split the input dim because this reduces the number
+        # of atomic reduction calls inside the actual kernel
+
+        # set local shapes according to distributed mode:
+        self.nlat_in_local = self.lat_in_shapes[self.comm_rank_polar]
+        self.nlat_out_local = self.nlat_out
+
+        # compute global convolution tensor
+        idx, vals, _ = _precompute_convolution_tensor_s2(
+            in_shape,
+            out_shape,
+            self.filter_basis,
+            grid_in=grid_in,
+            grid_out=grid_out,
+            theta_cutoff=theta_cutoff,
+            transpose_normalization=False,
+            basis_norm_mode=basis_norm_mode,
+            merge_quadrature=True,
+        )
+
+        # split the convolution tensor along latitude
+        idx, vals = _split_distributed_convolution_tensor_s2(idx, vals, in_shape, out_shape)
+
+        # sort the values
+        ker_idx = idx[0, ...].contiguous()
+        row_idx = idx[1, ...].contiguous()
+        col_idx = idx[2, ...].contiguous()
+        vals = vals.contiguous()
+
+        if _cuda_extension_available:
+            # preprocessed data-structure for GPU kernel
+            roff_idx = preprocess_psi(self.kernel_size, self.nlat_out_local, ker_idx, row_idx, col_idx, vals).contiguous()
+            self.register_buffer("psi_roff_idx", roff_idx, persistent=False)
+
+        # save all datastructures
+        self.register_buffer("psi_ker_idx", ker_idx, persistent=False)
+        self.register_buffer("psi_row_idx", row_idx, persistent=False)
+        self.register_buffer("psi_col_idx", col_idx, persistent=False)
+        self.register_buffer("psi_vals", vals, persistent=False)
+
+        # store psi jic:
+        self.psi = _get_psi(self.kernel_size, self.psi_idx, self.psi_vals, self.nlat_in, self.nlon_in, self.nlat_out, self.nlon_out, self.nlat_in_local, self.nlat_out_local)
+
+    def extra_repr(self):
+        return f"in_shape={(self.nlat_in, self.nlon_in)}, out_shape={(self.nlat_out, self.nlon_out)}, in_chans={self.groupsize * self.groups}, out_chans={self.weight.shape[0]}, filter_basis={self.filter_basis}, kernel_shape={self.kernel_shape}, groups={self.groups}"
+
+    @property
+    def psi_idx(self):
+        return torch.stack([self.psi_ker_idx, self.psi_row_idx, self.psi_col_idx], dim=0).contiguous()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+
+        # store number of channels
+        num_chans = x.shape[1]
+
+        # h and w is split. First we make w local by transposing into channel dim
+        if self.comm_size_azimuth > 1:
+            x = distributed_transpose_azimuth.apply(x, (1, -1), self.lon_in_shapes)
+
+        if x.is_cuda and _cuda_extension_available:
+            x = _disco_s2_contraction_cuda(
+                x, self.psi_roff_idx, self.psi_ker_idx, self.psi_row_idx, self.psi_col_idx, self.psi_vals, self.kernel_size, self.nlat_out_local, self.nlon_out
+            )
+        else:
+            if x.is_cuda:
+                warn("couldn't find CUDA extension, falling back to slow PyTorch implementation")
+
+            x = _disco_s2_contraction_torch(x, self.psi.to(x.device), self.nlon_out)
+
+        # perform reduce scatter in polar region
+        x = reduce_from_polar_region(x)
+        x = scatter_to_polar_region(x, -2)
+
+        # now we can transpose back the result, so that lon is split and channels are local
+        if self.comm_size_azimuth > 1:
+            chan_shapes = compute_split_shapes(num_chans, self.comm_size_azimuth)
+            x = distributed_transpose_azimuth.apply(x, (-1, 1), chan_shapes)
+
+        # extract shape
+        B, C, K, H, W = x.shape
+        x = x.reshape(B, self.groups, self.groupsize, K, H, W)
+
+        # do weight multiplication
+        out = torch.einsum("bgckxy,gock->bgoxy", x, self.weight.reshape(self.groups, -1, self.weight.shape[1], self.weight.shape[2])).contiguous()
+        out = out.reshape(out.shape[0], -1, H, W)
+
+        if self.bias is not None:
+            out = out + self.bias.reshape(1, -1, 1, 1)
+
+        return out
+
+
+class DistributedDiscreteContinuousConvTransposeS2(DiscreteContinuousConv):
+    """
+    Discrete-continuous transpose convolutions (DISCO) on the 2-Sphere as described in [1].
+
+    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
+
+    Returns
+    -------
+    out: torch.Tensor
+        Output tensor
+
+    References
+    ----------
+    [1] Ocampo, Price, McEwen, Scalable and equivariant spherical CNNs by discrete-continuous (DISCO) convolutions, ICLR (2023), arXiv:2209.13603
+
+    We assume the data can be splitted in polar and azimuthal directions.
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        in_shape: Tuple[int],
+        out_shape: Tuple[int],
+        kernel_shape: Union[int, Tuple[int], Tuple[int, int]],
+        basis_type: Optional[str] = "piecewise linear",
+        basis_norm_mode: Optional[str] = "mean",
+        groups: Optional[int] = 1,
+        grid_in: Optional[str] = "equiangular",
+        grid_out: Optional[str] = "equiangular",
+        bias: Optional[bool] = True,
+        theta_cutoff: Optional[float] = None,
+    ):
+        super().__init__(in_channels, out_channels, kernel_shape, basis_type, groups, bias)
+
+        self.nlat_in, self.nlon_in = in_shape
+        self.nlat_out, self.nlon_out = out_shape
+
+        # get the comms grid:
+        self.comm_size_polar = polar_group_size()
+        self.comm_rank_polar = polar_group_rank()
+        self.comm_size_azimuth = azimuth_group_size()
+        self.comm_rank_azimuth = azimuth_group_rank()
+
+        # we need those shapes:
+        self.lat_in_shapes = compute_split_shapes(self.nlat_in, self.comm_size_polar)
+        self.lon_in_shapes = compute_split_shapes(self.nlon_in, self.comm_size_azimuth)
+        self.lat_out_shapes = compute_split_shapes(self.nlat_out, self.comm_size_polar)
+        self.lon_out_shapes = compute_split_shapes(self.nlon_out, self.comm_size_azimuth)
+
+        # bandlimit
+        if theta_cutoff is None:
+            theta_cutoff = torch.pi / float(self.nlat_in - 1)
+
+        if theta_cutoff <= 0.0:
+            raise ValueError("Error, theta_cutoff has to be positive.")
+
+        # Note that the psi matrix is of shape nlat_out x nlat_in * nlon_in. Since the contraction in nlon direction is a convolution,
+        # we will keep local to all nodes and split the computation up along nlat. We further split the input dim because this reduces the number
+        # of atomic reduction calls inside the actual kernel
+
+        # set local shapes according to distributed mode:
+        self.nlat_in_local = self.nlat_in
+        self.nlat_out_local = self.lat_out_shapes[self.comm_rank_polar]
+
+        # compute global convolution tensor
+        # switch in_shape and out_shape since we want transpose conv
+        # distributed mode here is swapped because of the transpose
+        idx, vals, _ = _precompute_convolution_tensor_s2(
+            out_shape,
+            in_shape,
+            self.filter_basis,
+            grid_in=grid_out,
+            grid_out=grid_in,
+            theta_cutoff=theta_cutoff,
+            transpose_normalization=True,
+            basis_norm_mode=basis_norm_mode,
+            merge_quadrature=True,
+        )
+
+        # split the convolution tensor along latitude, again, we need to swap the meaning
+        # of in_shape and out_shape
+        idx, vals = _split_distributed_convolution_tensor_s2(idx, vals, out_shape, in_shape)
+
+        # sort the values
+        ker_idx = idx[0, ...].contiguous()
+        row_idx = idx[1, ...].contiguous()
+        col_idx = idx[2, ...].contiguous()
+        vals = vals.contiguous()
+
+        if _cuda_extension_available:
+            # preprocessed data-structure for GPU kernel
+            roff_idx = preprocess_psi(self.kernel_size, self.nlat_in_local, ker_idx, row_idx, col_idx, vals).contiguous()
+            self.register_buffer("psi_roff_idx", roff_idx, persistent=False)
+
+        # save all datastructures
+        self.register_buffer("psi_ker_idx", ker_idx, persistent=False)
+        self.register_buffer("psi_row_idx", row_idx, persistent=False)
+        self.register_buffer("psi_col_idx", col_idx, persistent=False)
+        self.register_buffer("psi_vals", vals, persistent=False)
+
+        # store psi as COO
+        self.psi_st = _get_psi(self.kernel_size, self.psi_idx, self.psi_vals, self.nlat_in, self.nlon_in, self.nlat_out, self.nlon_out, self.nlat_in_local, self.nlat_out_local, semi_transposed=True)
+
+    def extra_repr(self):
+        return f"in_shape={(self.nlat_in, self.nlon_in)}, out_shape={(self.nlat_out, self.nlon_out)}, in_chans={self.groupsize * self.groups}, out_chans={self.weight.shape[0]}, filter_basis={self.filter_basis}, kernel_shape={self.kernel_shape}, groups={self.groups}"
+
+    @property
+    def psi_idx(self):
+        return torch.stack([self.psi_ker_idx, self.psi_row_idx, self.psi_col_idx], dim=0).contiguous()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+
+        # extract shape
+        B, C, H, W = x.shape
+        x = x.reshape(B, self.groups, self.groupsize, H, W)
+
+        # do weight multiplication
+        x = torch.einsum("bgcxy,gock->bgokxy", x, self.weight.reshape(self.groups, -1, self.weight.shape[1], self.weight.shape[2])).contiguous()
+        x = x.reshape(B, -1, x.shape[-3], H, W)
+        num_chans = x.shape[1]
+
+        # transpose such that lon is local, channels are split
+        if self.comm_size_azimuth > 1:
+            x = distributed_transpose_azimuth.apply(x, (1, -1), self.lon_in_shapes)
+
+        # gather input tensor and set up backward reduction hooks
+        x = gather_from_polar_region(x, -2, self.lat_in_shapes)
+        x = copy_to_polar_region(x)
+
+        if x.is_cuda and _cuda_extension_available:
+            out = _disco_s2_transpose_contraction_cuda(
+                x, self.psi_roff_idx, self.psi_ker_idx, self.psi_row_idx, self.psi_col_idx, self.psi_vals, self.kernel_size, self.nlat_out_local, self.nlon_out
+            )
+        else:
+            if x.is_cuda:
+                warn("couldn't find CUDA extension, falling back to slow PyTorch implementation")
+            out = _disco_s2_transpose_contraction_torch(x, self.psi_st.to(x.device), self.nlon_out)
+
+        # now we can transpose back the result, so that lon is split and channels are local
+        if self.comm_size_azimuth > 1:
+            chan_shapes = compute_split_shapes(num_chans, self.comm_size_azimuth)
+            out = distributed_transpose_azimuth.apply(out, (-1, 1), chan_shapes)
+
+        if self.bias is not None:
+            out = out + self.bias.reshape(1, -1, 1, 1)
+
+        return out

.deps/torch_harmonics/distributed/distributed_resample.py

@@ -0,0 +1,247 @@
+# coding=utf-8
+
+# SPDX-FileCopyrightText: Copyright (c) 2022 The torch-harmonics Authors. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+#
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+#
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+#
+
+from typing import List, Tuple, Union, Optional
+import math
+
+import torch
+import torch.nn as nn
+
+from torch_harmonics.quadrature import _precompute_latitudes, _precompute_longitudes
+from torch_harmonics.distributed import polar_group_size, azimuth_group_size, distributed_transpose_azimuth, distributed_transpose_polar
+from torch_harmonics.distributed import reduce_from_azimuth_region, copy_to_azimuth_region
+from torch_harmonics.distributed import polar_group_rank, azimuth_group_rank
+from torch_harmonics.distributed import compute_split_shapes
+
+
+class DistributedResampleS2(nn.Module):
+    """
+    Distributed resampling module for spherical data on the 2-sphere.
+    
+    This module performs distributed resampling of spherical data across multiple processes,
+    supporting both upscaling and downscaling operations. The data is distributed across
+    polar and azimuthal directions, and the module handles the necessary communication
+    and interpolation operations.
+    
+    Parameters
+    -----------
+    nlat_in : int
+        Number of input latitude points
+    nlon_in : int
+        Number of input longitude points
+    nlat_out : int
+        Number of output latitude points
+    nlon_out : int
+        Number of output longitude points
+    grid_in : str, optional
+        Input grid type, by default "equiangular"
+    grid_out : str, optional
+        Output grid type, by default "equiangular"
+    mode : str, optional
+        Interpolation mode ("bilinear" or "bilinear-spherical"), by default "bilinear"
+    """
+
+    def __init__(
+        self,
+        nlat_in: int,
+        nlon_in: int,
+        nlat_out: int,
+        nlon_out: int,
+        grid_in: Optional[str] = "equiangular",
+        grid_out: Optional[str] = "equiangular",
+        mode: Optional[str] = "bilinear",
+    ):
+
+        super().__init__()
+
+        # currently only bilinear is supported
+        if mode in ["bilinear", "bilinear-spherical"]:
+            self.mode = mode
+        else:
+            raise NotImplementedError(f"unknown interpolation mode {mode}")
+
+        self.nlat_in, self.nlon_in = nlat_in, nlon_in
+        self.nlat_out, self.nlon_out = nlat_out, nlon_out
+
+        self.grid_in = grid_in
+        self.grid_out = grid_out
+
+        # get the comms grid:
+        self.comm_size_polar = polar_group_size()
+        self.comm_rank_polar = polar_group_rank()
+        self.comm_size_azimuth = azimuth_group_size()
+        self.comm_rank_azimuth = azimuth_group_rank()
+
+        # compute splits: is this correct even when expanding the poles?
+        self.lat_in_shapes = compute_split_shapes(self.nlat_in, self.comm_size_polar)
+        self.lon_in_shapes = compute_split_shapes(self.nlon_in, self.comm_size_azimuth)
+        self.lat_out_shapes = compute_split_shapes(self.nlat_out, self.comm_size_polar)
+        self.lon_out_shapes = compute_split_shapes(self.nlon_out, self.comm_size_azimuth)
+
+        # for upscaling the latitudes we will use interpolation
+        self.lats_in, _ = _precompute_latitudes(nlat_in, grid=grid_in)
+        self.lons_in = _precompute_longitudes(nlon_in)
+        self.lats_out, _ = _precompute_latitudes(nlat_out, grid=grid_out)
+        self.lons_out = _precompute_longitudes(nlon_out)
+
+        # in the case where some points lie outside of the range spanned by lats_in,
+        # we need to expand the solution to the poles before interpolating
+        self.expand_poles = (self.lats_out > self.lats_in[-1]).any() or (self.lats_out < self.lats_in[0]).any()
+        if self.expand_poles:
+            self.lats_in = torch.cat([torch.tensor([0.], dtype=torch.float64),
+                                      self.lats_in,
+                                      torch.tensor([math.pi], dtype=torch.float64)]).contiguous()
+
+        # prepare the interpolation by computing indices to the left and right of each output latitude
+        lat_idx = torch.searchsorted(self.lats_in, self.lats_out, side="right") - 1
+        # make sure that we properly treat the last point if they coincide with the pole
+        lat_idx = torch.where(self.lats_out == self.lats_in[-1], lat_idx - 1, lat_idx)
+
+        # lat_idx = np.where(self.lats_out > self.lats_in[-1], lat_idx - 1, lat_idx)
+        # lat_idx = np.where(self.lats_out < self.lats_in[0], 0, lat_idx)
+
+        # compute the interpolation weights along the latitude
+        lat_weights = ((self.lats_out - self.lats_in[lat_idx]) / torch.diff(self.lats_in)[lat_idx]).to(torch.float32)
+        lat_weights = lat_weights.unsqueeze(-1)
+
+        # register buffers
+        self.register_buffer("lat_idx", lat_idx, persistent=False)
+        self.register_buffer("lat_weights", lat_weights, persistent=False)
+
+        # get left and right indices but this time make sure periodicity in the longitude is handled
+        lon_idx_left = torch.searchsorted(self.lons_in, self.lons_out, side="right") - 1
+        lon_idx_right = torch.where(self.lons_out >= self.lons_in[-1], torch.zeros_like(lon_idx_left), lon_idx_left + 1)
+
+        # get the difference
+        diff = self.lons_in[lon_idx_right] - self.lons_in[lon_idx_left]
+        diff = torch.where(diff < 0.0, diff + 2 * math.pi, diff)
+        lon_weights = ((self.lons_out - self.lons_in[lon_idx_left]) / diff).to(torch.float32)
+
+        # register buffers
+        self.register_buffer("lon_idx_left", lon_idx_left, persistent=False)
+        self.register_buffer("lon_idx_right", lon_idx_right, persistent=False)
+        self.register_buffer("lon_weights", lon_weights, persistent=False)
+
+        self.skip_resampling = (nlon_in == nlon_out) and (nlat_in == nlat_out) and (grid_in == grid_out)
+
+    def extra_repr(self):
+        return f"in_shape={(self.nlat_in, self.nlon_in)}, out_shape={(self.nlat_out, self.nlon_out)}"
+
+    def _upscale_longitudes(self, x: torch.Tensor):
+        """Upscale the longitude dimension using interpolation."""
+        # do the interpolation
+        lwgt = self.lon_weights.to(x.dtype)
+        if self.mode == "bilinear":
+            x = torch.lerp(x[..., self.lon_idx_left], x[..., self.lon_idx_right], lwgt)
+        else:
+            omega = x[..., self.lon_idx_right] - x[..., self.lon_idx_left]
+            somega = torch.sin(omega)
+            start_prefac = torch.where(somega > 1e-4, torch.sin((1.0 - lwgt) * omega) / somega, (1.0 - lwgt))
+            end_prefac = torch.where(somega > 1e-4, torch.sin(lwgt * omega) / somega, lwgt)
+            x = start_prefac * x[..., self.lon_idx_left] + end_prefac * x[..., self.lon_idx_right]
+
+        return x
+
+    def _expand_poles(self, x: torch.Tensor):
+        """Expand the data to include pole values for interpolation."""
+        x_north = x[...,  0, :].sum(dim=-1, keepdims=True)
+        x_south = x[..., -1, :].sum(dim=-1, keepdims=True)
+        x_count = torch.tensor([x.shape[-1]], dtype=torch.long, device=x.device, requires_grad=False)
+        
+        if self.comm_size_azimuth > 1:
+            x_north = reduce_from_azimuth_region(x_north.contiguous())
+            x_south = reduce_from_azimuth_region(x_south.contiguous())
+            x_count = reduce_from_azimuth_region(x_count)
+        x_north = x_north / x_count
+        x_south = x_south / x_count
+
+        if self.comm_size_azimuth > 1:
+            x_north = copy_to_azimuth_region(x_north)
+            x_south = copy_to_azimuth_region(x_south)
+            
+        x = nn.functional.pad(x, pad=[0, 0, 1, 1], mode='constant')
+        x[..., 0, :] = x_north[...]
+        x[..., -1, :] = x_south[...]
+
+        return x
+
+    def _upscale_latitudes(self, x: torch.Tensor):
+        """Upscale the latitude dimension using interpolation."""
+        # do the interpolation
+        lwgt = self.lat_weights.to(x.dtype)
+        if self.mode == "bilinear":
+            x = torch.lerp(x[..., self.lat_idx, :], x[..., self.lat_idx + 1, :], lwgt)
+        else:
+            omega = x[..., self.lat_idx + 1, :] - x[..., self.lat_idx, :]
+            somega = torch.sin(omega)
+            start_prefac = torch.where(somega > 1e-4, torch.sin((1.0 - lwgt) * omega) / somega, (1.0 - lwgt))
+            end_prefac = torch.where(somega > 1e-4, torch.sin(lwgt * omega) / somega, lwgt)
+            x = start_prefac * x[..., self.lat_idx, :] + end_prefac * x[..., self.lat_idx + 1, :]
+
+        return x
+
+    def forward(self, x: torch.Tensor):
+
+        if self.skip_resampling:
+            return x
+
+        # transpose data so that h is local, and channels are split
+        num_chans = x.shape[-3]
+        
+        # h and w is split. First we make w local by transposing into channel dim
+        if self.comm_size_polar > 1:
+            channels_shapes = compute_split_shapes(num_chans, self.comm_size_polar)
+            x = distributed_transpose_polar.apply(x, (-3, -2), self.lat_in_shapes)
+
+        # expand poles if requested
+        if self.expand_poles:
+            x = self._expand_poles(x)
+
+        # upscaling
+        x = self._upscale_latitudes(x)
+
+        # now, transpose back
+        if self.comm_size_polar > 1:
+            x = distributed_transpose_polar.apply(x, (-2, -3), channels_shapes)
+
+        # now, transpose in w:
+        if self.comm_size_azimuth > 1:
+            channels_shapes = compute_split_shapes(num_chans, self.comm_size_azimuth)
+            x = distributed_transpose_azimuth.apply(x, (-3, -1), self.lon_in_shapes)
+
+        # upscale
+        x = self._upscale_longitudes(x)
+
+        # transpose back
+        if self.comm_size_azimuth > 1:
+            x = distributed_transpose_azimuth.apply(x, (-1, -3), channels_shapes)
+
+        return x

.deps/torch_harmonics/distributed/distributed_sht.py

@@ -0,0 +1,612 @@
+# coding=utf-8
+
+# SPDX-FileCopyrightText: Copyright (c) 2022 The torch-harmonics Authors. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+#
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+#
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+#
+
+import os
+import torch
+import torch.nn as nn
+import torch.fft
+import torch.nn.functional as F
+
+from torch_harmonics.quadrature import legendre_gauss_weights, lobatto_weights, clenshaw_curtiss_weights
+from torch_harmonics.legendre import _precompute_legpoly, _precompute_dlegpoly
+from torch_harmonics.distributed import polar_group_size, azimuth_group_size, distributed_transpose_azimuth, distributed_transpose_polar
+from torch_harmonics.distributed import polar_group_rank, azimuth_group_rank
+from torch_harmonics.distributed import compute_split_shapes, split_tensor_along_dim
+
+
+class DistributedRealSHT(nn.Module):
+    """
+    Defines a module for computing the forward (real-valued) SHT.
+    Precomputes Legendre Gauss nodes, weights and associated Legendre polynomials on these nodes.
+    The SHT is applied to the last two dimensions of the input
+
+    Parameters
+    ----------
+    nlat: int
+        Number of latitude points
+    nlon: int
+        Number of longitude points
+    lmax: int
+        Maximum spherical harmonic degree
+    mmax: int
+        Maximum spherical harmonic order
+    grid: str
+        Grid type ("equiangular", "legendre-gauss", "lobatto", "equidistant"), by default "equiangular"
+    norm: str
+        Normalization type ("ortho", "schmidt", "unnorm"), by default "ortho"
+    csphase: bool
+        Whether to apply the Condon-Shortley phase factor, by default True
+
+    Returns
+    -------
+    x: torch.Tensor
+        Tensor of shape (..., lmax, mmax)
+
+    References
+    ----------
+    [1] Schaeffer, N. Efficient spherical harmonic transforms aimed at pseudospectral numerical simulations, G3: Geochemistry, Geophysics, Geosystems.
+    [2] Wang, B., Wang, L., Xie, Z.; Accurate calculation of spherical and vector spherical harmonic expansions via spectral element grids; Adv Comput Math.
+    """
+
+    def __init__(self, nlat, nlon, lmax=None, mmax=None, grid="equiangular", norm="ortho", csphase=True):
+        
+        super().__init__()
+
+        self.nlat = nlat
+        self.nlon = nlon
+        self.grid = grid
+        self.norm = norm
+        self.csphase = csphase
+
+        # TODO: include assertions regarding the dimensions
+
+        # compute quadrature points
+        if self.grid == "legendre-gauss":
+            cost, weights = legendre_gauss_weights(nlat, -1, 1)
+            self.lmax = lmax or self.nlat
+        elif self.grid == "lobatto":
+            cost, weights = lobatto_weights(nlat, -1, 1)
+            self.lmax = lmax or self.nlat-1
+        elif self.grid == "equiangular":
+            cost, weights = clenshaw_curtiss_weights(nlat, -1, 1)
+            # cost, w = fejer2_weights(nlat, -1, 1)
+            self.lmax = lmax or self.nlat
+        else:
+            raise(ValueError("Unknown quadrature mode"))
+
+        # get the comms grid:
+        self.comm_size_polar = polar_group_size()
+        self.comm_rank_polar = polar_group_rank()
+        self.comm_size_azimuth = azimuth_group_size()
+        self.comm_rank_azimuth = azimuth_group_rank()
+
+        # apply cosine transform and flip them
+        tq = torch.flip(torch.arccos(cost), dims=(0,))
+
+        # determine the dimensions
+        self.mmax = mmax or self.nlon // 2 + 1
+
+        # compute splits
+        self.lat_shapes = compute_split_shapes(self.nlat, self.comm_size_polar)
+        self.lon_shapes = compute_split_shapes(self.nlon, self.comm_size_azimuth)
+        self.l_shapes = compute_split_shapes(self.lmax, self.comm_size_polar)
+        self.m_shapes = compute_split_shapes(self.mmax, self.comm_size_azimuth)
+
+        # combine quadrature weights with the legendre weights
+        pct = _precompute_legpoly(self.mmax, self.lmax, tq, norm=self.norm, csphase=self.csphase)
+        weights = torch.einsum('mlk,k->mlk', pct, weights)
+
+        # split weights
+        weights = split_tensor_along_dim(weights, dim=0, num_chunks=self.comm_size_azimuth)[self.comm_rank_azimuth].contiguous()
+
+        # remember quadrature weights
+        self.register_buffer('weights', weights, persistent=False)
+
+    def extra_repr(self):
+        return f'nlat={self.nlat}, nlon={self.nlon},\n lmax={self.lmax}, mmax={self.mmax},\n grid={self.grid}, csphase={self.csphase}'
+
+    def forward(self, x: torch.Tensor):
+
+        if x.dim() < 3:
+            raise ValueError(f"Expected tensor with at least 3 dimensions but got {x.dim()} instead")
+
+        # we need to ensure that we can split the channels evenly
+        num_chans = x.shape[-3]
+
+        # h and w is split. First we make w local by transposing into channel dim
+        if self.comm_size_azimuth > 1:
+            x = distributed_transpose_azimuth.apply(x, (-3, -1), self.lon_shapes)
+
+        # apply real fft in the longitudinal direction: make sure to truncate to nlon
+        x = 2.0 * torch.pi * torch.fft.rfft(x, n=self.nlon, dim=-1, norm="forward")
+
+        # truncate
+        x = x[..., :self.mmax]
+
+        # transpose: after this, m is split and c is local
+        if self.comm_size_azimuth > 1:
+            chan_shapes = compute_split_shapes(num_chans, self.comm_size_azimuth)
+            x = distributed_transpose_azimuth.apply(x, (-1, -3), chan_shapes)
+
+        # transpose: after this, c is split and h is local
+        if self.comm_size_polar > 1:
+            x = distributed_transpose_polar.apply(x, (-3, -2), self.lat_shapes)
+
+        # do the Legendre-Gauss quadrature
+        x = torch.view_as_real(x)
+
+        # contraction
+        xs = torch.einsum('...kmr,mlk->...lmr', x, self.weights.to(x.dtype)).contiguous()
+
+        # cast to complex
+        x = torch.view_as_complex(xs)
+
+        # transpose: after this, l is split and c is local
+        if self.comm_size_polar	> 1:
+            chan_shapes = compute_split_shapes(num_chans, self.comm_size_polar)
+            x = distributed_transpose_polar.apply(x, (-2, -3), chan_shapes)
+
+        return x
+
+
+class DistributedInverseRealSHT(nn.Module):
+    """
+    Defines a module for computing the inverse (real-valued) SHT.
+    Precomputes Legendre Gauss nodes, weights and associated Legendre polynomials on these nodes.
+
+    Parameters
+    ----------
+    nlat: int
+        Number of latitude points
+    nlon: int
+        Number of longitude points
+    lmax: int
+        Maximum spherical harmonic degree
+    mmax: int
+        Maximum spherical harmonic order
+    grid: str
+        Grid type ("equiangular", "legendre-gauss", "lobatto", "equidistant"), by default "equiangular"
+    norm: str
+        Normalization type ("ortho", "schmidt", "unnorm"), by default "ortho"
+    csphase: bool
+        Whether to apply the Condon-Shortley phase factor, by default True
+
+    Returns
+    -------
+    x: torch.Tensor
+        Tensor of shape (..., lmax, mmax)
+
+    References
+    ----------
+    [1] Schaeffer, N. Efficient spherical harmonic transforms aimed at pseudospectral numerical simulations, G3: Geochemistry, Geophysics, Geosystems.
+    [2] Wang, B., Wang, L., Xie, Z.; Accurate calculation of spherical and vector spherical harmonic expansions via spectral element grids; Adv Comput Math.
+    """
+
+    def __init__(self, nlat, nlon, lmax=None, mmax=None, grid="equiangular", norm="ortho", csphase=True):
+
+        super().__init__()
+
+        self.nlat = nlat
+        self.nlon = nlon
+        self.grid = grid
+        self.norm = norm
+        self.csphase = csphase
+
+        # compute quadrature points
+        if self.grid == "legendre-gauss":
+            cost, _ = legendre_gauss_weights(nlat, -1, 1)
+            self.lmax = lmax or self.nlat
+        elif self.grid == "lobatto":
+            cost, _ = lobatto_weights(nlat, -1, 1)
+            self.lmax = lmax or self.nlat-1
+        elif self.grid == "equiangular":
+            cost, _ = clenshaw_curtiss_weights(nlat, -1, 1)
+            self.lmax = lmax or self.nlat
+        else:
+            raise(ValueError("Unknown quadrature mode"))
+
+        # get the comms grid:
+        self.comm_size_polar = polar_group_size()
+        self.comm_rank_polar = polar_group_rank()
+        self.comm_size_azimuth = azimuth_group_size()
+        self.comm_rank_azimuth = azimuth_group_rank()
+
+        # apply cosine transform and flip them
+        t = torch.flip(torch.arccos(cost), dims=(0,))
+
+        # determine the dimensions
+        self.mmax = mmax or self.nlon // 2 + 1
+
+        # compute splits
+        self.lat_shapes = compute_split_shapes(self.nlat, self.comm_size_polar)
+        self.lon_shapes = compute_split_shapes(self.nlon, self.comm_size_azimuth)
+        self.l_shapes = compute_split_shapes(self.lmax, self.comm_size_polar)
+        self.m_shapes = compute_split_shapes(self.mmax, self.comm_size_azimuth)
+
+        # compute legende polynomials
+        pct = _precompute_legpoly(self.mmax, self.lmax, t, norm=self.norm, inverse=True, csphase=self.csphase)
+
+        # split in m
+        pct = split_tensor_along_dim(pct, dim=0, num_chunks=self.comm_size_azimuth)[self.comm_rank_azimuth].contiguous()
+
+        # register
+        self.register_buffer('pct', pct, persistent=False)
+
+    def extra_repr(self):
+        return f'nlat={self.nlat}, nlon={self.nlon},\n lmax={self.lmax}, mmax={self.mmax},\n grid={self.grid}, csphase={self.csphase}'
+
+    def forward(self, x: torch.Tensor):
+
+        if x.dim() < 3:
+            raise ValueError(f"Expected tensor with at least 3 dimensions but got {x.dim()} instead")
+
+        # we need to ensure that we can split the channels evenly
+        num_chans = x.shape[-3]
+
+        # transpose: after that, channels are split, l is local:
+        if self.comm_size_polar > 1:
+            x = distributed_transpose_polar.apply(x, (-3, -2), self.l_shapes)
+
+        # Evaluate associated Legendre functions on the output nodes
+        x = torch.view_as_real(x)
+
+        # einsum
+        xs = torch.einsum('...lmr, mlk->...kmr', x, self.pct.to(x.dtype)).contiguous()
+
+        # inverse FFT
+        x = torch.view_as_complex(xs)
+
+        if self.comm_size_polar > 1:
+            chan_shapes = compute_split_shapes(num_chans, self.comm_size_polar)
+            x = distributed_transpose_polar.apply(x, (-2, -3), chan_shapes)
+
+        # transpose: after this, channels are split and m is local
+        if self.comm_size_azimuth > 1:
+            x = distributed_transpose_azimuth.apply(x, (-3, -1), self.m_shapes)
+
+        # set DCT and nyquist frequencies to 0:
+        x[..., 0].imag = 0.0
+        if (self.nlon % 2 == 0) and (self.nlon // 2 < x.shape[-1]):
+            x[..., self.nlon // 2].imag = 0.0
+            
+        # apply the inverse (real) FFT
+        x = torch.fft.irfft(x, n=self.nlon, dim=-1, norm="forward")
+
+        # transpose: after this, m is split and channels are local
+        if self.comm_size_azimuth > 1:
+            chan_shapes = compute_split_shapes(num_chans, self.comm_size_azimuth)
+            x = distributed_transpose_azimuth.apply(x, (-1, -3), chan_shapes)
+
+        return x
+
+
+class DistributedRealVectorSHT(nn.Module):
+    """
+    Defines a module for computing the forward (real) vector SHT.
+    Precomputes Legendre Gauss nodes, weights and associated Legendre polynomials on these nodes.
+    The SHT is applied to the last three dimensions of the input.
+
+    Parameters
+    ----------
+    nlat: int
+        Number of latitude points
+    nlon: int
+        Number of longitude points  
+    lmax: int
+        Maximum spherical harmonic degree
+    mmax: int
+        Maximum spherical harmonic order
+    grid: str
+        Grid type ("equiangular", "legendre-gauss", "lobatto", "equidistant"), by default "equiangular"
+    norm: str
+        Normalization type ("ortho", "schmidt", "unnorm"), by default "ortho"
+    csphase: bool
+        Whether to apply the Condon-Shortley phase factor, by default True
+
+    Returns
+    -------
+    x: torch.Tensor
+        Tensor of shape (..., lmax, mmax)
+
+    References
+    ----------
+    [1] Schaeffer, N. Efficient spherical harmonic transforms aimed at pseudospectral numerical simulations, G3: Geochemistry, Geophysics, Geosystems.
+    [2] Wang, B., Wang, L., Xie, Z.; Accurate calculation of spherical and vector spherical harmonic expansions via spectral element grids; Adv Comput Math.
+    """
+
+    def __init__(self, nlat, nlon, lmax=None, mmax=None, grid="equiangular", norm="ortho", csphase=True):
+
+        super().__init__()
+
+        self.nlat = nlat
+        self.nlon = nlon
+        self.grid = grid
+        self.norm = norm
+        self.csphase = csphase
+
+        # compute quadrature points
+        if self.grid == "legendre-gauss":
+            cost, weights = legendre_gauss_weights(nlat, -1, 1)
+            self.lmax = lmax or self.nlat
+        elif self.grid == "lobatto":
+            cost, weights = lobatto_weights(nlat, -1, 1)
+            self.lmax = lmax or self.nlat-1
+        elif self.grid == "equiangular":
+            cost, weights = clenshaw_curtiss_weights(nlat, -1, 1)
+            # cost, w = fejer2_weights(nlat, -1, 1)
+            self.lmax = lmax or self.nlat
+        else:
+            raise(ValueError("Unknown quadrature mode"))
+
+        # get the comms grid:
+        self.comm_size_polar = polar_group_size()
+        self.comm_rank_polar = polar_group_rank()
+        self.comm_size_azimuth = azimuth_group_size()
+        self.comm_rank_azimuth = azimuth_group_rank()
+
+        # apply cosine transform and flip them
+        tq = torch.flip(torch.arccos(cost), dims=(0,))
+
+        # determine the dimensions
+        self.mmax = mmax or self.nlon // 2 + 1
+
+        # compute splits
+        self.lat_shapes = compute_split_shapes(self.nlat, self.comm_size_polar)
+        self.lon_shapes = compute_split_shapes(self.nlon, self.comm_size_azimuth)
+        self.l_shapes = compute_split_shapes(self.lmax, self.comm_size_polar)
+        self.m_shapes = compute_split_shapes(self.mmax, self.comm_size_azimuth)
+
+        # compute weights
+        dpct = _precompute_dlegpoly(self.mmax, self.lmax, tq, norm=self.norm, csphase=self.csphase)
+
+        # combine integration weights, normalization factor in to one:
+        l = torch.arange(0, self.lmax)
+        norm_factor = 1. / l / (l+1)
+        norm_factor[0] = 1.
+        weights = torch.einsum('dmlk,k,l->dmlk', dpct, weights, norm_factor)
+        # since the second component is imaginary, we need to take complex conjugation into account
+        weights[1] = -1 * weights[1]
+
+        # we need to split in m, pad before:
+        weights = split_tensor_along_dim(weights, dim=1, num_chunks=self.comm_size_azimuth)[self.comm_rank_azimuth].contiguous()
+
+        # remember quadrature weights
+        self.register_buffer('weights', weights, persistent=False)
+
+
+    def extra_repr(self):
+        return f'nlat={self.nlat}, nlon={self.nlon},\n lmax={self.lmax}, mmax={self.mmax},\n grid={self.grid}, csphase={self.csphase}'
+
+    def forward(self, x: torch.Tensor):
+
+        if x.dim() < 4:
+            raise ValueError(f"Expected tensor with at least 4 dimensions but got {x.dim()} instead")
+
+        # we need to ensure that we can split the channels evenly
+        num_chans = x.shape[-4]
+
+        # h and w is split. First we make w local by transposing into channel dim
+        if self.comm_size_azimuth > 1:
+            x = distributed_transpose_azimuth.apply(x, (-4, -1), self.lon_shapes)
+
+        # apply real fft in the longitudinal direction: make sure to truncate to nlon
+        x = 2.0 * torch.pi * torch.fft.rfft(x, n=self.nlon, dim=-1, norm="forward")
+
+        # truncate
+        x = x[..., :self.mmax]
+
+        # transpose: after this, m is split and c is local
+        if self.comm_size_azimuth > 1:
+            chan_shapes = compute_split_shapes(num_chans, self.comm_size_azimuth)
+            x = distributed_transpose_azimuth.apply(x, (-1, -4), chan_shapes)
+
+        # transpose: after this, c is split and h is local
+        if self.comm_size_polar > 1:
+            x = distributed_transpose_polar.apply(x, (-4, -2), self.lat_shapes)
+
+        # do the Legendre-Gauss quadrature
+        x = torch.view_as_real(x)
+
+        # create output array
+        xs = torch.zeros_like(x, dtype=x.dtype, device=x.device)
+
+        # contraction - spheroidal component
+        # real component
+        xs[..., 0, :, :, 0] =   torch.einsum('...km,mlk->...lm', x[..., 0, :, :, 0], self.weights[0].to(xs.dtype)) \
+                              - torch.einsum('...km,mlk->...lm', x[..., 1, :, :, 1], self.weights[1].to(xs.dtype))
+        # imag component
+        xs[..., 0, :, :, 1] =   torch.einsum('...km,mlk->...lm', x[..., 0, :, :, 1], self.weights[0].to(xs.dtype)) \
+                              + torch.einsum('...km,mlk->...lm', x[..., 1, :, :, 0], self.weights[1].to(xs.dtype))
+
+        # contraction - toroidal component
+        # real component
+        xs[..., 1, :, :, 0] = - torch.einsum('...km,mlk->...lm', x[..., 0, :, :, 1], self.weights[1].to(xs.dtype)) \
+                              - torch.einsum('...km,mlk->...lm', x[..., 1, :, :, 0], self.weights[0].to(xs.dtype))
+        # imag component
+        xs[..., 1, :, :, 1] =   torch.einsum('...km,mlk->...lm', x[..., 0, :, :, 0], self.weights[1].to(xs.dtype)) \
+                              - torch.einsum('...km,mlk->...lm', x[..., 1, :, :, 1], self.weights[0].to(xs.dtype))
+
+        # pad if required
+        x = torch.view_as_complex(xs)
+
+        # transpose: after this, l is split and c is local
+        if self.comm_size_polar > 1:
+            chan_shapes = compute_split_shapes(num_chans, self.comm_size_polar)
+            x = distributed_transpose_polar.apply(x, (-2, -4), chan_shapes)
+
+        return x
+
+
+class DistributedInverseRealVectorSHT(nn.Module):
+    """
+    Defines a module for computing the inverse (real-valued) vector SHT.
+    Precomputes Legendre Gauss nodes, weights and associated Legendre polynomials on these nodes.
+
+    Parameters
+    ----------
+    nlat: int
+        Number of latitude points
+    nlon: int
+        Number of longitude points
+    lmax: int
+        Maximum spherical harmonic degree
+    mmax: int
+        Maximum spherical harmonic order
+    grid: str
+        Grid type ("equiangular", "legendre-gauss", "lobatto", "equidistant"), by default "equiangular"
+    norm: str
+        Normalization type ("ortho", "schmidt", "unnorm"), by default "ortho"
+    csphase: bool
+        Whether to apply the Condon-Shortley phase factor, by default True
+
+    Returns
+    -------
+    x: torch.Tensor
+        Tensor of shape (..., lmax, mmax)
+
+    References
+    ----------
+    [1] Schaeffer, N. Efficient spherical harmonic transforms aimed at pseudospectral numerical simulations, G3: Geochemistry, Geophysics, Geosystems.
+    [2] Wang, B., Wang, L., Xie, Z.; Accurate calculation of spherical and vector spherical harmonic expansions via spectral element grids; Adv Comput Math.
+    """
+    def __init__(self, nlat, nlon, lmax=None, mmax=None, grid="equiangular", norm="ortho", csphase=True):
+
+        super().__init__()
+
+        self.nlat = nlat
+        self.nlon = nlon
+        self.grid = grid
+        self.norm = norm
+        self.csphase = csphase
+
+        # compute quadrature points
+        if self.grid == "legendre-gauss":
+            cost, _ = legendre_gauss_weights(nlat, -1, 1)
+            self.lmax = lmax or self.nlat
+        elif self.grid == "lobatto":
+            cost, _ = lobatto_weights(nlat, -1, 1)
+            self.lmax = lmax or self.nlat-1
+        elif self.grid == "equiangular":
+            cost, _ = clenshaw_curtiss_weights(nlat, -1, 1)
+            self.lmax = lmax or self.nlat
+        else:
+            raise(ValueError("Unknown quadrature mode"))
+
+        self.comm_size_polar = polar_group_size()
+        self.comm_rank_polar = polar_group_rank()
+        self.comm_size_azimuth = azimuth_group_size()
+        self.comm_rank_azimuth = azimuth_group_rank()
+
+        # apply cosine transform and flip them
+        t = torch.flip(torch.arccos(cost), dims=(0,))
+
+        # determine the dimensions
+        self.mmax = mmax or self.nlon // 2 + 1
+
+        # compute splits
+        self.lat_shapes = compute_split_shapes(self.nlat, self.comm_size_polar)
+        self.lon_shapes = compute_split_shapes(self.nlon, self.comm_size_azimuth)
+        self.l_shapes = compute_split_shapes(self.lmax, self.comm_size_polar)
+        self.m_shapes = compute_split_shapes(self.mmax, self.comm_size_azimuth)
+
+        # compute legende polynomials
+        dpct = _precompute_dlegpoly(self.mmax, self.lmax, t, norm=self.norm, inverse=True, csphase=self.csphase)
+
+        # split in m
+        dpct = split_tensor_along_dim(dpct, dim=1, num_chunks=self.comm_size_azimuth)[self.comm_rank_azimuth].contiguous()
+
+        # register buffer
+        self.register_buffer('dpct', dpct, persistent=False)
+
+    def extra_repr(self):
+        return f'nlat={self.nlat}, nlon={self.nlon},\n lmax={self.lmax}, mmax={self.mmax},\n grid={self.grid}, csphase={self.csphase}'
+
+    def forward(self, x: torch.Tensor):
+
+        if x.dim() < 4:
+            raise ValueError(f"Expected tensor with at least 4 dimensions but got {x.dim()} instead")
+
+        # store num channels
+        num_chans = x.shape[-4]
+
+        # transpose: after that, channels are split, l is local:
+        if self.comm_size_polar > 1:
+            x = distributed_transpose_polar.apply(x, (-4, -2), self.l_shapes)
+
+        # Evaluate associated Legendre functions on the output nodes
+        x = torch.view_as_real(x)
+
+        # contraction - spheroidal component
+        # real component
+        srl =   torch.einsum('...lm,mlk->...km', x[..., 0, :, :, 0], self.dpct[0].to(x.dtype)) \
+              - torch.einsum('...lm,mlk->...km', x[..., 1, :, :, 1], self.dpct[1].to(x.dtype))
+        # imag component
+        sim =   torch.einsum('...lm,mlk->...km', x[..., 0, :, :, 1], self.dpct[0].to(x.dtype)) \
+              + torch.einsum('...lm,mlk->...km', x[..., 1, :, :, 0], self.dpct[1].to(x.dtype))
+
+        # contraction - toroidal component
+        # real component
+        trl = - torch.einsum('...lm,mlk->...km', x[..., 0, :, :, 1], self.dpct[1].to(x.dtype)) \
+              - torch.einsum('...lm,mlk->...km', x[..., 1, :, :, 0], self.dpct[0].to(x.dtype))
+        # imag component
+        tim =   torch.einsum('...lm,mlk->...km', x[..., 0, :, :, 0], self.dpct[1].to(x.dtype)) \
+              - torch.einsum('...lm,mlk->...km', x[..., 1, :, :, 1], self.dpct[0].to(x.dtype))
+
+        # reassemble
+        s = torch.stack((srl, sim), -1)
+        t = torch.stack((trl, tim), -1)
+        xs = torch.stack((s, t), -4)
+
+        # convert to complex
+        x = torch.view_as_complex(xs)
+
+        if self.comm_size_polar > 1:
+            chan_shapes = compute_split_shapes(num_chans, self.comm_size_polar)
+            x = distributed_transpose_polar.apply(x, (-2, -4), chan_shapes)
+
+        # transpose: after this, channels are split and m is local
+        if self.comm_size_azimuth > 1:
+            x = distributed_transpose_azimuth.apply(x, (-4, -1), self.m_shapes)
+
+        # set DCT and nyquist frequencies to zero
+        x[..., 0].imag = 0.0
+        if (self.nlon % 2 == 0) and (self.nlon // 2 < x.shape[-1]):
+            x[..., self.nlon // 2].imag = 0.0
+
+        # apply the inverse (real) FFT
+        x = torch.fft.irfft(x, n=self.nlon, dim=-1, norm="forward")
+
+        # transpose: after this, m is split and channels are local
+        if self.comm_size_azimuth > 1:
+            chan_shapes = compute_split_shapes(num_chans, self.comm_size_azimuth)
+            x = distributed_transpose_azimuth.apply(x, (-1, -4), chan_shapes)
+
+        return x

.deps/torch_harmonics/distributed/primitives.py

@@ -0,0 +1,470 @@
+# coding=utf-8
+
+# SPDX-FileCopyrightText: Copyright (c) 2022 The torch-harmonics Authors. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+# 
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+#
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+#
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+#
+from typing import List
+
+import torch
+import torch.distributed as dist
+from torch.amp import custom_fwd, custom_bwd
+
+from .utils import polar_group, azimuth_group, polar_group_size
+from .utils import is_initialized, is_distributed_polar, is_distributed_azimuth
+
+# helper routine to compute uneven splitting in balanced way:
+def compute_split_shapes(size: int, num_chunks: int) -> List[int]:
+    """Compute the split shapes for a given size and number of chunks."""
+    
+    # treat trivial case first
+    if num_chunks == 1:
+        return [size]
+    
+    # first, check if we can split using div-up to balance the load: 
+    chunk_size = (size + num_chunks - 1) // num_chunks
+    last_chunk_size = max(0, size - chunk_size * (num_chunks - 1))
+    if last_chunk_size == 0:
+        # in this case, the last shard would be empty, split with floor instead:
+        chunk_size = size // num_chunks
+        last_chunk_size = size - chunk_size * (num_chunks-1)
+
+    # generate sections list
+    sections = [chunk_size for _ in range(num_chunks - 1)] + [last_chunk_size]
+
+    return sections
+
+    
+def split_tensor_along_dim(tensor, dim, num_chunks):
+    """Split a tensor along a given dimension into a given number of chunks."""
+    
+    assert dim < tensor.dim(), f"Error, tensor dimension is {tensor.dim()} which cannot be split along {dim}"
+    assert (tensor.shape[dim] >= num_chunks), f"Error, cannot split dim {dim} of size {tensor.shape[dim]} into \
+                                              {num_chunks} chunks. Empty slices are currently not supported."
+    
+    # get split
+    sections = compute_split_shapes(tensor.shape[dim], num_chunks)
+    tensor_list = torch.split(tensor, sections, dim=dim)
+    
+    return tensor_list
+
+
+def _transpose(tensor, dim0, dim1, dim1_split_sizes, group=None, async_op=False):
+    
+    # get comm params
+    comm_size = dist.get_world_size(group=group)
+    comm_rank = dist.get_rank(group=group)
+
+    # split and local transposition
+    tsplit = split_tensor_along_dim(tensor, num_chunks=comm_size, dim=dim0)
+    x_send = [y.contiguous() for y in tsplit]
+    x_send_shapes = [x.shape for x in x_send]
+    x_recv = []
+    x_shape = list(x_send_shapes[comm_rank])
+    for dim1_len in dim1_split_sizes:
+        x_shape[dim1] = dim1_len
+        x_recv.append(torch.empty(x_shape, dtype=tensor.dtype, device=tensor.device))
+        
+    # global transposition
+    req = dist.all_to_all(x_recv, x_send, group=group, async_op=async_op)
+
+    # get dim0 split sizes
+    dim0_split_sizes = [x[dim0] for x in x_send_shapes]
+    
+    return x_recv, dim0_split_sizes, req
+
+
+class distributed_transpose_azimuth(torch.autograd.Function):
+
+    @staticmethod
+    @custom_fwd(device_type="cuda")
+    def forward(ctx, x, dims, dim1_split_sizes):
+
+        # WAR for a potential contig check torch bug for channels last contig tensors
+        xlist, dim0_split_sizes, _ = _transpose(x, dims[0], dims[1], dim1_split_sizes, group=azimuth_group())
+        x = torch.cat(xlist, dim=dims[1])
+        ctx.dims = dims
+        ctx.dim0_split_sizes = dim0_split_sizes
+        
+        return x
+
+    @staticmethod
+    @custom_bwd(device_type="cuda")
+    def backward(ctx, go):
+        dims = ctx.dims
+        dim0_split_sizes = ctx.dim0_split_sizes
+        # WAR for a potential contig check torch bug for channels last contig tensors 
+        gilist, _, _ = _transpose(go, dims[1], dims[0], dim0_split_sizes, group=azimuth_group())
+        gi = torch.cat(gilist, dim=dims[0])
+        
+        return gi, None, None
+
+    
+class distributed_transpose_polar(torch.autograd.Function):
+
+    @staticmethod
+    @custom_fwd(device_type="cuda")
+    def forward(ctx, x, dim, dim1_split_sizes):
+
+        # WAR for a potential contig check torch bug for channels last contig tensors 
+        xlist, dim0_split_sizes, _ = _transpose(x, dim[0], dim[1], dim1_split_sizes, group=polar_group())
+        x = torch.cat(xlist, dim=dim[1])
+        ctx.dim = dim
+        ctx.dim0_split_sizes = dim0_split_sizes
+        return x
+
+    @staticmethod
+    @custom_bwd(device_type="cuda")
+    def backward(ctx, go):
+
+        dim = ctx.dim
+        dim0_split_sizes = ctx.dim0_split_sizes
+        # WAR for a potential contig check torch bug for channels last contig tensors 
+        gilist, _, _ = _transpose(go, dim[1], dim[0], dim0_split_sizes, group=polar_group())
+        gi = torch.cat(gilist, dim=dim[0])
+        return gi, None, None
+
+    
+# we need those additional primitives for distributed matrix multiplications
+def _reduce(input_, use_fp32=True, group=None):
+
+    # Bypass the function if we are using only 1 GPU.
+    if dist.get_world_size(group=group) == 1:
+        return input_
+    
+    # All-reduce.
+    if use_fp32:
+        dtype = input_.dtype
+        inputf_ = input_.float()
+        inputf_ = inputf_.contiguous()
+        dist.all_reduce(inputf_, group=group)
+        input_ = inputf_.to(dtype)
+    else:
+        input_ = input_.contiguous()
+        dist.all_reduce(input_, group=group)
+        
+    return input_
+    
+
+def _split(input_, dim_, group=None):
+    # Bypass the function if we are using only 1 GPU.
+    comm_size = dist.get_world_size(group=group)
+    if comm_size == 1:
+        return input_
+    
+    # Split along last dimension.
+    input_list = split_tensor_along_dim(input_, dim_, comm_size)
+    
+    # Note: torch.split does not create contiguous tensors by default.
+    rank = dist.get_rank(group=group)
+    output = input_list[rank]
+    
+    return output
+
+
+def _gather(input_, dim_, shapes_, group=None):
+    
+    comm_size = dist.get_world_size(group=group)
+
+    if (shapes_ is not None) and (len(shapes_) != comm_size):
+        raise ValueError()
+    if dim_ >= input_.dim():
+        raise ValueError()
+
+    if comm_size == 1:
+        return input_
+
+    # make contiguous:
+    input_ = input_.contiguous()
+    input_shape = list(input_.shape)
+
+    if shapes_ is not None:
+        input_list = []
+        for src in range(comm_size):
+            input_shape[dim_] = shapes_[src]
+            input_list.append(torch.empty(input_shape, dtype=input_.dtype, device=input_.device))
+    else:
+        # assume equal shape on all ranks
+        input_list = [torch.empty_like(input_) for _ in range(comm_size)]
+
+    dist.all_gather(input_list, input_, group=group)
+
+    output = torch.cat(input_list, dim=dim_)
+
+    return output
+
+
+def _reduce_scatter(input_, dim_, use_fp32=True, group=None):
+
+    # Bypass the function if we are using only 1 GPU.
+    if dist.get_world_size(group=group) == 1:
+        return input_
+
+    # make input contiguous
+    comm_size = dist.get_world_size(group=group)
+    comm_rank = dist.get_rank(group=group)
+    input_list = split_tensor_along_dim(input_, dim_, comm_size)
+
+    dtype = input_.dtype
+    if (use_fp32 and (dtype != torch.float32)):
+        input_list = [x.to(torch.float32) for x in input_list]
+
+    input_list = [x.contiguous() for x in input_list]
+
+    # perform reduce_scatter
+    output = torch.empty_like(input_list[comm_rank])
+    dist.reduce_scatter(output, input_list, group=group)
+
+    # convert dtype if necessary
+    if use_fp32:
+        output = output.to(dtype=dtype)
+
+    return output
+
+
+class _CopyToPolarRegion(torch.autograd.Function):
+    
+    @staticmethod
+    def symbolic(graph, input_):
+        return input_
+    
+    @staticmethod
+    @custom_fwd(device_type="cuda")
+    def forward(ctx, input_):
+        
+        return input_
+    
+    @staticmethod
+    @custom_bwd(device_type="cuda")
+    def backward(ctx, grad_output):
+
+        if is_distributed_polar():
+            return _reduce(grad_output, group=polar_group())
+        else:
+            return grad_output, None
+
+
+class _CopyToAzimuthRegion(torch.autograd.Function):
+
+    @staticmethod
+    def symbolic(graph, input_):
+        return input_
+
+    @staticmethod
+    @custom_fwd(device_type="cuda")
+    def forward(ctx, input_):
+
+        return input_
+
+    @staticmethod
+    @custom_bwd(device_type="cuda")
+    def backward(ctx, grad_output):
+
+        if is_distributed_azimuth():
+            return _reduce(grad_output, group=azimuth_group())
+        else:
+            return grad_output, None
+
+
+class _ScatterToPolarRegion(torch.autograd.Function):
+
+    @staticmethod
+    def symbolic(graph, input_, dim_):
+        return _split(input_, dim_, group=polar_group())
+
+    @staticmethod
+    @custom_fwd(device_type="cuda")
+    def forward(ctx, input_, dim_):
+        if is_distributed_polar():
+            ctx.dim = dim_
+            ctx.split_shapes = compute_split_shapes(
+                input_.shape[dim_], polar_group_size()
+            )
+            return _split(input_, dim_, group=polar_group())
+        else:
+            return input_
+
+    @staticmethod
+    @custom_bwd(device_type="cuda")
+    def backward(ctx, grad_output):
+        if is_distributed_polar():
+            return _gather(grad_output, ctx.dim, ctx.split_shapes, polar_group()), None
+        else:
+            return grad_output, None
+
+
+class _GatherFromPolarRegion(torch.autograd.Function):
+
+    @staticmethod
+    def symbolic(graph, input_, dim_, shapes_):
+        return _gather(input_, dim_, shapes_, polar_group())
+
+    @staticmethod
+    @custom_fwd(device_type="cuda")
+    def forward(ctx, input_, dim_, shapes_):
+        if is_distributed_polar():
+            ctx.dim = dim_
+            return _gather(input_, dim_, shapes_, group=polar_group())
+        else:
+            return input_
+
+    @staticmethod
+    @custom_bwd(device_type="cuda")
+    def backward(ctx, grad_output):
+        if is_distributed_polar():
+            return _split(grad_output, ctx.dim, group=polar_group()), None, None
+        else:
+            return grad_output, None, None
+
+    
+class _ReduceFromPolarRegion(torch.autograd.Function):
+    
+    @staticmethod
+    def symbolic(graph, input_):
+        if is_distributed_polar():
+            return _reduce(input_, group=polar_group())
+        else:
+            return input_
+
+    @staticmethod
+    @custom_fwd(device_type="cuda")
+    def forward(ctx, input_):
+        if is_distributed_polar():
+            return _reduce(input_, group=polar_group())
+        else:
+            return input_
+
+    @staticmethod
+    @custom_bwd(device_type="cuda")
+    def backward(ctx, grad_output):
+        return grad_output
+
+    
+class _ReduceFromAzimuthRegion(torch.autograd.Function):
+ 
+    @staticmethod
+    def symbolic(graph, input_):
+        if is_distributed_azimuth():
+            return _reduce(input_, group=azimuth_group())
+        else:
+            return input_
+
+    @staticmethod
+    @custom_fwd(device_type="cuda")
+    def forward(ctx, input_):
+        if is_distributed_azimuth():
+            return _reduce(input_, group=azimuth_group())
+        else:
+            return input_
+
+    @staticmethod
+    @custom_bwd(device_type="cuda")
+    def backward(ctx, grad_output):
+        return grad_output
+
+
+class _ReduceFromScatterToPolarRegion(torch.autograd.Function):
+
+    @staticmethod
+    def symbolic(graph, input_, dim_):
+        if is_distributed_polar():
+            return _reduce_scatter(input_, dim_, group=polar_group())
+        else:
+            return input_
+
+    @staticmethod
+    @custom_fwd(device_type="cuda")
+    def forward(ctx, input_, dim_):
+        if is_distributed_polar():
+            ctx.dim = dim_
+            ctx.split_shapes = compute_split_shapes(
+                input_.shape[dim_], polar_group_size()
+            )
+            return _reduce_scatter(input_, dim_, group=polar_group())
+        else:
+            return input_
+
+    @staticmethod
+    @custom_bwd(device_type="cuda")
+    def backward(ctx, grad_output):
+        if is_distributed_polar():
+            return _gather(grad_output, ctx.dim, ctx.split_shapes, polar_group()), None
+        else:
+            return grad_output, None
+
+
+class _GatherFromCopyToPolarRegion(torch.autograd.Function):
+
+    @staticmethod
+    def symbolic(graph, input_, dim_, shapes_):
+        if is_distributed_polar():
+            return _gather(input_, dim_, shapes_, polar_group())
+        else:
+            return input_
+
+    @staticmethod
+    @custom_fwd(device_type="cuda")
+    def forward(ctx, input_, dim_, shapes_):
+        if is_distributed_polar():
+            ctx.dim = dim_
+            return _gather(input_, dim_, shapes_, group=polar_group())
+        else:
+            return input_
+
+    @staticmethod
+    @custom_bwd(device_type="cuda")
+    def backward(ctx, grad_output):
+        if is_distributed_polar():
+            return _reduce_scatter(grad_output, ctx.dim, use_fp32=True, group=polar_group()), None, None
+        else:
+            return grad_output, None, None
+        
+
+        
+def copy_to_polar_region(input_):
+    return _CopyToPolarRegion.apply(input_)
+
+def copy_to_azimuth_region(input_):
+    return _CopyToAzimuthRegion.apply(input_)
+        
+def reduce_from_polar_region(input_):
+    return _ReduceFromPolarRegion.apply(input_)
+
+def reduce_from_azimuth_region(input_):
+    return _ReduceFromAzimuthRegion.apply(input_)
+
+def scatter_to_polar_region(input_, dim_):
+    return _ScatterToPolarRegion.apply(input_, dim_)
+
+def gather_from_polar_region(input_, dim_, shapes_):
+    return _GatherFromPolarRegion.apply(input_, dim_, shapes_)
+
+def reduce_from_scatter_to_polar_region(input_, dim_):
+    return _ReduceFromScatterToPolarRegion.apply(input_, dim_)
+
+def gather_from_copy_to_polar_region(input_, dim_, shapes_):
+    return _GatherFromCopyToPolarRegion.apply(input_, dim_, shapes_)

.deps/torch_harmonics/distributed/utils.py

@@ -0,0 +1,92 @@
+# coding=utf-8
+
+# SPDX-FileCopyrightText: Copyright (c) 2022 The torch-harmonics Authors. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+# 
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+#
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+#
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+#
+
+# we need this in order to enable distributed
+import torch
+import torch.distributed as dist
+
+# those need to be global
+_POLAR_PARALLEL_GROUP = None
+_AZIMUTH_PARALLEL_GROUP = None
+_IS_INITIALIZED = False
+
+def polar_group():
+    return _POLAR_PARALLEL_GROUP
+
+def azimuth_group():
+    return _AZIMUTH_PARALLEL_GROUP
+
+def init(polar_process_group, azimuth_process_group):
+    global _POLAR_PARALLEL_GROUP
+    global _AZIMUTH_PARALLEL_GROUP
+    _POLAR_PARALLEL_GROUP = polar_process_group
+    _AZIMUTH_PARALLEL_GROUP = azimuth_process_group
+    _IS_INITIALIZED = True
+
+def finalize():
+    if is_initialized():
+        if is_distributed_polar():
+            dist.destroy_process_group(_POLAR_PARALLEL_GROUP)
+        if is_distributed_azimuth():
+            ist.destroy_process_group(_AZIMUTH_PARALLEL_GROUP)
+
+def is_initialized() -> bool:
+    return _IS_INITIALIZED
+
+def is_distributed_polar() -> bool:
+    return (_POLAR_PARALLEL_GROUP is not None)
+
+def is_distributed_azimuth() -> bool:
+    return (_AZIMUTH_PARALLEL_GROUP is not None)
+
+def polar_group_size() -> int:
+    if not is_distributed_polar():
+        return 1
+    else:
+        return dist.get_world_size(group = _POLAR_PARALLEL_GROUP)
+
+def azimuth_group_size() -> int:
+    if not is_distributed_azimuth():
+        return 1
+    else:
+        return dist.get_world_size(group = _AZIMUTH_PARALLEL_GROUP)
+
+def polar_group_rank() -> int:
+    if not is_distributed_polar():
+        return 0
+    else:
+        return dist.get_rank(group = _POLAR_PARALLEL_GROUP)
+
+def azimuth_group_rank() -> int:
+    if not is_distributed_azimuth():
+        return 0
+    else:
+        return dist.get_rank(group = _AZIMUTH_PARALLEL_GROUP)

.deps/torch_harmonics/examples/__init__.py

@@ -0,0 +1,36 @@
+# coding=utf-8
+
+# SPDX-FileCopyrightText: Copyright (c) 2022 The torch-harmonics Authors. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+#
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+#
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+#
+
+from .pde_sphere import SphereSolver
+from .shallow_water_equations import ShallowWaterSolver
+
+from .pde_dataset import PdeDataset
+from .stanford_2d3ds_dataset import StanfordSegmentationDataset, StanfordDepthDataset, Stanford2D3DSDownloader, compute_stats_s2, StanfordDatasetSubset

.deps/torch_harmonics/examples/losses.py

@@ -0,0 +1,417 @@
+# coding=utf-8
+
+# SPDX-FileCopyrightText: Copyright (c) 2025 The torch-harmonics Authors. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+#
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+#
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+#
+
+import torch
+import torch.nn as nn
+import torch.amp as amp
+import torch.nn.functional as F
+from typing import Optional
+from abc import ABC, abstractmethod
+
+from torch_harmonics.quadrature import _precompute_latitudes
+
+
+def get_quadrature_weights(nlat: int, nlon: int, grid: str, tile: bool = False, normalized: bool = True) -> torch.Tensor:
+    # area weights
+    _, q = _precompute_latitudes(nlat=nlat, grid=grid)
+    q = q.reshape(-1, 1) * 2 * torch.pi / nlon
+
+    # numerical precision can be an issue here, make sure it sums to 1:
+    if normalized:
+        q = q / torch.sum(q) / float(nlon)
+
+    if tile:
+        q = torch.tile(q, (1, nlon)).contiguous()
+
+    return q.to(torch.float32)
+
+
+class DiceLossS2(nn.Module):
+    """
+    Dice loss for spherical segmentation tasks.
+
+    Parameters
+    -----------
+    nlat : int
+        Number of latitude points
+    nlon : int
+        Number of longitude points
+    grid : str, optional
+        Grid type, by default "equiangular"
+    weight : torch.Tensor, optional
+        Class weights, by default None
+    smooth : float, optional
+        Smoothing factor, by default 0
+    ignore_index : int, optional
+        Index to ignore in loss computation, by default -100
+    mode : str, optional
+        Aggregation mode ("micro" or "macro"), by default "micro"
+    """
+
+    def __init__(self, nlat: int, nlon: int, grid: str = "equiangular", weight: torch.Tensor = None, smooth: float = 0, ignore_index: int = -100, mode: str = "micro"):
+
+        super().__init__()
+
+        self.smooth = smooth
+        self.ignore_index = ignore_index
+        self.mode = mode
+
+        # area weights
+        q = get_quadrature_weights(nlat=nlat, nlon=nlon, grid=grid)
+        self.register_buffer("quad_weights", q)
+
+        if weight is None:
+            self.weight = None
+        else:
+            self.register_buffer("weight", weight.unsqueeze(0))
+
+    def forward(self, prd: torch.Tensor, tar: torch.Tensor) -> torch.Tensor:
+        prd = nn.functional.softmax(prd, dim=1)
+
+        # mask values
+        if self.ignore_index is not None:
+            mask = torch.where(tar == self.ignore_index, 0, 1)
+            prd = prd * mask.unsqueeze(1)
+            tar = tar * mask
+
+        # one hot encode
+        taroh = nn.functional.one_hot(tar, num_classes=prd.shape[1]).permute(0, 3, 1, 2)
+
+        # compute numerator and denominator
+        intersection = torch.sum((prd * taroh) * self.quad_weights, dim=(-2, -1))
+        union = torch.sum((prd + taroh) * self.quad_weights, dim=(-2, -1))
+
+        if self.mode == "micro":
+            if self.weight is not None:
+                intersection = torch.sum(intersection * self.weight, dim=1)
+                union = torch.sum(union * self.weight, dim=1)
+            else:
+                intersection = torch.mean(intersection, dim=1)
+                union = torch.mean(union, dim=1)
+
+        # compute score
+        dice = (2 * intersection + self.smooth) / (union + self.smooth)
+
+        # compute average over classes
+        if self.mode == "macro":
+            if self.weight is not None:
+                dice = torch.sum(dice * self.weight, dim=1)
+            else:
+                dice = torch.mean(dice, dim=1)
+
+        # average over batch
+        dice = torch.mean(dice)
+
+        return 1 - dice
+
+
+class CrossEntropyLossS2(nn.Module):
+    """
+    Cross-entropy loss for spherical classification tasks.
+
+    Parameters
+    -----------
+    nlat : int
+        Number of latitude points
+    nlon : int
+        Number of longitude points
+    grid : str, optional
+        Grid type, by default "equiangular"
+    weight : torch.Tensor, optional
+        Class weights, by default None
+    smooth : float, optional
+        Label smoothing factor, by default 0
+    ignore_index : int, optional
+        Index to ignore in loss computation, by default -100
+    """
+
+    def __init__(self, nlat: int, nlon: int, grid: str = "equiangular", weight: torch.Tensor = None, smooth: float = 0, ignore_index: int = -100):
+
+        super().__init__()
+
+        self.smooth = smooth
+        self.ignore_index = ignore_index
+
+        if weight is None:
+            self.weight = None
+        else:
+            self.register_buffer("weight", weight)
+
+        q = get_quadrature_weights(nlat=nlat, nlon=nlon, grid=grid)
+        self.register_buffer("quad_weights", q)
+
+    def forward(self, prd: torch.Tensor, tar: torch.Tensor) -> torch.Tensor:
+
+        # compute log softmax
+        logits = nn.functional.log_softmax(prd, dim=1)
+        ce = nn.functional.cross_entropy(logits, tar, weight=self.weight, reduction="none", ignore_index=self.ignore_index, label_smoothing=self.smooth)
+        ce = (ce * self.quad_weights).sum(dim=(-1, -2))
+        ce = torch.mean(ce)
+
+        return ce
+
+
+class FocalLossS2(nn.Module):
+    """
+    Focal loss for spherical classification tasks.
+
+    Parameters
+    -----------
+    nlat : int
+        Number of latitude points
+    nlon : int
+        Number of longitude points
+    grid : str, optional
+        Grid type, by default "equiangular"
+    weight : torch.Tensor, optional
+        Class weights, by default None
+    smooth : float, optional
+        Label smoothing factor, by default 0
+    ignore_index : int, optional
+        Index to ignore in loss computation, by default -100
+    """
+
+    def __init__(self, nlat: int, nlon: int, grid: str = "equiangular", weight: torch.Tensor = None, smooth: float = 0, ignore_index: int = -100):
+
+        super().__init__()
+
+        self.smooth = smooth
+        self.ignore_index = ignore_index
+
+        if weight is None:
+            self.weight = None
+        else:
+            self.register_buffer("weight", weight)
+
+        q = get_quadrature_weights(nlat=nlat, nlon=nlon, grid=grid)
+        self.register_buffer("quad_weights", q)
+
+    def forward(self, prd: torch.Tensor, tar: torch.Tensor, alpha: float = 0.25, gamma: float = 2):
+
+        # compute logits
+        logits = nn.functional.log_softmax(prd, dim=1)
+
+        # w = (1.0 - nn.functional.softmax(prd, dim=-3)).pow(gamma)
+        # w = torch.where(tar == self.ignore_index, 0.0, w.gather(-3, tar.unsqueeze(-3)).squeeze(-3))
+        ce = nn.functional.cross_entropy(logits, tar, weight=self.weight, reduction="none", ignore_index=self.ignore_index, label_smoothing=self.smooth)
+        fl = alpha * (1 - torch.exp(-ce)) ** gamma * ce
+        # fl = w * ce
+        fl = (fl * self.quad_weights).sum(dim=(-1, -2))
+        fl = fl.mean()
+
+        return fl
+
+
+class SphericalLossBase(nn.Module, ABC):
+    """Abstract base class for spherical losses that handles common initialization and integration."""
+
+    def __init__(self, nlat: int, nlon: int, grid: str = "equiangular", normalized: bool = True):
+        super().__init__()
+
+        self.nlat = nlat
+        self.nlon = nlon
+        self.grid = grid
+
+        # get quadrature weights - these sum to 1!
+        q = get_quadrature_weights(nlat=nlat, nlon=nlon, grid=grid, normalized=normalized)
+        self.register_buffer("quad_weights", q)
+
+    def _integrate_sphere(self, ugrid, mask=None):
+        if mask is None:
+            out = torch.sum(ugrid * self.quad_weights, dim=(-2, -1))
+        elif mask is not None:
+            out = torch.sum(mask * ugrid * self.quad_weights, dim=(-2, -1)) / torch.sum(mask * self.quad_weights, dim=(-2, -1))
+        return out
+
+    @abstractmethod
+    def _compute_loss_term(self, prd: torch.Tensor, tar: torch.Tensor) -> torch.Tensor:
+        """Abstract method that must be implemented by child classes to compute loss terms.
+
+        Args:
+            prd (torch.Tensor): Prediction tensor
+            tar (torch.Tensor): Target tensor
+
+        Returns:
+            torch.Tensor: Computed loss term before integration
+        """
+        pass
+
+    def _post_integration_hook(self, loss: torch.Tensor) -> torch.Tensor:
+        """Post-integration hook. Commonly used for the roots in Lp norms"""
+        return loss
+
+    def forward(self, prd: torch.Tensor, tar: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor:
+        """Common forward pass that handles masking and reduction.
+
+        Args:
+            prd (torch.Tensor): Prediction tensor
+            tar (torch.Tensor): Target tensor
+            mask (Optional[torch.Tensor], optional): Mask tensor. Defaults to None.
+
+        Returns:
+            torch.Tensor: Final loss value
+        """
+        loss_term = self._compute_loss_term(prd, tar)
+        # Integrate over the sphere for each item in the batch
+        loss = self._integrate_sphere(loss_term, mask)
+        # potentially call root
+        loss = self._post_integration_hook(loss)
+        # Average the loss over the batch dimension
+        return torch.mean(loss)
+
+
+class SquaredL2LossS2(SphericalLossBase):
+    def _compute_loss_term(self, prd: torch.Tensor, tar: torch.Tensor) -> torch.Tensor:
+        return torch.square(prd - tar)
+
+
+class L1LossS2(SphericalLossBase):
+    def _compute_loss_term(self, prd: torch.Tensor, tar: torch.Tensor) -> torch.Tensor:
+        return torch.abs(prd - tar)
+
+
+class L2LossS2(SquaredL2LossS2):
+    def _post_integration_hook(self, loss: torch.Tensor) -> torch.Tensor:
+        return torch.sqrt(loss)
+
+
+class W11LossS2(SphericalLossBase):
+    def __init__(self, nlat: int, nlon: int, grid: str = "equiangular"):
+        super().__init__(nlat=nlat, nlon=nlon, grid=grid)
+        # Set up grid and domain for FFT
+        l_phi = 2 * torch.pi  # domain size
+        l_theta = torch.pi  # domain size
+
+        k_phi = torch.fft.fftfreq(nlon, d=l_phi / (2 * torch.pi * nlon))
+        k_theta = torch.fft.fftfreq(nlat, d=l_theta / (2 * torch.pi * nlat))
+        k_theta_mesh, k_phi_mesh = torch.meshgrid(k_theta, k_phi, indexing="ij")
+        self.register_buffer("k_phi_mesh", k_phi_mesh)
+        self.register_buffer("k_theta_mesh", k_theta_mesh)
+
+    def _compute_loss_term(self, prd: torch.Tensor, tar: torch.Tensor) -> torch.Tensor:
+        prdtype = prd.dtype
+        with amp.autocast(device_type="cuda", enabled=False):
+            prd = prd.to(torch.float32)
+            prd_prime_fft2_phi_h = torch.fft.ifft2(1j * self.k_phi_mesh * torch.fft.fft2(prd)).real
+            prd_prime_fft2_theta_h = torch.fft.ifft2(1j * self.k_theta_mesh * torch.fft.fft2(prd)).real
+
+            tar_prime_fft2_phi_h = torch.fft.ifft2(1j * self.k_phi_mesh * torch.fft.fft2(tar)).real
+            tar_prime_fft2_theta_h = torch.fft.ifft2(1j * self.k_theta_mesh * torch.fft.fft2(tar)).real
+
+        # Return the element-wise loss term
+        return torch.abs(prd_prime_fft2_phi_h - tar_prime_fft2_phi_h) + torch.abs(prd_prime_fft2_theta_h - tar_prime_fft2_theta_h)
+
+
+class NormalLossS2(SphericalLossBase):
+    """Combined L1 and Surface Normal Consistency Loss for spherical data.
+
+    This loss function combines an L1 loss term with a surface normal alignment term.
+
+    The loss consists of:
+    1. L1 Loss: Absolute difference between predicted and target values
+    2. Normal Consistency Loss: 1 - cosine similarity between surface normals
+       (equivalent to cosine distance between normal vectors)
+
+    Surface normals are computed by calculating gradients in latitude and longitude
+    directions using FFT, then constructing 3D normal vectors that are normalized.
+
+    Parameters
+    ----------
+    nlat : int
+        Number of latitude points
+    nlon : int
+        Number of longitude points
+    grid : str, optional
+        Grid type, by default "equiangular"
+
+    Returns
+    -------
+    torch.Tensor
+        Combined loss term
+    """
+
+    def __init__(self, nlat: int, nlon: int, grid: str = "equiangular"):
+        super().__init__(nlat=nlat, nlon=nlon, grid=grid)
+        # Set up grid and domain for FFT
+        l_phi = 2 * torch.pi  # domain size
+        l_theta = torch.pi  # domain size
+
+        k_phi = torch.fft.fftfreq(nlon, d=l_phi / (2 * torch.pi * nlon))
+        k_theta = torch.fft.fftfreq(nlat, d=l_theta / (2 * torch.pi * nlat))
+        k_theta_mesh, k_phi_mesh = torch.meshgrid(k_theta, k_phi, indexing="ij")
+        self.register_buffer("k_phi_mesh", k_phi_mesh)
+        self.register_buffer("k_theta_mesh", k_theta_mesh)
+
+    def compute_gradients(self, x):
+        # Make sure x is reshaped to have a batch dimension if it's missing
+        if x.dim() == 2:
+            x = x.unsqueeze(0)  # Add batch dimension
+
+        x_prime_fft2_phi_h = torch.fft.ifft2(1j * self.k_phi_mesh * torch.fft.fft2(x)).real
+        x_prime_fft2_theta_h = torch.fft.ifft2(1j * self.k_theta_mesh * torch.fft.fft2(x)).real
+        return x_prime_fft2_theta_h, x_prime_fft2_phi_h
+
+    def compute_normals(self, x):
+        x = x.to(torch.float32)
+        # Ensure x has a batch dimension
+        if x.dim() == 2:
+            x = x.unsqueeze(0)
+
+        grad_lat, grad_lon = self.compute_gradients(x)
+
+        # Create 3D normal vectors
+        ones = torch.ones_like(x)
+        normals = torch.stack([-grad_lon, -grad_lat, ones], dim=1)
+
+        # Normalize along component dimension
+        normals = F.normalize(normals, p=2, dim=1)
+        return normals
+
+    def _compute_loss_term(self, prd: torch.Tensor, tar: torch.Tensor) -> torch.Tensor:
+        # Handle dimensions for both prediction and target
+        # Ensure we have at least a batch dimension
+        if prd.dim() == 2:
+            prd = prd.unsqueeze(0)
+        if tar.dim() == 2:
+            tar = tar.unsqueeze(0)
+
+        # For 4D tensors (batch, channel, height, width), remove channel if it's 1
+        if prd.dim() == 4 and prd.size(1) == 1:
+            prd = prd.squeeze(1)
+        if tar.dim() == 4 and tar.size(1) == 1:
+            tar = tar.squeeze(1)
+
+        pred_normals = self.compute_normals(prd)
+        tar_normals = self.compute_normals(tar)
+
+        # Compute cosine similarity
+        normal_loss = 1 - torch.sum(pred_normals * tar_normals, dim=1, keepdim=True)
+        return normal_loss

.deps/torch_harmonics/examples/metrics.py

@@ -0,0 +1,291 @@
+# coding=utf-8
+
+# SPDX-FileCopyrightText: Copyright (c) 2025 The torch-harmonics Authors. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+#
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+#
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+#
+
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+
+from torch_harmonics.quadrature import _precompute_latitudes
+from .losses import get_quadrature_weights
+
+
+# routine to compute multiclass labels on the sphere
+# the routine follows the implementation in
+# https://github.com/qubvel-org/segmentation_models.pytorch/blob/4aa36c6ad13f8a12552e4ea4131af2a86e564962/segmentation_models_pytorch/metrics/functional.py
+# but uses quadrature weights
+def _get_stats_multiclass(
+    output: torch.LongTensor,
+    target: torch.LongTensor,
+    num_classes: int,
+    quad_weights: torch.Tensor,
+    ignore_index: Optional[int],
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """
+    Compute multiclass statistics (TP, FP, FN, TN) on the sphere using quadrature weights.
+    
+    This function computes true positives, false positives, false negatives, and true negatives
+    for multiclass classification on spherical data, properly weighted by quadrature weights
+    to account for the spherical geometry.
+    
+    Parameters
+    -----------
+    output : torch.LongTensor
+        Predicted class labels
+    target : torch.LongTensor
+        Ground truth class labels
+    num_classes : int
+        Number of classes in the classification task
+    quad_weights : torch.Tensor
+        Quadrature weights for spherical integration
+    ignore_index : Optional[int]
+        Index to ignore in the computation (e.g., for padding or invalid regions)
+        
+    Returns
+    -------
+    Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]
+        Tuple containing (tp_count, fp_count, fn_count, tn_count) for each class
+    """
+    batch_size, *dims = output.shape
+    num_elements = torch.prod(torch.tensor(dims)).long()
+
+    if ignore_index is not None:
+        ignore = target == ignore_index
+        output = torch.where(ignore, -1, output)
+        target = torch.where(ignore, -1, target)
+        ignore_per_sample = ignore.view(batch_size, -1).sum(1)
+
+    tp_count = torch.zeros(batch_size, num_classes, dtype=torch.float32, device=output.device)
+    fp_count = torch.zeros(batch_size, num_classes, dtype=torch.float32, device=output.device)
+    fn_count = torch.zeros(batch_size, num_classes, dtype=torch.float32, device=output.device)
+    tn_count = torch.zeros(batch_size, num_classes, dtype=torch.float32, device=output.device)
+
+    matched = target == output
+    not_matched = target != output
+    for i in range(batch_size):
+        matched_i = matched[i, ...]
+        not_matched_i = not_matched[i, ...]
+        target_i = target[i, ...]
+        output_i = output[i, ...]
+        for c in range(num_classes):
+            # compute weights
+            qwt_c = quad_weights[target_i == c]
+            qwo_c = quad_weights[output_i == c]
+
+            # true positives
+            tp_count[i, c] = torch.sum(matched_i[target_i == c] * qwt_c)
+            # false positives
+            fp_count[i, c] = torch.sum(not_matched_i[output_i == c] * qwo_c)
+            # false negatives
+            fn_count[i, c] = torch.sum(not_matched_i[target_i == c] * qwt_c)
+
+    # true negatives is the leftovers
+    tn_count = torch.sum(quad_weights) - tp_count - fp_count - fn_count
+    return tp_count, fp_count, fn_count, tn_count
+
+
+def _predict_classes(logits: torch.Tensor) -> torch.Tensor:
+ """
+    Convert logits to class predictions using softmax and argmax.
+    
+    Parameters
+    -----------
+    logits : torch.Tensor
+        Input logits tensor
+        
+    Returns
+    -------
+    torch.Tensor
+        Predicted class labels
+    """
+    return torch.argmax(torch.softmax(logits, dim=1), dim=1, keepdim=False)
+
+
+class BaseMetricS2(nn.Module):
+    """
+    Base class for spherical metrics that properly handle spherical geometry.
+    
+    This class provides the foundation for computing metrics on spherical data
+    by using quadrature weights to account for the non-uniform area distribution
+    on the sphere.
+    
+    Parameters
+    -----------
+    nlat : int
+        Number of latitude points
+    nlon : int
+        Number of longitude points
+    grid : str, optional
+        Grid type ("equiangular", "legendre-gauss", etc.), by default "equiangular"
+    weight : torch.Tensor, optional
+        Class weights for weighted averaging, by default None
+    ignore_index : int, optional
+        Index to ignore in computations, by default -100
+    mode : str, optional
+        Averaging mode ("micro" or "macro"), by default "micro"
+    """
+    
+    def __init__(self, nlat: int, nlon: int, grid: str = "equiangular", weight: torch.Tensor = None, ignore_index: int = -100, mode: str = "micro"):
+        super().__init__()
+
+        self.ignore_index = ignore_index
+        self.mode = mode
+
+        # area weights
+        q = get_quadrature_weights(nlat=nlat, nlon=nlon, grid=grid, tile=True)
+        self.register_buffer("quad_weights", q)
+
+        if weight is None:
+            self.weight = None
+        else:
+            self.register_buffer("weight", weight.unsqueeze(0))
+
+    def _forward(self, pred: torch.Tensor, truth: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+
+        # convert logits to class predictions
+        pred_class = _predict_classes(pred)
+
+        # get true positive, false positive, etc
+        tp, fp, fn, tn = _get_stats_multiclass(pred_class, truth, pred.shape[1], self.quad_weights, self.ignore_index)
+
+        # compute averages:
+        if self.mode == "micro":
+            if self.weight is not None:
+                # weighted average
+                tp = torch.sum(tp * self.weight)
+                fp = torch.sum(fp * self.weight)
+                fn = torch.sum(fn * self.weight)
+                tn = torch.sum(tn * self.weight)
+            else:
+                # normal average
+                tp = torch.mean(tp)
+                fp = torch.mean(fp)
+                fn = torch.mean(fn)
+                tn = torch.mean(tn)
+        else:
+            tp = torch.mean(tp, dim=0)
+            fp = torch.mean(fp, dim=0)
+            fn = torch.mean(fn, dim=0)
+            tn = torch.mean(tn, dim=0)
+
+        return tp, fp, fn, tn
+
+
+class IntersectionOverUnionS2(BaseMetricS2):
+    """
+    Intersection over Union (IoU) metric for spherical data.
+    
+    Computes the IoU score for multiclass classification on the sphere,
+    properly weighted by quadrature weights to account for spherical geometry.
+    
+    Parameters
+    -----------
+    nlat : int
+        Number of latitude points
+    nlon : int
+        Number of longitude points
+    grid : str, optional
+        Grid type ("equiangular", "legendre-gauss", etc.), by default "equiangular"
+    weight : torch.Tensor, optional
+        Class weights for weighted averaging, by default None
+    ignore_index : int, optional
+        Index to ignore in computations, by default -100
+    mode : str, optional
+        Averaging mode ("micro" or "macro"), by default "micro"
+    """
+    
+    def __init__(self, nlat: int, nlon: int, grid: str = "equiangular", weight: torch.Tensor = None, ignore_index: int = -100, mode: str = "micro"):
+        super().__init__(nlat, nlon, grid, weight, ignore_index, mode)
+
+    def forward(self, pred: torch.Tensor, truth: torch.Tensor) -> torch.Tensor:
+
+        tp, fp, fn, tn = self._forward(pred, truth)
+
+        # compute score
+        score = tp / (tp + fp + fn)
+
+        if self.mode == "macro":
+            # we need to do some averaging still:
+            # be careful with zeros
+            score = torch.where(torch.isnan(score), 0.0, score)
+
+            if self.weight is not None:
+                score = torch.sum(score * self.weight)
+            else:
+                score = torch.mean(score)
+
+        return score
+
+
+class AccuracyS2(BaseMetricS2):
+    """
+    Accuracy metric for spherical data.
+    
+    Computes the accuracy score for multiclass classification on the sphere,
+    properly weighted by quadrature weights to account for spherical geometry.
+    
+    Parameters
+    -----------
+    nlat : int
+        Number of latitude points
+    nlon : int
+        Number of longitude points
+    grid : str, optional
+        Grid type ("equiangular", "legendre-gauss", etc.), by default "equiangular"
+    weight : torch.Tensor, optional
+        Class weights for weighted averaging, by default None
+    ignore_index : int, optional
+        Index to ignore in computations, by default -100
+    mode : str, optional
+        Averaging mode ("micro" or "macro"), by default "micro"
+    """
+    
+    def __init__(self, nlat: int, nlon: int, grid: str = "equiangular", weight: torch.Tensor = None, ignore_index: int = -100, mode: str = "micro"):
+        super().__init__(nlat, nlon, grid, weight, ignore_index, mode)
+
+    def forward(self, pred: torch.Tensor, truth: torch.Tensor) -> torch.Tensor:
+
+        tp, fp, fn, tn = self._forward(pred, truth)
+
+        # compute score
+        score = (tp + tn) / (tp + fp + fn + tn)
+
+        if self.mode == "macro":
+            # we need to do some averaging still:
+            # be careful with zeros
+            score = torch.where(torch.isnan(score), 0.0, score)
+
+            if self.weight is not None:
+                score = torch.sum(score * self.weight)
+            else:
+                score = torch.mean(score)
+
+        return score

.deps/torch_harmonics/examples/models/__init__.py

@@ -0,0 +1,36 @@
+# coding=utf-8
+
+# SPDX-FileCopyrightText: Copyright (c) 2022 The torch-harmonics Authors. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+#
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+#
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+#
+
+from .sfno import SphericalFourierNeuralOperator
+from .lsno import LocalSphericalNeuralOperator
+from .s2unet import SphericalUNet
+from .s2transformer import SphericalTransformer
+from .s2segformer import SphericalSegformer

.deps/torch_harmonics/examples/models/_layers.py

@@ -0,0 +1,621 @@
+# coding=utf-8
+
+# SPDX-FileCopyrightText: Copyright (c) 2022 The torch-harmonics Authors. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+#
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+#
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+#
+
+import abc
+import math
+
+import torch
+import torch.nn as nn
+import torch.fft
+from torch.utils.checkpoint import checkpoint
+
+from torch_harmonics import InverseRealSHT
+
+
+def _no_grad_trunc_normal_(tensor, mean, std, a, b):
+   
+    # Cut & paste from PyTorch official master until it's in a few official releases - RW
+    # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf
+    def norm_cdf(x):
+       
+        # Computes standard normal cumulative distribution function
+        return (1.0 + math.erf(x / math.sqrt(2.0))) / 2.0
+
+    if (mean < a - 2 * std) or (mean > b + 2 * std):
+        warnings.warn("mean is more than 2 std from [a, b] in nn.init.trunc_normal_. " "The distribution of values may be incorrect.", stacklevel=2)
+
+    with torch.no_grad():
+        # Values are generated by using a truncated uniform distribution and
+        # then using the inverse CDF for the normal distribution.
+        # Get upper and lower cdf values
+        l = norm_cdf((a - mean) / std)
+        u = norm_cdf((b - mean) / std)
+
+        # Uniformly fill tensor with values from [l, u], then translate to
+        # [2l-1, 2u-1].
+        tensor.uniform_(2 * l - 1, 2 * u - 1)
+
+        # Use inverse cdf transform for normal distribution to get truncated
+        # standard normal
+        tensor.erfinv_()
+
+        # Transform to proper mean, std
+        tensor.mul_(std * math.sqrt(2.0))
+        tensor.add_(mean)
+
+        # Clamp to ensure it's in the proper range
+        tensor.clamp_(min=a, max=b)
+        return tensor
+
+
+def trunc_normal_(tensor, mean=0.0, std=1.0, a=-2.0, b=2.0):
+    """Fills the input Tensor with values drawn from a truncated
+    normal distribution. The values are effectively drawn from the
+    normal distribution :math:`\mathcal{N}(\text{mean}, \text{std}^2)`
+    with values outside :math:`[a, b]` redrawn until they are within
+    the bounds. The method used for generating the random values works
+    best when :math:`a \leq \text{mean} \leq b`.
+    
+    Parameters
+    -----------
+    tensor: torch.Tensor
+        an n-dimensional `torch.Tensor`
+    mean: float
+        the mean of the normal distribution
+    std: float
+        the standard deviation of the normal distribution
+    a: float
+        the minimum cutoff value, by default -2.0
+    b: float
+        the maximum cutoff value
+    Examples
+    --------
+    >>> w = torch.empty(3, 5)
+    >>> nn.init.trunc_normal_(w)
+    """
+    return _no_grad_trunc_normal_(tensor, mean, std, a, b)
+
+
+@torch.jit.script
+def drop_path(x: torch.Tensor, drop_prob: float = 0.0, training: bool = False) -> torch.Tensor:
+    """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
+    This is the same as the DropConnect impl I created for EfficientNet, etc networks, however,
+    the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...
+    See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for
+    changing the layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use
+    'survival rate' as the argument.
+
+    Parameters
+    ----------
+    x : torch.Tensor
+        Input tensor
+    drop_prob : float, optional
+        Probability of dropping a path, by default 0.0
+    training : bool, optional
+        Whether the model is in training mode, by default False
+
+    Returns
+    -------
+    torch.Tensor
+        Output tensor
+    """
+    if drop_prob == 0.0 or not training:
+        return x
+    keep_prob = 1.0 - drop_prob
+    shape = (x.shape[0],) + (1,) * (x.ndim - 1)  # work with diff dim tensors, not just 2d ConvNets
+    random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)
+    random_tensor.floor_()  # binarize
+    output = x.div(keep_prob) * random_tensor
+    return output
+
+
+class DropPath(nn.Module):
+    """
+    Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
+    
+    This module implements stochastic depth regularization by randomly dropping
+    entire residual paths during training, which helps with regularization and
+    training of very deep networks.
+    
+    Parameters
+    ----------
+    drop_prob : float, optional
+        Probability of dropping a path, by default None
+    """
+    
+    def __init__(self, drop_prob=None):
+        super(DropPath, self).__init__()
+        self.drop_prob = drop_prob
+
+    def forward(self, x):
+
+        return drop_path(x, self.drop_prob, self.training)
+
+
+class PatchEmbed(nn.Module):
+    """
+    Patch embedding layer for vision transformers.
+    
+    This module splits input images into patches and projects them to a
+    higher dimensional embedding space using convolutional layers.
+    
+    Parameters
+    ----------
+    img_size : tuple, optional
+        Input image size (height, width), by default (224, 224)
+    patch_size : tuple, optional
+        Patch size (height, width), by default (16, 16)
+    in_chans : int, optional
+        Number of input channels, by default 3
+    embed_dim : int, optional
+        Embedding dimension, by default 768
+    """
+    
+    def __init__(self, img_size=(224, 224), patch_size=(16, 16), in_chans=3, embed_dim=768):
+        super(PatchEmbed, self).__init__()
+        self.red_img_size = ((img_size[0] // patch_size[0]), (img_size[1] // patch_size[1]))
+        num_patches = self.red_img_size[0] * self.red_img_size[1]
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.num_patches = num_patches
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size, bias=True)
+        self.proj.weight.is_shared_mp = ["spatial"]
+        self.proj.bias.is_shared_mp = ["spatial"]
+
+    def forward(self, x):
+
+        # gather input
+        B, C, H, W = x.shape
+        assert H == self.img_size[0] and W == self.img_size[1], f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
+        # new: B, C, H*W
+        x = self.proj(x).flatten(2)
+        return x
+
+
+class MLP(nn.Module):
+    """
+    Multi-layer perceptron with optional checkpointing.
+    
+    This module implements a feed-forward network with two linear layers
+    and an activation function, with optional dropout and gradient checkpointing.
+    
+    Parameters
+    ----------
+    in_features : int
+        Number of input features
+    hidden_features : int, optional
+        Number of hidden features, by default None (same as in_features)
+    out_features : int, optional
+        Number of output features, by default None (same as in_features)
+    act_layer : nn.Module, optional
+        Activation layer, by default nn.ReLU
+    output_bias : bool, optional
+        Whether to use bias in output layer, by default False
+    drop_rate : float, optional
+        Dropout rate, by default 0.0
+    checkpointing : bool, optional
+        Whether to use gradient checkpointing, by default False
+    gain : float, optional
+        Gain factor for weight initialization, by default 1.0
+    """
+    
+    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.ReLU, output_bias=False, drop_rate=0.0, checkpointing=False, gain=1.0):
+        super(MLP, self).__init__()
+        self.checkpointing = checkpointing
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+
+        # Fist dense layer
+        fc1 = nn.Conv2d(in_features, hidden_features, 1, bias=True)
+        # initialize the weights correctly
+        scale = math.sqrt(2.0 / in_features)
+        nn.init.normal_(fc1.weight, mean=0.0, std=scale)
+        if fc1.bias is not None:
+            nn.init.constant_(fc1.bias, 0.0)
+
+        # activation
+        act = act_layer()
+
+        # output layer
+        fc2 = nn.Conv2d(hidden_features, out_features, 1, bias=output_bias)
+        # gain factor for the output determines the scaling of the output init
+        scale = math.sqrt(gain / hidden_features)
+        nn.init.normal_(fc2.weight, mean=0.0, std=scale)
+        if fc2.bias is not None:
+            nn.init.constant_(fc2.bias, 0.0)
+
+        if drop_rate > 0.0:
+            drop = nn.Dropout2d(drop_rate)
+            self.fwd = nn.Sequential(fc1, act, drop, fc2, drop)
+        else:
+            self.fwd = nn.Sequential(fc1, act, fc2)
+
+    @torch.jit.ignore
+    def checkpoint_forward(self, x):
+
+        return checkpoint(self.fwd, x)
+
+    def forward(self, x):
+
+        if self.checkpointing:
+            return self.checkpoint_forward(x)
+        else:
+            return self.fwd(x)
+
+
+class RealFFT2(nn.Module):
+    """
+    Helper routine to wrap FFT similarly to the SHT.
+    
+    This module provides a wrapper around PyTorch's real FFT2D that mimics
+    the interface of spherical harmonic transforms for consistency.
+    
+    Parameters
+    -----------
+    nlat : int
+        Number of latitude points
+    nlon : int
+        Number of longitude points
+    lmax : int, optional
+        Maximum spherical harmonic degree, by default None (same as nlat)
+    mmax : int, optional
+        Maximum spherical harmonic order, by default None (nlon//2 + 1)
+    """
+    
+    def __init__(self, nlat, nlon, lmax=None, mmax=None):
+        super(RealFFT2, self).__init__()
+
+        self.nlat = nlat
+        self.nlon = nlon
+        self.lmax = lmax or self.nlat
+        self.mmax = mmax or self.nlon // 2 + 1
+
+    def forward(self, x):
+
+        y = torch.fft.rfft2(x, dim=(-2, -1), norm="ortho")
+        y = torch.cat((y[..., : math.ceil(self.lmax / 2), : self.mmax], y[..., -math.floor(self.lmax / 2) :, : self.mmax]), dim=-2)
+        return y
+
+
+class InverseRealFFT2(nn.Module):
+    """
+    Helper routine to wrap inverse FFT similarly to the SHT.
+    
+    This module provides a wrapper around PyTorch's inverse real FFT2D that mimics
+    the interface of inverse spherical harmonic transforms for consistency.
+    
+    Parameters
+    -----------
+    nlat : int
+        Number of latitude points
+    nlon : int
+        Number of longitude points
+    lmax : int, optional
+        Maximum spherical harmonic degree, by default None (same as nlat)
+    mmax : int, optional
+        Maximum spherical harmonic order, by default None (nlon//2 + 1)
+    """
+    
+    def __init__(self, nlat, nlon, lmax=None, mmax=None):
+        super(InverseRealFFT2, self).__init__()
+
+        self.nlat = nlat
+        self.nlon = nlon
+        self.lmax = lmax or self.nlat
+        self.mmax = mmax or self.nlon // 2 + 1
+
+    def forward(self, x):
+
+        return torch.fft.irfft2(x, dim=(-2, -1), s=(self.nlat, self.nlon), norm="ortho")
+
+
+class LayerNorm(nn.Module):
+    """
+    Wrapper class that moves the channel dimension to the end.
+    
+    This module provides a layer normalization that works with channel-first
+    tensors by temporarily transposing the channel dimension to the end,
+    applying normalization, and then transposing back.
+    
+    Parameters
+    ----------
+    in_channels : int
+        Number of input channels
+    eps : float, optional
+        Epsilon for numerical stability, by default 1e-05
+    elementwise_affine : bool, optional
+        Whether to use learnable affine parameters, by default True
+    bias : bool, optional
+        Whether to use bias, by default True
+    device : torch.device, optional
+        Device to place the module on, by default None
+    dtype : torch.dtype, optional
+        Data type for the module, by default None
+    """
+    
+    def __init__(self, in_channels, eps=1e-05, elementwise_affine=True, bias=True, device=None, dtype=None):
+        super().__init__()
+
+        self.channel_dim = -3
+
+        self.norm = nn.LayerNorm(normalized_shape=in_channels, eps=1e-6, elementwise_affine=elementwise_affine, bias=bias, device=device, dtype=dtype)
+
+    def forward(self, x):
+
+        return self.norm(x.transpose(self.channel_dim, -1)).transpose(-1, self.channel_dim)
+
+
+class SpectralConvS2(nn.Module):
+    """
+    Spectral Convolution according to Driscoll & Healy. Designed for convolutions on the two-sphere S2
+    using the Spherical Harmonic Transforms in torch-harmonics, but supports convolutions on the periodic
+    domain via the RealFFT2 and InverseRealFFT2 wrappers.
+    
+    Parameters
+    ----------
+    forward_transform : nn.Module
+        Forward transform (SHT or FFT)
+    inverse_transform : nn.Module
+        Inverse transform (ISHT or IFFT)
+    in_channels : int
+        Number of input channels
+    out_channels : int
+        Number of output channels
+    gain : float, optional
+        Gain factor for weight initialization, by default 2.0
+    operator_type : str, optional
+        Type of spectral operator ("driscoll-healy", "diagonal", "block-diagonal"), by default "driscoll-healy"
+    lr_scale_exponent : int, optional
+        Learning rate scaling exponent, by default 0
+    bias : bool, optional
+        Whether to use bias, by default False
+    """
+    
+    def __init__(self, forward_transform, inverse_transform, in_channels, out_channels, gain=2.0, operator_type="driscoll-healy", lr_scale_exponent=0, bias=False):
+        super().__init__()
+
+        self.forward_transform = forward_transform
+        self.inverse_transform = inverse_transform
+
+        self.modes_lat = self.inverse_transform.lmax
+        self.modes_lon = self.inverse_transform.mmax
+
+        self.scale_residual = (self.forward_transform.nlat != self.inverse_transform.nlat) or (self.forward_transform.nlon != self.inverse_transform.nlon)
+
+        # remember factorization details
+        self.operator_type = operator_type
+
+        assert self.inverse_transform.lmax == self.modes_lat
+        assert self.inverse_transform.mmax == self.modes_lon
+
+        weight_shape = [out_channels, in_channels]
+
+        if self.operator_type == "diagonal":
+            weight_shape += [self.modes_lat, self.modes_lon]
+            self.contract_func = "...ilm,oilm->...olm"
+        elif self.operator_type == "block-diagonal":
+            weight_shape += [self.modes_lat, self.modes_lon, self.modes_lon]
+            self.contract_func = "...ilm,oilnm->...oln"
+        elif self.operator_type == "driscoll-healy":
+            weight_shape += [self.modes_lat]
+            self.contract_func = "...ilm,oil->...olm"
+        else:
+            raise NotImplementedError(f"Unkonw operator type f{self.operator_type}")
+
+        # form weight tensors
+        scale = math.sqrt(gain / in_channels)
+        self.weight = nn.Parameter(scale * torch.randn(*weight_shape, dtype=torch.complex64))
+        if bias:
+            self.bias = nn.Parameter(torch.zeros(1, out_channels, 1, 1))
+
+    def forward(self, x):
+
+        dtype = x.dtype
+        x = x.float()
+        residual = x
+
+        with torch.autocast(device_type="cuda", enabled=False):
+            x = self.forward_transform(x)
+            if self.scale_residual:
+                residual = self.inverse_transform(x)
+
+        x = torch.einsum(self.contract_func, x, self.weight)
+
+        with torch.autocast(device_type="cuda", enabled=False):
+            x = self.inverse_transform(x)
+
+        if hasattr(self, "bias"):
+            x = x + self.bias
+        x = x.type(dtype)
+
+        return x, residual
+
+
+class PositionEmbedding(nn.Module, metaclass=abc.ABCMeta):
+    """
+    Abstract base class for position embeddings.
+    
+    This class defines the interface for position embedding modules
+    that add positional information to input tensors.
+    
+    Parameters
+    ----------
+    img_shape : tuple, optional
+        Image shape (height, width), by default (480, 960)
+    grid : str, optional
+        Grid type, by default "equiangular"
+    num_chans : int, optional
+        Number of channels, by default 1
+    """
+    
+    def __init__(self, img_shape=(480, 960), grid="equiangular", num_chans=1):
+        super().__init__()
+
+        self.img_shape = img_shape
+        self.num_chans = num_chans
+
+    def forward(self, x: torch.Tensor):
+
+        return x + self.position_embeddings
+
+
+class SequencePositionEmbedding(PositionEmbedding):
+    """
+    Standard sequence-based position embedding.
+    
+    This module implements sinusoidal position embeddings similar to those
+    used in the original Transformer paper, adapted for 2D spatial data.
+    
+    Parameters
+    ----------
+    img_shape : tuple, optional
+        Image shape (height, width), by default (480, 960)
+    grid : str, optional
+        Grid type, by default "equiangular"
+    num_chans : int, optional
+        Number of channels, by default 1
+    """
+    
+    def __init__(self, img_shape=(480, 960), grid="equiangular", num_chans=1):
+        super().__init__(img_shape=img_shape, grid=grid, num_chans=num_chans)
+
+        with torch.no_grad():
+            # alternating custom position embeddings
+            pos = torch.arange(self.img_shape[0] * self.img_shape[1]).reshape(1, 1, *self.img_shape).repeat(1, self.num_chans, 1, 1)
+            k = torch.arange(self.num_chans).reshape(1, self.num_chans, 1, 1)
+            denom = torch.pow(10000, 2 * k / self.num_chans)
+
+            pos_embed = torch.where(k % 2 == 0, torch.sin(pos / denom), torch.cos(pos / denom))
+
+        # register tensor
+        self.register_buffer("position_embeddings", pos_embed.float())
+
+
+class SpectralPositionEmbedding(PositionEmbedding):
+    """
+    Spectral position embeddings for spherical transformers.
+    
+    This module creates position embeddings in the spectral domain using
+    spherical harmonic functions, which are particularly suitable for
+    spherical data processing.
+    
+    Parameters
+    -----------
+    img_shape : tuple, optional
+        Image shape (height, width), by default (480, 960)
+    grid : str, optional
+        Grid type, by default "equiangular"
+    num_chans : int, optional
+        Number of channels, by default 1
+    """
+    
+    def __init__(self, img_shape=(480, 960), grid="equiangular", num_chans=1):
+        super().__init__(img_shape=img_shape, grid=grid, num_chans=num_chans)
+
+        # compute maximum required frequency and prepare isht
+        lmax = math.floor(math.sqrt(self.num_chans)) + 1
+        isht = InverseRealSHT(nlat=self.img_shape[0], nlon=self.img_shape[1], lmax=lmax, mmax=lmax, grid=grid)
+
+        # fill position embedding
+        with torch.no_grad():
+            pos_embed_freq = torch.zeros(1, self.num_chans, isht.lmax, isht.mmax, dtype=torch.complex64)
+
+            for i in range(self.num_chans):
+                l = math.floor(math.sqrt(i))
+                m = i - l**2 - l
+
+                if m < 0:
+                    pos_embed_freq[0, i, l, -m] = 1.0j
+                else:
+                    pos_embed_freq[0, i, l, m] = 1.0
+
+        # compute spatial position embeddings
+        pos_embed = isht(pos_embed_freq)
+
+        # normalization
+        pos_embed = pos_embed / torch.amax(pos_embed.abs(), dim=(-1, -2), keepdim=True)
+
+        # register tensor
+        self.register_buffer("position_embeddings", pos_embed)
+
+
+class LearnablePositionEmbedding(PositionEmbedding):
+    """
+    Learnable position embeddings for spherical transformers.
+    
+    This module provides learnable position embeddings that can be either
+    latitude-only or full latitude-longitude embeddings.
+    
+    Parameters
+    ----------
+    img_shape : tuple, optional
+        Image shape (height, width), by default (480, 960)
+    grid : str, optional
+        Grid type, by default "equiangular"
+    num_chans : int, optional
+        Number of channels, by default 1
+    embed_type : str, optional
+        Embedding type ("lat" or "latlon"), by default "lat"
+    """
+    
+    def __init__(self, img_shape=(480, 960), grid="equiangular", num_chans=1, embed_type="lat"):
+        super().__init__(img_shape=img_shape, grid=grid, num_chans=num_chans)
+
+        if embed_type == "latlon":
+            self.position_embeddings = nn.Parameter(torch.zeros(1, self.num_chans, self.img_shape[0], self.img_shape[1]))
+        elif embed_type == "lat":
+            self.position_embeddings = nn.Parameter(torch.zeros(1, self.num_chans, self.img_shape[0], 1))
+        else:
+            raise ValueError(f"Unknown learnable position embedding type {embed_type}")
+
+# class SpiralPositionEmbedding(PositionEmbedding):
+#     """
+#     Returns position embeddings on the torus
+#     """
+
+#     def __init__(self, img_shape=(480, 960), grid="equiangular", num_chans=1):
+
+#         super().__init__(img_shape=img_shape, grid=grid, num_chans=num_chans)
+
+#         with torch.no_grad():
+
+#             # alternating custom position embeddings
+#             lats, _ = _precompute_latitudes(img_shape[0], grid=grid)
+#             lats = lats.reshape(-1, 1)
+#             lons = torch.linspace(0, 2 * math.pi, img_shape[1] + 1)[:-1]
+#             lons = lons.reshape(1, -1)
+
+#             # channel index
+#             k = torch.arange(self.num_chans).reshape(1, -1, 1, 1)
+#             pos_embed = torch.where(k % 2 == 0, torch.sin(k * (lons + lats)), torch.cos(k * (lons - lats)))
+
+#         # register tensor
+#         self.register_buffer("position_embeddings", pos_embed.float())

.deps/torch_harmonics/examples/models/lsno.py

@@ -0,0 +1,623 @@
+# coding=utf-8
+
+# SPDX-FileCopyrightText: Copyright (c) 2024 The torch-harmonics Authors. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+#
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+#
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+#
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.amp as amp
+
+from torch_harmonics import RealSHT, InverseRealSHT
+from torch_harmonics import DiscreteContinuousConvS2, DiscreteContinuousConvTransposeS2
+from torch_harmonics import ResampleS2
+
+from torch_harmonics.examples.models._layers import MLP, SpectralConvS2, SequencePositionEmbedding, SpectralPositionEmbedding, LearnablePositionEmbedding
+
+from functools import partial
+
+# heuristic for finding theta_cutoff
+def _compute_cutoff_radius(nlat, kernel_shape, basis_type):
+    theta_cutoff_factor = {"piecewise linear": 0.5, "morlet": 0.5, "zernike": math.sqrt(2.0)}
+
+    return (kernel_shape[0] + 1) * theta_cutoff_factor[basis_type] * math.pi / float(nlat - 1)
+
+class DiscreteContinuousEncoder(nn.Module):
+    """
+    Discrete-continuous encoder for spherical neural operators.
+    
+    This module performs downsampling using discrete-continuous convolutions on the sphere,
+    reducing the spatial resolution while maintaining the spectral properties of the data.
+    
+    Parameters
+    ----------
+    in_shape : tuple, optional
+        Input shape (nlat, nlon), by default (721, 1440)
+    out_shape : tuple, optional
+        Output shape (nlat, nlon), by default (480, 960)
+    grid_in : str, optional
+        Input grid type, by default "equiangular"
+    grid_out : str, optional
+        Output grid type, by default "equiangular"
+    inp_chans : int, optional
+        Number of input channels, by default 2
+    out_chans : int, optional
+        Number of output channels, by default 2
+    kernel_shape : tuple, optional
+        Kernel shape for convolution, by default (3, 3)
+    basis_type : str, optional
+        Filter basis type, by default "morlet"
+    groups : int, optional
+        Number of groups for grouped convolution, by default 1
+    bias : bool, optional
+        Whether to use bias, by default False
+    """
+    def __init__(
+        self,
+        in_shape=(721, 1440),
+        out_shape=(480, 960),
+        grid_in="equiangular",
+        grid_out="equiangular",
+        inp_chans=2,
+        out_chans=2,
+        kernel_shape=(3, 3),
+        basis_type="morlet",
+        groups=1,
+        bias=False,
+    ):
+        super().__init__()
+
+        # set up local convolution
+        self.conv = DiscreteContinuousConvS2(
+            inp_chans,
+            out_chans,
+            in_shape=in_shape,
+            out_shape=out_shape,
+            kernel_shape=kernel_shape,
+            basis_type=basis_type,
+            grid_in=grid_in,
+            grid_out=grid_out,
+            groups=groups,
+            bias=bias,
+            theta_cutoff=_compute_cutoff_radius(in_shape[0], kernel_shape, basis_type),
+        )
+
+    def forward(self, x):
+
+        dtype = x.dtype
+
+        with amp.autocast(device_type="cuda", enabled=False):
+            x = x.float()
+            x = self.conv(x)
+            x = x.to(dtype=dtype)
+
+        return x
+
+
+class DiscreteContinuousDecoder(nn.Module):
+    """
+    Discrete-continuous decoder for spherical neural operators.
+    
+    This module performs upsampling using either spherical harmonic transforms or resampling,
+    followed by discrete-continuous convolutions to restore spatial resolution.
+    
+    Parameters
+    ----------
+    in_shape : tuple, optional
+        Input shape (nlat, nlon), by default (480, 960)
+    out_shape : tuple, optional
+        Output shape (nlat, nlon), by default (721, 1440)
+    grid_in : str, optional
+        Input grid type, by default "equiangular"
+    grid_out : str, optional
+        Output grid type, by default "equiangular"
+    inp_chans : int, optional
+        Number of input channels, by default 2
+    out_chans : int, optional
+        Number of output channels, by default 2
+    kernel_shape : tuple, optional
+        Kernel shape for convolution, by default (3, 3)
+    basis_type : str, optional
+        Filter basis type, by default "morlet"
+    groups : int, optional
+        Number of groups for grouped convolution, by default 1
+    bias : bool, optional
+        Whether to use bias, by default False
+    upsample_sht : bool, optional
+        Whether to use SHT for upsampling, by default False
+    """
+    def __init__(
+        self,
+        in_shape=(480, 960),
+        out_shape=(721, 1440),
+        grid_in="equiangular",
+        grid_out="equiangular",
+        inp_chans=2,
+        out_chans=2,
+        kernel_shape=(3, 3),
+        basis_type="morlet",
+        groups=1,
+        bias=False,
+        upsample_sht=False,
+    ):
+        super().__init__()
+
+        # set up upsampling
+        if upsample_sht:
+            self.sht = RealSHT(*in_shape, grid=grid_in).float()
+            self.isht = InverseRealSHT(*out_shape, lmax=self.sht.lmax, mmax=self.sht.mmax, grid=grid_out).float()
+            self.upsample = nn.Sequential(self.sht, self.isht)
+        else:
+            self.upsample = ResampleS2(*in_shape, *out_shape, grid_in=grid_in, grid_out=grid_out)
+
+        # set up DISCO convolution
+        self.conv = DiscreteContinuousConvS2(
+            inp_chans,
+            out_chans,
+            in_shape=out_shape,
+            out_shape=out_shape,
+            kernel_shape=kernel_shape,
+            basis_type=basis_type,
+            grid_in=grid_out,
+            grid_out=grid_out,
+            groups=groups,
+            bias=False,
+            theta_cutoff=_compute_cutoff_radius(in_shape[0], kernel_shape, basis_type),
+        )
+
+    def forward(self, x):
+
+        dtype = x.dtype
+
+        with amp.autocast(device_type="cuda", enabled=False):
+            x = x.float()
+            x = self.upsample(x)
+            x = self.conv(x)
+            x = x.to(dtype=dtype)
+
+        return x
+
+
+class SphericalNeuralOperatorBlock(nn.Module):
+    """
+    Helper module for a single SFNO/FNO block. Can use both FFTs and SHTs to represent either FNO or SFNO blocks.
+
+    Parameters
+    ----------
+    forward_transform : torch.nn.Module
+        Forward transform to use for the block
+    inverse_transform : torch.nn.Module
+        Inverse transform to use for the block
+    input_dim : int
+        Input dimension
+    output_dim : int
+        Output dimension
+    conv_type : str, optional
+        Type of convolution to use, by default "local"
+    mlp_ratio : float, optional
+        MLP expansion ratio, by default 2.0
+    drop_rate : float, optional
+        Dropout rate, by default 0.0
+    drop_path : float, optional
+        Drop path rate, by default 0.0
+    act_layer : torch.nn.Module, optional
+        Activation function to use, by default nn.GELU
+    norm_layer : str, optional
+        Type of normalization to use, by default "none"
+    inner_skip : str, optional
+        Type of inner skip connection to use, by default "none"
+    outer_skip : str, optional
+        Type of outer skip connection to use, by default "identity"
+    use_mlp : bool, optional
+        Whether to use MLP layers, by default True
+    disco_kernel_shape : tuple, optional
+        Kernel shape for discrete-continuous convolution, by default (3, 3)
+    disco_basis_type : str, optional
+        Filter basis type for discrete-continuous convolution, by default "morlet"
+    bias : bool, optional
+        Whether to use bias, by default False
+
+    Returns
+    -------
+    torch.Tensor
+        Output tensor
+    """
+
+    def __init__(
+        self,
+        forward_transform,
+        inverse_transform,
+        input_dim,
+        output_dim,
+        conv_type="local",
+        mlp_ratio=2.0,
+        drop_rate=0.0,
+        drop_path=0.0,
+        act_layer=nn.GELU,
+        norm_layer="none",
+        inner_skip="none",
+        outer_skip="identity",
+        use_mlp=True,
+        disco_kernel_shape=(3, 3),
+        disco_basis_type="morlet",
+        bias=False,
+    ):
+        super().__init__()
+
+        if act_layer == nn.Identity:
+            gain_factor = 1.0
+        else:
+            gain_factor = 2.0
+
+        if inner_skip == "linear" or inner_skip == "identity":
+            gain_factor /= 2.0
+
+        # convolution layer
+        if conv_type == "local":
+            theta_cutoff = 2.0 * _compute_cutoff_radius(forward_transform.nlat, disco_kernel_shape, disco_basis_type)
+            self.local_conv = DiscreteContinuousConvS2(
+                input_dim,
+                output_dim,
+                in_shape=(forward_transform.nlat, forward_transform.nlon),
+                out_shape=(inverse_transform.nlat, inverse_transform.nlon),
+                kernel_shape=disco_kernel_shape,
+                basis_type=disco_basis_type,
+                grid_in=forward_transform.grid,
+                grid_out=inverse_transform.grid,
+                bias=bias,
+                theta_cutoff=theta_cutoff,
+            )
+        elif conv_type == "global":
+            self.global_conv = SpectralConvS2(forward_transform, inverse_transform, input_dim, output_dim, gain=gain_factor, bias=bias)
+        else:
+            raise ValueError(f"Unknown convolution type {conv_type}")
+
+        if inner_skip == "linear":
+            self.inner_skip = nn.Conv2d(input_dim, output_dim, 1, 1)
+            nn.init.normal_(self.inner_skip.weight, std=math.sqrt(gain_factor / input_dim))
+        elif inner_skip == "identity":
+            assert input_dim == output_dim
+            self.inner_skip = nn.Identity()
+        elif inner_skip == "none":
+            pass
+        else:
+            raise ValueError(f"Unknown skip connection type {inner_skip}")
+
+        # normalisation layer
+        if norm_layer == "layer_norm":
+            self.norm = nn.LayerNorm(normalized_shape=(inverse_transform.nlat, inverse_transform.nlon), eps=1e-6)
+        elif norm_layer == "instance_norm":
+            self.norm = nn.InstanceNorm2d(num_features=output_dim, eps=1e-6, affine=True, track_running_stats=False)
+        elif norm_layer == "none":
+            self.norm = nn.Identity()
+        else:
+            raise NotImplementedError(f"Error, normalization {norm_layer} not implemented.")
+
+        # dropout
+        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+
+        gain_factor = 1.0
+        if outer_skip == "linear" or inner_skip == "identity":
+            gain_factor /= 2.0
+
+        if use_mlp == True:
+            mlp_hidden_dim = int(output_dim * mlp_ratio)
+            self.mlp = MLP(
+                in_features=output_dim,
+                out_features=input_dim,
+                hidden_features=mlp_hidden_dim,
+                act_layer=act_layer,
+                drop_rate=drop_rate,
+                checkpointing=False,
+                gain=gain_factor,
+            )
+
+        if outer_skip == "linear":
+            self.outer_skip = nn.Conv2d(input_dim, input_dim, 1, 1)
+            torch.nn.init.normal_(self.outer_skip.weight, std=math.sqrt(gain_factor / input_dim))
+        elif outer_skip == "identity":
+            assert input_dim == output_dim
+            self.outer_skip = nn.Identity()
+        elif outer_skip == "none":
+            pass
+        else:
+            raise ValueError(f"Unknown skip connection type {outer_skip}")
+
+    def forward(self, x):
+
+        residual = x
+
+        if hasattr(self, "global_conv"):
+            x, _ = self.global_conv(x)
+        elif hasattr(self, "local_conv"):
+            x = self.local_conv(x)
+
+        x = self.norm(x)
+
+        if hasattr(self, "inner_skip"):
+            x = x + self.inner_skip(residual)
+
+        if hasattr(self, "mlp"):
+            x = self.mlp(x)
+
+        x = self.drop_path(x)
+
+        if hasattr(self, "outer_skip"):
+            x = x + self.outer_skip(residual)
+
+        return x
+
+
+class LocalSphericalNeuralOperator(nn.Module):
+    """
+    LocalSphericalNeuralOperator module. A spherical neural operator which uses both local and global integral
+    operators to accureately model both types of solution operators [1]. The architecture is based on the Spherical
+    Fourier Neural Operator [2] and improves upon it with local integral operators in both the Neural Operator blocks,
+    as well as in the encoder and decoders.
+    
+    Parameters
+    ----------
+    img_size : tuple, optional
+        Input image size (nlat, nlon), by default (128, 256)
+    grid : str, optional
+        Grid type for input/output, by default "equiangular"
+    grid_internal : str, optional
+        Grid type for internal processing, by default "legendre-gauss"
+    scale_factor : int, optional
+        Scale factor for resolution changes, by default 3
+    in_chans : int, optional
+        Number of input channels, by default 3
+    out_chans : int, optional
+        Number of output channels, by default 3
+    embed_dim : int, optional
+        Embedding dimension, by default 256
+    num_layers : int, optional
+        Number of layers, by default 4
+    activation_function : str, optional
+        Activation function name, by default "gelu"
+    kernel_shape : tuple, optional
+        Kernel shape for convolutions, by default (3, 3)
+    encoder_kernel_shape : tuple, optional
+        Kernel shape for encoder, by default (3, 3)
+    filter_basis_type : str, optional
+        Filter basis type, by default "morlet"
+    use_mlp : bool, optional
+        Whether to use MLP layers, by default True
+    mlp_ratio : float, optional
+        MLP expansion ratio, by default 2.0
+    drop_rate : float, optional
+        Dropout rate, by default 0.0
+    drop_path_rate : float, optional
+        Drop path rate, by default 0.0
+    normalization_layer : str, optional
+        Type of normalization layer to use ("layer_norm", "instance_norm", "none"), by default "instance_norm"
+    sfno_block_frequency : int, optional
+        Frequency of SFNO blocks, by default 2
+    hard_thresholding_fraction : float, optional
+        Hard thresholding fraction, by default 1.0
+    residual_prediction : bool, optional
+        Whether to use residual prediction, by default False
+    pos_embed : str, optional
+        Position embedding type, by default "none"
+    upsample_sht : bool, optional
+        Use SHT upsampling if true, else linear interpolation
+    bias : bool, optional
+        Whether to use a bias, by default False
+
+    Example
+    ----------
+    >>> model = LocalSphericalNeuralOperator(
+    ...         img_shape=(128, 256),
+    ...         scale_factor=4,
+    ...         in_chans=2,
+    ...         out_chans=2,
+    ...         embed_dim=16,
+    ...         num_layers=4,
+    ...         use_mlp=True,)
+    >>> model(torch.randn(1, 2, 128, 256)).shape
+    torch.Size([1, 2, 128, 256])
+
+    References
+    ----------
+    .. [1] Liu-Schiaffini M., Berner J., Bonev B., Kurth T., Azizzadenesheli K., Anandkumar A.;
+        "Neural Operators with Localized Integral and Differential Kernels" (2024).
+        ICML 2024, https://arxiv.org/pdf/2402.16845.
+
+    .. [2] Bonev B., Kurth T., Hundt C., Pathak, J., Baust M., Kashinath K., Anandkumar A.;
+        "Spherical Fourier Neural Operators: Learning Stable Dynamics on the Sphere" (2023).
+        ICML 2023, https://arxiv.org/abs/2306.03838.
+
+    """
+
+    def __init__(
+        self,
+        img_size=(128, 256),
+        grid="equiangular",
+        grid_internal="legendre-gauss",
+        scale_factor=3,
+        in_chans=3,
+        out_chans=3,
+        embed_dim=256,
+        num_layers=4,
+        activation_function="gelu",
+        kernel_shape=(3, 3),
+        encoder_kernel_shape=(3, 3),
+        filter_basis_type="morlet",
+        use_mlp=True,
+        mlp_ratio=2.0,
+        drop_rate=0.0,
+        drop_path_rate=0.0,
+        normalization_layer="none",
+        sfno_block_frequency=2,
+        hard_thresholding_fraction=1.0,
+        residual_prediction=False,
+        pos_embed="none",
+        upsample_sht=False,
+        bias=False,
+    ):
+        super().__init__()
+
+        self.img_size = img_size
+        self.grid = grid
+        self.grid_internal = grid_internal
+        self.scale_factor = scale_factor
+        self.in_chans = in_chans
+        self.out_chans = out_chans
+        self.embed_dim = embed_dim
+        self.num_layers = num_layers
+        self.encoder_kernel_shape = encoder_kernel_shape
+        self.hard_thresholding_fraction = hard_thresholding_fraction
+        self.normalization_layer = normalization_layer
+        self.use_mlp = use_mlp
+        self.residual_prediction = residual_prediction
+
+        # activation function
+        if activation_function == "relu":
+            self.activation_function = nn.ReLU
+        elif activation_function == "gelu":
+            self.activation_function = nn.GELU
+        # for debugging purposes
+        elif activation_function == "identity":
+            self.activation_function = nn.Identity
+        else:
+            raise ValueError(f"Unknown activation function {activation_function}")
+
+        # compute downsampled image size. We assume that the latitude-grid includes both poles
+        self.h = (self.img_size[0] - 1) // scale_factor + 1
+        self.w = self.img_size[1] // scale_factor
+
+        # dropout
+        self.pos_drop = nn.Dropout(p=drop_rate) if drop_rate > 0.0 else nn.Identity()
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, self.num_layers)]
+
+        if pos_embed == "sequence":
+            self.pos_embed = SequencePositionEmbedding((self.h, self.w), num_chans=self.embed_dim, grid=grid_internal)
+        elif pos_embed == "spectral":
+            self.pos_embed = SpectralPositionEmbedding((self.h, self.w), num_chans=self.embed_dim, grid=grid_internal)
+        elif pos_embed == "learnable lat":
+            self.pos_embed = LearnablePositionEmbedding((self.h, self.w), num_chans=self.embed_dim, grid=grid_internal, embed_type="lat")
+        elif pos_embed == "learnable latlon":
+            self.pos_embed = LearnablePositionEmbedding((self.h, self.w), num_chans=self.embed_dim, grid=grid_internal, embed_type="latlon")
+        elif pos_embed == "none":
+            self.pos_embed = nn.Identity()
+        else:
+            raise ValueError(f"Unknown position embedding type {pos_embed}")
+
+        # encoder
+        self.encoder = DiscreteContinuousEncoder(
+            in_shape=self.img_size,
+            out_shape=(self.h, self.w),
+            grid_in=grid,
+            grid_out=grid_internal,
+            inp_chans=self.in_chans,
+            out_chans=self.embed_dim,
+            kernel_shape=self.encoder_kernel_shape,
+            basis_type=filter_basis_type,
+            groups=1,
+            bias=False,
+        )
+
+        # compute the modes for the sht
+        modes_lat = self.h
+        # due to some spectral artifacts with cufft, we substract one mode here
+        modes_lon = (self.w // 2 + 1) - 1
+
+        modes_lat = modes_lon = int(min(modes_lat, modes_lon) * self.hard_thresholding_fraction)
+
+        self.trans = RealSHT(self.h, self.w, lmax=modes_lat, mmax=modes_lon, grid=grid_internal).float()
+        self.itrans = InverseRealSHT(self.h, self.w, lmax=modes_lat, mmax=modes_lon, grid=grid_internal).float()
+
+        self.blocks = nn.ModuleList([])
+        for i in range(self.num_layers):
+
+            block = SphericalNeuralOperatorBlock(
+                self.trans,
+                self.itrans,
+                self.embed_dim,
+                self.embed_dim,
+                conv_type="global" if i % sfno_block_frequency == (sfno_block_frequency-1) else "local",
+                mlp_ratio=mlp_ratio,
+                drop_rate=drop_rate,
+                drop_path=dpr[i],
+                act_layer=self.activation_function,
+                norm_layer=self.normalization_layer,
+                use_mlp=use_mlp,
+                disco_kernel_shape=kernel_shape,
+                disco_basis_type=filter_basis_type,
+                bias=bias,
+            )
+
+            self.blocks.append(block)
+
+        # decoder
+        self.decoder = DiscreteContinuousDecoder(
+            in_shape=(self.h, self.w),
+            out_shape=self.img_size,
+            grid_in=grid_internal,
+            grid_out=grid,
+            inp_chans=self.embed_dim,
+            out_chans=self.out_chans,
+            kernel_shape=self.encoder_kernel_shape,
+            basis_type=filter_basis_type,
+            groups=1,
+            bias=False,
+            upsample_sht=upsample_sht,
+        )
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        return {"pos_embed", "cls_token"}
+
+    def forward_features(self, x):
+        x = self.pos_drop(x)
+
+        for blk in self.blocks:
+            x = blk(x)
+
+        return x
+
+    def forward(self, x):
+
+        if self.residual_prediction:
+            residual = x
+
+        x = self.encoder(x)
+
+        if self.pos_embed is not None:
+            x = self.pos_embed(x)
+
+        x = self.forward_features(x)
+
+        x = self.decoder(x)
+
+        if self.residual_prediction:
+            x = x + residual
+
+        return x

.deps/torch_harmonics/examples/models/s2segformer.py

@@ -0,0 +1,810 @@
+# coding=utf-8
+
+# SPDX-FileCopyrightText: Copyright (c) 2025 The torch-harmonics Authors. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+#
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+#
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+#
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.amp as amp
+
+from torch_harmonics import DiscreteContinuousConvS2, DiscreteContinuousConvTransposeS2
+from torch_harmonics import AttentionS2, NeighborhoodAttentionS2
+from torch_harmonics import ResampleS2
+from torch_harmonics import RealSHT, InverseRealSHT
+from torch_harmonics.quadrature import _precompute_latitudes
+
+from torch_harmonics.examples.models._layers import MLP, LayerNorm, DropPath
+
+from functools import partial
+
+
+# heuristic for finding theta_cutoff
+def _compute_cutoff_radius(nlat, kernel_shape, basis_type):
+    theta_cutoff_factor = {"piecewise linear": 0.5, "morlet": 0.5, "zernike": math.sqrt(2.0)}
+
+    return (kernel_shape[0] + 1) * theta_cutoff_factor[basis_type] * math.pi / float(nlat - 1)
+
+
+class OverlapPatchMerging(nn.Module):
+    """
+    Overlap patch merging module for spherical segformer.
+    
+    This module performs patch merging with overlapping patches using discrete-continuous
+    convolutions on the sphere, followed by layer normalization.
+    
+    Parameters
+    -----------
+    in_shape : tuple, optional
+        Input shape (nlat, nlon), by default (721, 1440)
+    out_shape : tuple, optional
+        Output shape (nlat, nlon), by default (481, 960)
+    grid_in : str, optional
+        Input grid type, by default "equiangular"
+    grid_out : str, optional
+        Output grid type, by default "equiangular"
+    in_channels : int, optional
+        Number of input channels, by default 3
+    out_channels : int, optional
+        Number of output channels, by default 64
+    kernel_shape : tuple, optional
+        Kernel shape for convolution, by default (3, 3)
+    basis_type : str, optional
+        Filter basis type, by default "morlet"
+    bias : bool, optional
+        Whether to use bias, by default False
+    """
+    
+    def __init__(
+        self,
+        in_shape=(721, 1440),
+        out_shape=(481, 960),
+        grid_in="equiangular",
+        grid_out="equiangular",
+        in_channels=3,
+        out_channels=64,
+        kernel_shape=(3, 3),
+        basis_type="morlet",
+        bias=False,
+    ):
+        super().__init__()
+
+        # convolution for patches, curtoff radius inferred from kernel shape
+        theta_cutoff = _compute_cutoff_radius(out_shape[0], kernel_shape, basis_type)
+        self.conv = DiscreteContinuousConvS2(
+            in_channels,
+            out_channels,
+            in_shape=in_shape,
+            out_shape=out_shape,
+            kernel_shape=kernel_shape,
+            basis_type=basis_type,
+            grid_in=grid_in,
+            grid_out=grid_out,
+            bias=bias,
+            theta_cutoff=theta_cutoff,
+        )
+
+        # layer norm
+        self.norm = nn.LayerNorm((out_channels), eps=1e-05, elementwise_affine=True, bias=True)
+
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+
+        if isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    def forward(self, x):
+
+        dtype = x.dtype
+
+        with amp.autocast(device_type="cuda", enabled=False):
+            x = x.float()
+            x = self.conv(x).to(dtype=dtype)
+
+        # permute
+        x = x.permute(0, 2, 3, 1)
+        x = self.norm(x)
+        out = x.permute(0, 3, 1, 2)
+
+        return out
+
+
+class MixFFN(nn.Module):
+    """
+    Mix FFN module for spherical segformer.
+    
+    This module implements a feed-forward network that combines MLP operations
+    with discrete-continuous convolutions on the sphere.
+    
+    Parameters
+    -----------
+    shape : tuple
+        Shape (nlat, nlon) of the input
+    inout_channels : int
+        Number of input/output channels
+    hidden_channels : int
+        Number of hidden channels in MLP
+    mlp_bias : bool, optional
+        Whether to use bias in MLP, by default True
+    grid : str, optional
+        Grid type, by default "equiangular"
+    kernel_shape : tuple, optional
+        Kernel shape for convolution, by default (3, 3)
+    basis_type : str, optional
+        Filter basis type, by default "morlet"
+    conv_bias : bool, optional
+        Whether to use bias in convolution, by default False
+    activation : nn.Module, optional
+        Activation function, by default nn.GELU
+    use_mlp : bool, optional
+        Whether to use MLP instead of linear layers, by default False
+    drop_path : float, optional
+        Drop path rate, by default 0.0
+    """
+    
+    def __init__(
+        self,
+        shape,
+        inout_channels,
+        hidden_channels,
+        mlp_bias=True,
+        grid="equiangular",
+        kernel_shape=(3, 3),
+        basis_type="morlet",
+        conv_bias=False,
+        activation=nn.GELU,
+        use_mlp=False,
+        drop_path=0.0,
+    ):
+        super().__init__()
+
+        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+
+        self.norm = nn.LayerNorm((inout_channels), eps=1e-05, elementwise_affine=True, bias=True)
+
+        if use_mlp:
+            # although the paper says MLP, it uses a single linear layer
+            self.mlp_in = MLP(inout_channels, hidden_features=hidden_channels, out_features=inout_channels, act_layer=activation, output_bias=False, drop_rate=0.0)
+        else:
+            self.mlp_in = nn.Conv2d(in_channels=inout_channels, out_channels=inout_channels, kernel_size=1, bias=True)
+
+        # convolution for patches, curtoff radius inferred from kernel shape
+        theta_cutoff = _compute_cutoff_radius(shape[0], kernel_shape, basis_type)
+        self.conv = DiscreteContinuousConvS2(
+            inout_channels,
+            inout_channels,
+            in_shape=shape,
+            out_shape=shape,
+            kernel_shape=kernel_shape,
+            basis_type=basis_type,
+            grid_in=grid,
+            grid_out=grid,
+            groups=inout_channels,
+            bias=conv_bias,
+            theta_cutoff=theta_cutoff,
+        )
+
+        if use_mlp:
+            self.mlp_out = MLP(inout_channels, hidden_features=hidden_channels, out_features=inout_channels, act_layer=activation, output_bias=False, drop_rate=0.0)
+        else:
+            self.mlp_out = nn.Conv2d(in_channels=inout_channels, out_channels=inout_channels, kernel_size=1, bias=True)
+
+        self.act = activation()
+
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+
+        if isinstance(m, nn.Conv2d):
+            nn.init.trunc_normal_(m.weight, std=0.02)
+            if m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        residual = x
+
+        # norm
+        x = x.permute(0, 2, 3, 1)
+        x = self.norm(x)
+        x = x.permute(0, 3, 1, 2)
+
+        # NOTE: we add another activation here
+        # because in the paper they only use depthwise conv,
+        # but without this activation it would just be a fused MM
+        # with the disco conv
+        x = self.mlp_in(x)
+
+        # conv parth
+        x = self.act(self.conv(x))
+
+        # second linear
+        x = self.mlp_out(x)
+
+        return residual + self.drop_path(x)
+
+
+class AttentionWrapper(nn.Module):
+    """
+    Attention wrapper for spherical segformer.
+    
+    This module wraps attention mechanisms (neighborhood or global) with optional
+    normalization and drop path regularization.
+    
+    Parameters
+    -----------
+    channels : int
+        Number of channels
+    shape : tuple
+        Shape (nlat, nlon) of the input
+    grid : str
+        Grid type
+    heads : int
+        Number of attention heads
+    pre_norm : bool, optional
+        Whether to apply normalization before attention, by default False
+    attention_drop_rate : float, optional
+        Dropout rate for attention, by default 0.0
+    drop_path : float, optional
+        Drop path rate, by default 0.0
+    attention_mode : str, optional
+        Attention mode ("neighborhood" or "global"), by default "neighborhood"
+    theta_cutoff : float, optional
+        Cutoff radius for neighborhood attention, by default None
+    bias : bool, optional
+        Whether to use bias, by default True
+    """
+    def __init__(
+        self,
+        channels,
+        shape,
+        grid,
+        heads,
+        pre_norm=False,
+        attention_drop_rate=0.0,
+        drop_path=0.0,
+        attention_mode="neighborhood",
+        theta_cutoff=None,
+        bias=True
+    ):
+        super().__init__()
+
+        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+        self.attention_mode = attention_mode
+
+        if attention_mode == "neighborhood":
+            if theta_cutoff is None:
+                theta_cutoff = (7.0 / math.sqrt(math.pi)) * math.pi / (shape[0] - 1)
+            self.att = NeighborhoodAttentionS2(
+                in_channels=channels,
+                in_shape=shape,
+                out_shape=shape,
+                grid_in=grid,
+                grid_out=grid,
+                theta_cutoff=theta_cutoff,
+                out_channels=channels,
+                num_heads=heads,
+                bias=bias
+                # drop_rate=attention_drop_rate,
+            )
+        else:
+            self.att = AttentionS2(
+                in_channels=channels,
+                num_heads=heads,
+                in_shape=shape,
+                out_shape=shape,
+                grid_in=grid,
+                grid_out=grid,
+                out_channels=channels,
+                drop_rate=attention_drop_rate,
+                bias=bias
+            )
+
+        self.norm = None
+        if pre_norm:
+            self.norm = nn.LayerNorm((channels), eps=1e-05, elementwise_affine=True, bias=True)
+
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+
+        residual = x
+        if self.norm is not None:
+            x = x.permute(0, 2, 3, 1)
+            x = self.norm(x)
+            x = x.permute(0, 3, 1, 2)
+
+        if self.attention_mode == "neighborhood":
+            dtype = x.dtype
+            with amp.autocast(device_type="cuda", enabled=False):
+                x = x.float()
+                x = self.att(x).to(dtype=dtype)
+        else:
+            x = self.att(x)
+
+        return residual + self.drop_path(x)
+
+
+class TransformerBlock(nn.Module):
+    """
+    Transformer block for spherical segformer.
+    
+    This block combines patch merging, attention, and Mix FFN operations
+    in a hierarchical structure for processing spherical data.
+    
+    Parameters
+    -----------
+    in_shape : tuple
+        Input shape (nlat, nlon)
+    out_shape : tuple
+        Output shape (nlat, nlon)
+    in_channels : int
+        Number of input channels
+    out_channels : int
+        Number of output channels
+    mlp_hidden_channels : int
+        Number of hidden channels in MLP
+    grid_in : str, optional
+        Input grid type, by default "equiangular"
+    grid_out : str, optional
+        Output grid type, by default "equiangular"
+    nrep : int, optional
+        Number of repetitions, by default 1
+    heads : int, optional
+        Number of attention heads, by default 1
+    kernel_shape : tuple, optional
+        Kernel shape for convolution, by default (3, 3)
+    basis_type : str, optional
+        Filter basis type, by default "morlet"
+    activation : nn.Module, optional
+        Activation function, by default nn.GELU
+    att_drop_rate : float, optional
+        Dropout rate for attention, by default 0.0
+    drop_path_rates : float, optional
+        Drop path rates, by default 0.0
+    attention_mode : str, optional
+        Attention mode ("neighborhood" or "global"), by default "neighborhood"
+    theta_cutoff : float, optional
+        Cutoff radius for neighborhood attention, by default None
+    bias : bool, optional
+        Whether to use bias, by default True
+    """
+    def __init__(
+        self,
+        in_shape,
+        out_shape,
+        in_channels,
+        out_channels,
+        mlp_hidden_channels,
+        grid_in="equiangular",
+        grid_out="equiangular",
+        nrep=1,
+        heads=1,
+        kernel_shape=(3, 3),
+        basis_type="morlet",
+        activation=nn.GELU,
+        att_drop_rate=0.0,
+        drop_path_rates=0.0,
+        attention_mode="neighborhood",
+        theta_cutoff=None,
+        bias=True
+    ):
+        super().__init__()
+
+        self.in_shape = in_shape
+        self.out_shape = out_shape
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+
+        if isinstance(drop_path_rates, float):
+            drop_path_rates = [x.item() for x in torch.linspace(0, drop_path_rates, nrep)]
+
+        assert len(drop_path_rates) == nrep
+
+        self.fwd = [
+            OverlapPatchMerging(
+                in_shape=in_shape,
+                out_shape=out_shape,
+                grid_in=grid_in,
+                grid_out=grid_out,
+                in_channels=in_channels,
+                out_channels=out_channels,
+                kernel_shape=kernel_shape,
+                basis_type=basis_type,
+                bias=False,
+            )
+        ]
+
+        for i in range(nrep):
+            self.fwd.append(
+                AttentionWrapper(
+                    channels=out_channels,
+                    shape=out_shape,
+                    grid=grid_out,
+                    heads=heads,
+                    pre_norm=True,
+                    attention_drop_rate=att_drop_rate,
+                    drop_path=drop_path_rates[i],
+                    attention_mode=attention_mode,
+                    theta_cutoff=theta_cutoff,
+                    bias=bias
+                )
+            )
+
+            self.fwd.append(
+                MixFFN(
+                    out_shape,
+                    inout_channels=out_channels,
+                    hidden_channels=mlp_hidden_channels,
+                    mlp_bias=True,
+                    grid=grid_out,
+                    kernel_shape=kernel_shape,
+                    basis_type=basis_type,
+                    conv_bias=False,
+                    activation=activation,
+                    use_mlp=False,
+                    drop_path=drop_path_rates[i],
+                )
+            )
+
+        # make sequential
+        self.fwd = nn.Sequential(*self.fwd)
+
+        # final norm
+        self.norm = nn.LayerNorm((out_channels), eps=1e-05, elementwise_affine=True, bias=True)
+
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.fwd(x)
+
+        # apply norm
+        x = x.permute(0, 2, 3, 1)
+        x = self.norm(x)
+        x = x.permute(0, 3, 1, 2)
+
+        return x
+
+
+class Upsampling(nn.Module):
+    """
+    Upsampling module for spherical segformer.
+    
+    This module performs upsampling using either discrete-continuous transposed convolutions
+    or bilinear resampling on spherical data.
+    
+    Parameters
+    -----------
+    in_shape : tuple
+        Input shape (nlat, nlon)
+    out_shape : tuple
+        Output shape (nlat, nlon)
+    in_channels : int
+        Number of input channels
+    out_channels : int
+        Number of output channels
+    hidden_channels : int
+        Number of hidden channels in MLP
+    mlp_bias : bool, optional
+        Whether to use bias in MLP, by default True
+    grid_in : str, optional
+        Input grid type, by default "equiangular"
+    grid_out : str, optional
+        Output grid type, by default "equiangular"
+    kernel_shape : tuple, optional
+        Kernel shape for convolution, by default (3, 3)
+    basis_type : str, optional
+        Filter basis type, by default "morlet"
+    conv_bias : bool, optional
+        Whether to use bias in convolution, by default False
+    activation : nn.Module, optional
+        Activation function, by default nn.GELU
+    use_mlp : bool, optional
+        Whether to use MLP instead of linear layers, by default False
+    upsampling_method : str, optional
+        Upsampling method ("conv" or "bilinear"), by default "conv"
+    """
+    def __init__(
+        self,
+        in_shape,
+        out_shape,
+        in_channels,
+        out_channels,
+        hidden_channels,
+        mlp_bias=True,
+        grid_in="equiangular",
+        grid_out="equiangular",
+        kernel_shape=(3, 3),
+        basis_type="morlet",
+        conv_bias=False,
+        activation=nn.GELU,
+        use_mlp=False,
+        upsampling_method="conv"
+    ):
+        super().__init__()
+
+        if use_mlp:
+            self.mlp = MLP(in_channels, hidden_features=hidden_channels, out_features=out_channels, act_layer=activation, output_bias=False, drop_rate=0.0)
+        else:
+            self.mlp = nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=1, bias=True)
+
+        if upsampling_method == "conv":
+            theta_cutoff = _compute_cutoff_radius(in_shape[0], kernel_shape, basis_type)
+            self.upsample = DiscreteContinuousConvTransposeS2(
+                out_channels,
+                out_channels,
+                in_shape=in_shape,
+                out_shape=out_shape,
+                kernel_shape=kernel_shape,
+                basis_type=basis_type,
+                grid_in=grid_in,
+                grid_out=grid_out,
+                bias=conv_bias,
+                theta_cutoff=theta_cutoff,
+            )
+        elif upsampling_method == "bilinear":
+            self.upsample = ResampleS2(*in_shape, *out_shape, grid_in=grid_in, grid_out=grid_out)
+        else:
+            raise ValueError(f"Unknown upsampling method {upsampling_method}")
+
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Conv2d):
+            nn.init.trunc_normal_(m.weight, std=0.02)
+            if m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+
+        x = self.upsample(self.mlp(x))
+
+        return x
+
+
+class SphericalSegformer(nn.Module):
+    """
+    Spherical segformer model designed to approximate mappings from spherical signals to spherical segmentation masks
+
+    Parameters
+    -----------
+    img_shape : tuple, optional
+        Shape of the input channels, by default (128, 256)
+    kernel_shape: tuple, int
+    scale_factor: int, optional
+        Scale factor to use, by default 2
+    in_chans : int, optional
+        Number of input channels, by default 3
+    out_chans : int, optional
+        Number of classes, by default 3
+    embed_dims : List[int], optional
+        Dimension of the embeddings for each block, has to be the same length as heads
+    heads : List[int], optional
+        Number of heads for each block in the network, has to be the same length as embed_dims
+    depths: List[in], optional
+        Number of repetitions of attentions blocks and ffn mixers per layer. Has to be the same length as embed_dims and heads
+    activation_function : str, optional
+        Activation function to use, by default "gelu"
+    embedder_kernel_shape : int, optional
+        size of the encoder kernel
+    filter_basis_type: Optional[str]: str, optional
+        filter basis type
+    use_mlp : int, optional
+        Whether to use MLPs in the SFNO blocks, by default True
+    mlp_ratio : int, optional
+        Ratio of MLP to use, by default 2.0
+    drop_rate : float, optional
+        Dropout rate, by default 0.0
+    drop_path_rate : float, optional
+        Dropout path rate, by default 0.0
+    normalization_layer : str, optional
+        Type of normalization layer to use ("layer_norm", "instance_norm", "none"), by default "instance_norm"
+    hard_thresholding_fraction : float, optional
+        Fraction of hard thresholding (frequency cutoff) to apply, by default 1.0
+    upsampling_method : str
+        Conv, bilinear
+
+    Example
+    -----------
+    >>> model = SphericalTransformer(
+    ...         img_shape=(128, 256),
+    ...         scale_factor=4,
+    ...         in_chans=2,
+    ...         out_chans=2,
+    ...         embed_dim=16,
+    ...         num_layers=4,
+    ...         use_mlp=True,)
+    >>> model(torch.randn(1, 2, 128, 256)).shape
+    torch.Size([1, 2, 128, 256])
+    """
+
+    def __init__(
+        self,
+        img_size=(128, 256),
+        grid="equiangular",
+        grid_internal="legendre-gauss",
+        in_chans=3,
+        out_chans=3,
+        embed_dims=[64, 128, 256, 512],
+        heads=[1, 2, 4, 8],
+        depths=[3, 4, 6, 3],
+        scale_factor=2,
+        activation_function="gelu",
+        kernel_shape=(3, 3),
+        filter_basis_type="morlet",
+        mlp_ratio=2.0,
+        att_drop_rate=0.0,
+        drop_path_rate=0.1,
+        attention_mode="neighborhood",
+        theta_cutoff=None,
+        upsampling_method="bilinear",
+        bias=True
+    ):
+        super().__init__()
+
+        self.img_size = img_size
+        self.grid = grid
+        self.grid_internal = grid_internal
+        self.in_chans = in_chans
+        self.out_chans = out_chans
+        self.embed_dims = embed_dims
+        self.heads = heads
+        self.num_blocks = len(self.embed_dims)
+        self.depths = depths
+        self.kernel_shape = kernel_shape
+
+        assert len(self.heads) == self.num_blocks
+        assert len(self.depths) == self.num_blocks
+
+        # activation function
+        if activation_function == "relu":
+            self.activation_function = nn.ReLU
+        elif activation_function == "gelu":
+            self.activation_function = nn.GELU
+        # for debugging purposes
+        elif activation_function == "identity":
+            self.activation_function = nn.Identity
+        else:
+            raise ValueError(f"Unknown activation function {activation_function}")
+
+        # set up drop path rates
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(self.depths))]
+
+        self.blocks = nn.ModuleList([])
+        out_shape = img_size
+        grid_in = grid
+        grid_out = grid_internal
+        in_channels = in_chans
+        cur = 0
+        for i in range(self.num_blocks):
+            out_shape_new = (out_shape[0] // scale_factor, out_shape[1] // scale_factor)
+            out_channels = self.embed_dims[i]
+            self.blocks.append(
+                TransformerBlock(
+                    in_shape=out_shape,
+                    out_shape=out_shape_new,
+                    in_channels=in_channels,
+                    out_channels=out_channels,
+                    mlp_hidden_channels=int(mlp_ratio * out_channels),
+                    grid_in=grid_in,
+                    grid_out=grid_out,
+                    nrep=self.depths[i],
+                    heads=self.heads[i],
+                    kernel_shape=kernel_shape,
+                    basis_type=filter_basis_type,
+                    activation=self.activation_function,
+                    att_drop_rate=att_drop_rate,
+                    drop_path_rates=dpr[cur : cur + self.depths[i]],
+                    attention_mode=attention_mode,
+                    theta_cutoff=theta_cutoff,
+                    bias=bias
+                )
+            )
+            cur += self.depths[i]
+            out_shape = out_shape_new
+            grid_in = grid_internal
+            in_channels = out_channels
+
+        self.upsamplers = nn.ModuleList([])
+        out_shape = img_size
+        grid_out = grid
+        for i in range(self.num_blocks):
+            in_shape = self.blocks[i].out_shape
+            self.upsamplers.append(
+                Upsampling(
+                    in_shape=in_shape,
+                    out_shape=out_shape,
+                    in_channels=self.embed_dims[i],
+                    out_channels=self.embed_dims[i],
+                    hidden_channels=int(mlp_ratio * self.embed_dims[i]),
+                    mlp_bias=True,
+                    grid_in=grid_internal,
+                    grid_out=grid,
+                    kernel_shape=kernel_shape,
+                    basis_type=filter_basis_type,
+                    conv_bias=False,
+                    activation=nn.GELU,
+                    upsampling_method=upsampling_method
+                )
+            )
+
+        segmentation_head_dim = sum(self.embed_dims)
+        self.segmentation_head = nn.Conv2d(in_channels=segmentation_head_dim, out_channels=out_chans, kernel_size=1, bias=True)
+
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+
+        if isinstance(m, nn.Conv2d):
+            nn.init.trunc_normal_(m.weight, std=0.02)
+            if m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    def forward(self, x):
+
+        # encoder:
+        features = []
+        feat = x
+        for block in self.blocks:
+            feat = block(feat)
+            features.append(feat)
+
+        # perform upsample
+        upfeats = []
+        for feat, upsampler in zip(features, self.upsamplers):
+            upfeats.append(upsampler(feat))
+
+        # perform concatenation
+        upfeats = torch.cat(upfeats, dim=1)
+
+        # final upsampling and prediction
+        out = self.segmentation_head(upfeats)
+
+        return out

.deps/torch_harmonics/examples/models/s2transformer.py

@@ -0,0 +1,589 @@
+# coding=utf-8
+
+# SPDX-FileCopyrightText: Copyright (c) 2025 The torch-harmonics Authors. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+#
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+#
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+#
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.amp as amp
+
+from torch_harmonics import DiscreteContinuousConvS2, DiscreteContinuousConvTransposeS2
+from torch_harmonics import NeighborhoodAttentionS2, AttentionS2
+from torch_harmonics import ResampleS2
+from torch_harmonics import RealSHT, InverseRealSHT
+from torch_harmonics.quadrature import _precompute_latitudes
+
+from torch_harmonics.examples.models._layers import MLP, DropPath, LayerNorm, SequencePositionEmbedding, SpectralPositionEmbedding, LearnablePositionEmbedding
+
+from functools import partial
+
+# heuristic for finding theta_cutoff
+def _compute_cutoff_radius(nlat, kernel_shape, basis_type):
+    theta_cutoff_factor = {"piecewise linear": 0.5, "morlet": 0.5, "zernike": math.sqrt(2.0)}
+
+    return (kernel_shape[0] + 1) * theta_cutoff_factor[basis_type] * math.pi / float(nlat - 1)
+
+class DiscreteContinuousEncoder(nn.Module):
+    """
+    Discrete-continuous encoder for spherical transformers.
+    
+    This module performs downsampling using discrete-continuous convolutions on the sphere,
+    reducing the spatial resolution while maintaining the spectral properties of the data.
+    
+    Parameters
+    -----------
+    in_shape : tuple, optional
+        Input shape (nlat, nlon), by default (721, 1440)
+    out_shape : tuple, optional
+        Output shape (nlat, nlon), by default (480, 960)
+    grid_in : str, optional
+        Input grid type, by default "equiangular"
+    grid_out : str, optional
+        Output grid type, by default "equiangular"
+    in_chans : int, optional
+        Number of input channels, by default 2
+    out_chans : int, optional
+        Number of output channels, by default 2
+    kernel_shape : tuple, optional
+        Kernel shape for convolution, by default (3, 3)
+    basis_type : str, optional
+        Filter basis type, by default "morlet"
+    groups : int, optional
+        Number of groups for grouped convolution, by default 1
+    bias : bool, optional
+        Whether to use bias, by default False
+    """
+    
+    def __init__(
+        self,
+        in_shape=(721, 1440),
+        out_shape=(480, 960),
+        grid_in="equiangular",
+        grid_out="equiangular",
+        in_chans=2,
+        out_chans=2,
+        kernel_shape=(3, 3),
+        basis_type="morlet",
+        groups=1,
+        bias=False,
+    ):
+        super().__init__()
+
+        # set up local convolution
+        self.conv = DiscreteContinuousConvS2(
+            in_chans,
+            out_chans,
+            in_shape=in_shape,
+            out_shape=out_shape,
+            kernel_shape=kernel_shape,
+            basis_type=basis_type,
+            grid_in=grid_in,
+            grid_out=grid_out,
+            groups=groups,
+            bias=bias,
+            theta_cutoff=_compute_cutoff_radius(in_shape[0], kernel_shape, basis_type),
+        )
+
+    def forward(self, x):
+
+        dtype = x.dtype
+
+        with amp.autocast(device_type="cuda", enabled=False):
+            x = x.float()
+            x = self.conv(x)
+            x = x.to(dtype=dtype)
+
+        return x
+
+
+class DiscreteContinuousDecoder(nn.Module):
+    """
+    Discrete-continuous decoder for spherical transformers.
+    
+    This module performs upsampling using either spherical harmonic transforms or resampling,
+    followed by discrete-continuous convolutions to restore spatial resolution.
+    
+    Parameters
+    -----------
+    in_shape : tuple, optional
+        Input shape (nlat, nlon), by default (480, 960)
+    out_shape : tuple, optional
+        Output shape (nlat, nlon), by default (721, 1440)
+    grid_in : str, optional
+        Input grid type, by default "equiangular"
+    grid_out : str, optional
+        Output grid type, by default "equiangular"
+    in_chans : int, optional
+        Number of input channels, by default 2
+    out_chans : int, optional
+        Number of output channels, by default 2
+    kernel_shape : tuple, optional
+        Kernel shape for convolution, by default (3, 3)
+    basis_type : str, optional
+        Filter basis type, by default "morlet"
+    groups : int, optional
+        Number of groups for grouped convolution, by default 1
+    bias : bool, optional
+        Whether to use bias, by default False
+    upsample_sht : bool, optional
+        Whether to use SHT for upsampling, by default False
+    """
+    
+    def __init__(
+        self,
+        in_shape=(480, 960),
+        out_shape=(721, 1440),
+        grid_in="equiangular",
+        grid_out="equiangular",
+        in_chans=2,
+        out_chans=2,
+        kernel_shape=(3, 3),
+        basis_type="morlet",
+        groups=1,
+        bias=False,
+        upsample_sht=False,
+    ):
+        super().__init__()
+
+        # set up upsampling
+        if upsample_sht:
+            self.sht = RealSHT(*in_shape, grid=grid_in).float()
+            self.isht = InverseRealSHT(*out_shape, lmax=self.sht.lmax, mmax=self.sht.mmax, grid=grid_out).float()
+            self.upsample = nn.Sequential(self.sht, self.isht)
+        else:
+            self.upsample = ResampleS2(*in_shape, *out_shape, grid_in=grid_in, grid_out=grid_out)
+
+        # set up DISCO convolution
+        self.conv = DiscreteContinuousConvS2(
+            in_chans,
+            out_chans,
+            in_shape=out_shape,
+            out_shape=out_shape,
+            kernel_shape=kernel_shape,
+            basis_type=basis_type,
+            grid_in=grid_out,
+            grid_out=grid_out,
+            groups=groups,
+            bias=False,
+            theta_cutoff=_compute_cutoff_radius(in_shape[0], kernel_shape, basis_type),
+        )
+
+    def forward(self, x):
+
+        dtype = x.dtype
+
+        with amp.autocast(device_type="cuda", enabled=False):
+            x = x.float()
+            x = self.upsample(x)
+            x = self.conv(x)
+            x = x.to(dtype=dtype)
+
+        return x
+
+
+class SphericalAttentionBlock(nn.Module):
+    """
+    Spherical attention block for transformers on the sphere.
+    
+    This module implements a single attention block that can use either global attention
+    or neighborhood attention on spherical data, followed by an optional MLP.
+    
+    Parameters
+    -----------
+    in_shape : tuple, optional
+        Input shape (nlat, nlon), by default (480, 960)
+    out_shape : tuple, optional
+        Output shape (nlat, nlon), by default (480, 960)
+    grid_in : str, optional
+        Input grid type, by default "equiangular"
+    grid_out : str, optional
+        Output grid type, by default "equiangular"
+    in_chans : int, optional
+        Number of input channels, by default 2
+    out_chans : int, optional
+        Number of output channels, by default 2
+    num_heads : int, optional
+        Number of attention heads, by default 1
+    mlp_ratio : float, optional
+        Ratio of MLP hidden dimension to output dimension, by default 2.0
+    drop_rate : float, optional
+        Dropout rate, by default 0.0
+    drop_path : float, optional
+        Drop path rate, by default 0.0
+    act_layer : nn.Module, optional
+        Activation layer, by default nn.GELU
+    norm_layer : str, optional
+        Normalization layer type, by default "none"
+    use_mlp : bool, optional
+        Whether to use MLP after attention, by default True
+    bias : bool, optional
+        Whether to use bias, by default False
+    attention_mode : str, optional
+        Attention mode ("neighborhood" or "global"), by default "neighborhood"
+    theta_cutoff : float, optional
+        Cutoff radius for neighborhood attention, by default None
+    """
+
+    def __init__(
+        self,
+        in_shape=(480, 960),
+        out_shape=(480, 960),
+        grid_in="equiangular",
+        grid_out="equiangular",
+        in_chans=2,
+        out_chans=2,
+        num_heads=1,
+        mlp_ratio=2.0,
+        drop_rate=0.0,
+        drop_path=0.0,
+        act_layer=nn.GELU,
+        norm_layer="none",
+        use_mlp=True,
+        bias=False,
+        attention_mode="neighborhood",
+        theta_cutoff=None,
+    ):
+        super().__init__()
+
+        # normalisation layer
+        if norm_layer == "layer_norm":
+            self.norm0 = LayerNorm(in_channels=in_chans, eps=1e-6)
+            self.norm1 = LayerNorm(in_channels=out_chans, eps=1e-6)
+        elif norm_layer == "instance_norm":
+            self.norm0 = nn.InstanceNorm2d(num_features=in_chans, eps=1e-6, affine=True, track_running_stats=False)
+            self.norm1 = nn.InstanceNorm2d(num_features=out_chans, eps=1e-6, affine=True, track_running_stats=False)
+        elif norm_layer == "none":
+            self.norm0 = nn.Identity()
+            self.norm1 = nn.Identity()
+        else:
+            raise NotImplementedError(f"Error, normalization {norm_layer} not implemented.")
+
+        # determine radius for neighborhood attention
+        self.attention_mode = attention_mode
+        if attention_mode == "neighborhood":
+            if theta_cutoff is None:
+                theta_cutoff = (7.0 / math.sqrt(math.pi)) * math.pi / (in_shape[0] - 1)
+            self.self_attn = NeighborhoodAttentionS2(
+                in_channels=in_chans,
+                in_shape=in_shape,
+                out_shape=out_shape,
+                grid_in=grid_in,
+                grid_out=grid_out,
+                num_heads=num_heads,
+                theta_cutoff=theta_cutoff,
+                k_channels=None,
+                out_channels=out_chans,
+                bias=bias,
+            )
+        else:
+            self.self_attn = AttentionS2(
+                in_channels=in_chans,
+                num_heads=num_heads,
+                in_shape=in_shape,
+                out_shape=out_shape,
+                grid_in=grid_in,
+                grid_out=grid_out,
+                out_channels=out_chans,
+                drop_rate=drop_rate,
+                bias=bias,
+            )
+
+        self.skip0 = nn.Identity()
+
+        # dropout
+        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+
+        if use_mlp == True:
+            mlp_hidden_dim = int(out_chans * mlp_ratio)
+            self.mlp = MLP(
+                in_features=out_chans,
+                out_features=out_chans,
+                hidden_features=mlp_hidden_dim,
+                act_layer=act_layer,
+                drop_rate=drop_rate,
+                checkpointing=False,
+                gain=0.5,
+            )
+
+        self.skip1 = nn.Identity()
+
+    def forward(self, x):
+
+        residual = x
+
+        x = self.norm0(x)
+
+        if self.attention_mode == "neighborhood":
+            dtype = x.dtype
+            with amp.autocast(device_type="cuda", enabled=False):
+                x = x.float()
+                x = self.self_attn(x).to(dtype=dtype)
+        else:
+            x = self.self_attn(x)
+
+        if hasattr(self, "skip0"):
+            x = x + self.skip0(residual)
+
+        residual = x
+
+        x = self.norm1(x)
+
+        if hasattr(self, "mlp"):
+            x = self.mlp(x)
+
+        x = self.drop_path(x)
+
+        if hasattr(self, "skip1"):
+            x = x + self.skip1(residual)
+
+        return x
+
+
+class SphericalTransformer(nn.Module):
+    """
+    Spherical transformer model designed to approximate mappings from spherical signals to spherical signals
+
+    Parameters
+    -----------
+    img_shape : tuple, optional
+        Shape of the input channels, by default (128, 256)
+    kernel_shape: tuple, int
+    scale_factor : int, optional
+        Scale factor to use, by default 3
+    in_chans : int, optional
+        Number of input channels, by default 3
+    out_chans : int, optional
+        Number of output channels, by default 3
+    embed_dim : int, optional
+        Dimension of the embeddings, by default 256
+    num_layers : int, optional
+        Number of layers in the network, by default 4
+    activation_function : str, optional
+        Activation function to use, by default "gelu"
+    encoder_kernel_shape : int, optional
+        size of the encoder kernel
+    filter_basis_type: str, optional
+        filter basis type
+    num_heads: int, optional
+        number of attention heads
+    use_mlp : int, optional
+        Whether to use MLPs in the SFNO blocks, by default True
+    mlp_ratio : int, optional
+        Ratio of MLP to use, by default 2.0
+    drop_rate : float, optional
+        Dropout rate, by default 0.0
+    drop_path_rate : float, optional
+        Dropout path rate, by default 0.0
+    normalization_layer : str, optional
+        Type of normalization layer to use ("layer_norm", "instance_norm", "none"), by default "instance_norm"
+    hard_thresholding_fraction : float, optional
+        Fraction of hard thresholding (frequency cutoff) to apply, by default 1.0
+    residual_prediction : bool, optional
+        Whether to add a single large skip connection, by default True
+    pos_embed : bool, optional
+        Whether to use positional embedding, by default True
+    upsample_sht : bool, optional
+        Use SHT upsampling if true, else linear interpolation
+    bias : bool, optional
+        Whether to use a bias, by default False
+
+    Example
+    -----------
+    >>> model = SphericalTransformer(
+    ...         img_shape=(128, 256),
+    ...         scale_factor=4,
+    ...         in_chans=2,
+    ...         out_chans=2,
+    ...         embed_dim=16,
+    ...         num_layers=4,
+    ...         use_mlp=True,)
+    >>> model(torch.randn(1, 2, 128, 256)).shape
+    torch.Size([1, 2, 128, 256])
+    """
+
+    def __init__(
+        self,
+        img_size=(128, 256),
+        grid="equiangular",
+        grid_internal="legendre-gauss",
+        scale_factor=3,
+        in_chans=3,
+        out_chans=3,
+        embed_dim=256,
+        num_layers=4,
+        activation_function="gelu",
+        encoder_kernel_shape=(3, 3),
+        filter_basis_type="morlet",
+        num_heads=1,
+        use_mlp=True,
+        mlp_ratio=2.0,
+        drop_rate=0.0,
+        drop_path_rate=0.0,
+        normalization_layer="none",
+        hard_thresholding_fraction=1.0,
+        residual_prediction=False,
+        pos_embed="spectral",
+        upsample_sht=False,
+        attention_mode="neighborhood",
+        bias=False,
+        theta_cutoff=None,
+    ):
+        super().__init__()
+        self.img_size = img_size
+        self.grid = grid
+        self.grid_internal = grid_internal
+        self.scale_factor = scale_factor
+        self.in_chans = in_chans
+        self.out_chans = out_chans
+        self.embed_dim = embed_dim
+        self.num_layers = num_layers
+        self.encoder_kernel_shape = encoder_kernel_shape
+        self.hard_thresholding_fraction = hard_thresholding_fraction
+        self.normalization_layer = normalization_layer
+        self.use_mlp = use_mlp
+        self.residual_prediction = residual_prediction
+
+        # activation function
+        if activation_function == "relu":
+            self.activation_function = nn.ReLU
+        elif activation_function == "gelu":
+            self.activation_function = nn.GELU
+        # for debugging purposes
+        elif activation_function == "identity":
+            self.activation_function = nn.Identity
+        else:
+            raise ValueError(f"Unknown activation function {activation_function}")
+
+        # compute downsampled image size. We assume that the latitude-grid includes both poles
+        self.h = (self.img_size[0] - 1) // scale_factor + 1
+        self.w = self.img_size[1] // scale_factor
+
+        # dropout
+        self.pos_drop = nn.Dropout(p=drop_rate) if drop_rate > 0.0 else nn.Identity()
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, self.num_layers)]
+
+        if pos_embed == "sequence":
+            self.pos_embed = SequencePositionEmbedding((self.h, self.w), num_chans=self.embed_dim, grid=grid_internal)
+        elif pos_embed == "spectral":
+            self.pos_embed = SpectralPositionEmbedding((self.h, self.w), num_chans=self.embed_dim, grid=grid_internal)
+        elif pos_embed == "learnable lat":
+            self.pos_embed = LearnablePositionEmbedding((self.h, self.w), num_chans=self.embed_dim, grid=grid_internal, embed_type="lat")
+        elif pos_embed == "learnable latlon":
+            self.pos_embed = LearnablePositionEmbedding((self.h, self.w), num_chans=self.embed_dim, grid=grid_internal, embed_type="latlon")
+        elif pos_embed == "none":
+            self.pos_embed = nn.Identity()
+        else:
+            raise ValueError(f"Unknown position embedding type {pos_embed}")
+
+        # maybe keep for now becuase tr
+        # encoder
+        self.encoder = DiscreteContinuousEncoder(
+            in_shape=self.img_size,
+            out_shape=(self.h, self.w),
+            grid_in=grid,
+            grid_out=grid_internal,
+            in_chans=self.in_chans,
+            out_chans=self.embed_dim,
+            kernel_shape=self.encoder_kernel_shape,
+            basis_type=filter_basis_type,
+            groups=1,
+            bias=False,
+        )
+
+        self.blocks = nn.ModuleList([])
+        for i in range(self.num_layers):
+            block = SphericalAttentionBlock(
+                in_shape=(self.h, self.w),
+                out_shape=(self.h, self.w),
+                grid_in=grid_internal,
+                grid_out=grid_internal,
+                in_chans=self.embed_dim,
+                out_chans=self.embed_dim,
+                num_heads=num_heads,
+                mlp_ratio=mlp_ratio,
+                drop_rate=drop_rate,
+                drop_path=dpr[i],
+                act_layer=self.activation_function,
+                norm_layer=self.normalization_layer,
+                use_mlp=use_mlp,
+                bias=bias,
+                attention_mode=attention_mode,
+                theta_cutoff=theta_cutoff,
+            )
+
+            self.blocks.append(block)
+
+        # decoder
+        self.decoder = DiscreteContinuousDecoder(
+            in_shape=(self.h, self.w),
+            out_shape=self.img_size,
+            grid_in=grid_internal,
+            grid_out=grid,
+            in_chans=self.embed_dim,
+            out_chans=self.out_chans,
+            kernel_shape=self.encoder_kernel_shape,
+            basis_type=filter_basis_type,
+            groups=1,
+            bias=False,
+            upsample_sht=upsample_sht,
+        )
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        return {"pos_embed", "cls_token"}
+
+    def forward_features(self, x):
+        x = self.pos_drop(x)
+
+        for blk in self.blocks:
+            x = blk(x)
+
+        return x
+
+    def forward(self, x):
+
+        if self.residual_prediction:
+            residual = x
+
+        x = self.encoder(x)
+
+        if self.pos_embed is not None:
+            # x = x + self.pos_embed
+            x = self.pos_embed(x)
+
+        x = self.forward_features(x)
+
+        x = self.decoder(x)
+
+        if self.residual_prediction:
+            x = x + residual
+
+        return x

.deps/torch_harmonics/examples/models/s2unet.py

@@ -0,0 +1,610 @@
+# coding=utf-8
+
+# SPDX-FileCopyrightText: Copyright (c) 2025 The torch-harmonics Authors. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+#
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+#
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+#
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.amp as amp
+
+from torch_harmonics import DiscreteContinuousConvS2, DiscreteContinuousConvTransposeS2
+from torch_harmonics import NeighborhoodAttentionS2
+from torch_harmonics import ResampleS2
+from torch_harmonics import RealSHT, InverseRealSHT
+from torch_harmonics.quadrature import _precompute_latitudes
+
+from torch_harmonics.examples.models._layers import MLP, DropPath
+
+from functools import partial
+
+
+# heuristic for finding theta_cutoff
+def _compute_cutoff_radius(nlat, kernel_shape, basis_type):
+    theta_cutoff_factor = {"piecewise linear": 0.5, "morlet": 0.5, "zernike": math.sqrt(2.0)}
+
+    return (kernel_shape[0] + 1) * theta_cutoff_factor[basis_type] * math.pi / float(nlat - 1)
+
+
+class DownsamplingBlock(nn.Module):
+    """
+    Downsampling block for spherical U-Net architecture.
+    
+    This block performs convolution operations followed by downsampling on spherical data,
+    using discrete-continuous convolutions to maintain spectral properties.
+    
+    Parameters
+    -----------
+    in_shape : tuple
+        Input shape (nlat, nlon)
+    out_shape : tuple
+        Output shape (nlat, nlon)
+    in_channels : int
+        Number of input channels
+    out_channels : int
+        Number of output channels
+    grid_in : str, optional
+        Input grid type, by default "equiangular"
+    grid_out : str, optional
+        Output grid type, by default "equiangular"
+    nrep : int, optional
+        Number of convolution repetitions, by default 1
+    kernel_shape : tuple, optional
+        Kernel shape for convolution, by default (3, 3)
+    basis_type : str, optional
+        Filter basis type, by default "morlet"
+    activation : nn.Module, optional
+        Activation function, by default nn.ReLU
+    transform_skip : bool, optional
+        Whether to transform skip connection, by default False
+    drop_conv_rate : float, optional
+        Dropout rate for convolutions, by default 0.0
+    drop_path_rate : float, optional
+        Drop path rate, by default 0.0
+    drop_dense_rate : float, optional
+        Dropout rate for dense layers, by default 0.0
+    downsampling_mode : str, optional
+        Downsampling mode ("bilinear", "conv"), by default "bilinear"
+    """
+    
+    def __init__(
+        self,
+        in_shape,
+        out_shape,
+        in_channels,
+        out_channels,
+        grid_in="equiangular",
+        grid_out="equiangular",
+        nrep=1,
+        kernel_shape=(3, 3),
+        basis_type="morlet",
+        activation=nn.ReLU,
+        transform_skip=False,
+        drop_conv_rate=0.0,
+        drop_path_rate=0.0,
+        drop_dense_rate=0.0,
+        downsampling_mode="bilinear",
+    ):
+        super().__init__()
+
+        self.in_shape = in_shape
+        self.out_shape = out_shape
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.grid_in = grid_in
+        self.grid_out = grid_out
+
+        self.drop_path = DropPath(drop_path_rate) if drop_path_rate > 0.0 else nn.Identity()
+
+        self.fwd = []
+        for i in range(nrep):
+            # conv
+            theta_cutoff = _compute_cutoff_radius(in_shape[0], kernel_shape, basis_type)
+            self.fwd.append(
+                DiscreteContinuousConvS2(
+                    in_channels=(in_channels if i == 0 else out_channels),
+                    out_channels=out_channels,
+                    in_shape=in_shape,
+                    out_shape=in_shape,
+                    kernel_shape=kernel_shape,
+                    basis_type=basis_type,
+                    grid_in=grid_out,
+                    grid_out=grid_out,
+                    bias=False,
+                    theta_cutoff=theta_cutoff,
+                )
+            )
+
+            if drop_conv_rate > 0.0:
+                self.fwd.append(nn.Dropout2d(p=drop_conv_rate))
+
+            # batchnorm
+            self.fwd.append(nn.BatchNorm2d(out_channels, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True))
+
+            # activation
+            self.fwd.append(
+                activation(),
+            )
+
+        if downsampling_mode == "conv":
+            theta_cutoff = _compute_cutoff_radius(out_shape[0], kernel_shape, basis_type)
+            self.downsample = DiscreteContinuousConvS2(
+                out_channels,
+                out_channels,
+                in_shape=in_shape,
+                out_shape=out_shape,
+                kernel_shape=kernel_shape,
+                basis_type=basis_type,
+                grid_in=grid_in,
+                grid_out=grid_out,
+                bias=False,
+                theta_cutoff=theta_cutoff,
+            )
+        else:
+            self.downsample = ResampleS2(
+                nlat_in=in_shape[0],
+                nlon_in=in_shape[1],
+                nlat_out=out_shape[0],
+                nlon_out=out_shape[1],
+                grid_in=grid_in,
+                grid_out=grid_out,
+                mode=downsampling_mode,
+            )
+
+        # make sequential
+        self.fwd = nn.Sequential(*self.fwd)
+
+        # final norm
+        if transform_skip or (in_channels != out_channels):
+            self.transform_skip = nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=True)
+
+            if drop_dense_rate > 0.0:
+                self.transform_skip = nn.Sequential(
+                    self.transform_skip,
+                    nn.Dropout2d(p=drop_dense_rate),
+                )
+
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+
+        if isinstance(m, nn.Conv2d):
+            nn.init.trunc_normal_(m.weight, std=0.02)
+            if m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+
+        # skip connection
+        residual = x
+        if hasattr(self, "transform_skip"):
+            residual = self.transform_skip(residual)
+
+        # main path
+        x = self.fwd(x)
+
+        # add residual connection
+        x = residual + self.drop_path(x)
+
+        # downsample
+        x = self.downsample(x)
+
+        return x
+
+
+class UpsamplingBlock(nn.Module):
+    """
+    Upsampling block for spherical U-Net architecture.
+    
+    This block performs upsampling followed by convolution operations on spherical data,
+    using discrete-continuous convolutions to maintain spectral properties.
+    
+    Parameters
+    -----------
+    in_shape : tuple
+        Input shape (nlat, nlon)
+    out_shape : tuple
+        Output shape (nlat, nlon)
+    in_channels : int
+        Number of input channels
+    out_channels : int
+        Number of output channels
+    grid_in : str, optional
+        Input grid type, by default "equiangular"
+    grid_out : str, optional
+        Output grid type, by default "equiangular"
+    nrep : int, optional
+        Number of convolution repetitions, by default 1
+    kernel_shape : tuple, optional
+        Kernel shape for convolution, by default (3, 3)
+    basis_type : str, optional
+        Filter basis type, by default "morlet"
+    activation : nn.Module, optional
+        Activation function, by default nn.ReLU
+    transform_skip : bool, optional
+        Whether to transform skip connection, by default False
+    drop_conv_rate : float, optional
+        Dropout rate for convolutions, by default 0.0
+    drop_path_rate : float, optional
+        Drop path rate, by default 0.0
+    drop_dense_rate : float, optional
+        Dropout rate for dense layers, by default 0.0
+    upsampling_mode : str, optional
+        Upsampling mode ("bilinear", "conv"), by default "bilinear"
+    """
+
+    def __init__(
+        self,
+        in_shape,
+        out_shape,
+        in_channels,
+        out_channels,
+        grid_in="equiangular",
+        grid_out="equiangular",
+        nrep=1,
+        kernel_shape=(3, 3),
+        basis_type="morlet",
+        activation=nn.ReLU,
+        transform_skip=False,
+        drop_conv_rate=0.0,
+        drop_path_rate=0.0,
+        drop_dense_rate=0.0,
+        upsampling_mode="bilinear",
+    ):
+        super().__init__()
+
+        self.in_shape = in_shape
+        self.out_shape = out_shape
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+
+        self.drop_path = DropPath(drop_path_rate) if drop_path_rate > 0.0 else nn.Identity()
+
+        if in_shape != out_shape:
+            if upsampling_mode == "conv":
+                theta_cutoff = _compute_cutoff_radius(in_shape[0], kernel_shape, basis_type)
+                self.upsample = nn.Sequential(
+                    DiscreteContinuousConvTransposeS2(
+                        in_channels=out_channels,
+                        out_channels=out_channels,
+                        in_shape=in_shape,
+                        out_shape=out_shape,
+                        kernel_shape=kernel_shape,
+                        basis_type=basis_type,
+                        grid_in=grid_in,
+                        grid_out=grid_out,
+                        bias=False,
+                        theta_cutoff=theta_cutoff,
+                    ),
+                    nn.BatchNorm2d(out_channels, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True),
+                    activation(),
+                    DiscreteContinuousConvS2(
+                        in_channels=out_channels,
+                        out_channels=out_channels,
+                        in_shape=out_shape,
+                        out_shape=out_shape,
+                        kernel_shape=kernel_shape,
+                        basis_type=basis_type,
+                        grid_in=grid_in,
+                        grid_out=grid_out,
+                        bias=False,
+                        theta_cutoff=theta_cutoff,
+                    ),
+                )
+
+            else:
+                self.upsample = ResampleS2(
+                    nlat_in=in_shape[0],
+                    nlon_in=in_shape[1],
+                    nlat_out=out_shape[0],
+                    nlon_out=out_shape[1],
+                    grid_in=grid_in,
+                    grid_out=grid_out,
+                    mode=upsampling_mode,
+                )
+        else:
+            theta_cutoff = _compute_cutoff_radius(in_shape[0], kernel_shape, basis_type)
+            self.upsample = DiscreteContinuousConvS2(
+                in_channels=out_channels,
+                out_channels=out_channels,
+                in_shape=in_shape,
+                out_shape=in_shape,
+                kernel_shape=kernel_shape,
+                basis_type=basis_type,
+                grid_in=grid_in,
+                grid_out=grid_out,
+                bias=False,
+                theta_cutoff=theta_cutoff,
+            )
+
+        self.fwd = []
+        for i in range(nrep):
+            # conv
+            theta_cutoff = _compute_cutoff_radius(in_shape[0], kernel_shape, basis_type)
+            self.fwd.append(
+                DiscreteContinuousConvS2(
+                    in_channels=in_channels,
+                    out_channels=(out_channels if i == nrep - 1 else in_channels),
+                    in_shape=in_shape,
+                    out_shape=in_shape,
+                    kernel_shape=kernel_shape,
+                    basis_type=basis_type,
+                    grid_in=grid_in,
+                    grid_out=grid_in,
+                    bias=False,
+                    theta_cutoff=theta_cutoff,
+                )
+            )
+
+            if drop_conv_rate > 0.0:
+                self.fwd.append(nn.Dropout2d(p=drop_conv_rate))
+
+            # batchnorm
+            self.fwd.append(nn.BatchNorm2d((out_channels if i == nrep - 1 else in_channels), eps=1e-05, momentum=0.1, affine=True, track_running_stats=True))
+
+            # activation
+            self.fwd.append(
+                activation(),
+            )
+
+        # make sequential
+        self.fwd = nn.Sequential(*self.fwd)
+
+        # final norm
+        if transform_skip or (in_channels != out_channels):
+            self.transform_skip = nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=True)
+            if drop_dense_rate > 0.0:
+                self.transform_skip = nn.Sequential(
+                    self.transform_skip,
+                    nn.Dropout2d(p=drop_dense_rate),
+                )
+
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Conv2d):
+            nn.init.trunc_normal_(m.weight, std=0.02)
+            if m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # skip connection
+        residual = x
+        if hasattr(self, "transform_skip"):
+            residual = self.transform_skip(residual)
+
+        # main path
+        x = residual + self.drop_path(self.fwd(x))
+
+        # upsampling
+        x = self.upsample(x)
+
+        return x
+
+
+class SphericalUNet(nn.Module):
+    """
+    Spherical segformer model designed to approximate mappings from spherical signals to spherical segmentation masks
+
+    Parameters
+    -----------
+    img_shape : tuple, optional
+        Shape of the input channels, by default (128, 256)
+    kernel_shape: tuple, int
+    scale_factor: int, optional
+        Scale factor to use, by default 2
+    in_chans : int, optional
+        Number of input channels, by default 3
+    out_chans : int, optional
+        Number of classes, by default 3
+    embed_dims : List[int], optional
+        Dimension of the embeddings for each block, has to be the same length as depths
+    depths: List[in], optional
+        Number of repetitions of conv blocks and ffn mixers per layer. Has to be the same length as embed_dims
+    activation_function : str, optional
+        Activation function to use, by default "relu"
+    embedder_kernel_shape : int, optional
+        size of the encoder kernel
+    filter_basis_type: Optional[str]: str, optional
+        filter basis type
+    use_mlp : int, optional
+        Whether to use MLPs in the SFNO blocks, by default True
+    mlp_ratio : int, optional
+        Ratio of MLP to use, by default 2.0
+    drop_rate : float, optional
+        Dropout rate, by default 0.0
+    drop_path_rate : float, optional
+        Dropout path rate, by default 0.0
+    normalization_layer : str, optional
+        Type of normalization layer to use ("layer_norm", "instance_norm", "none"), by default "instance_norm"
+    hard_thresholding_fraction : float, optional
+        Fraction of hard thresholding (frequency cutoff) to apply, by default 1.0
+    upsample_sht : bool, optional
+        Use SHT upsampling if true, else linear interpolation
+
+    Example
+    -----------
+    >>> model = SphericalTransformer(
+    ...         img_shape=(128, 256),
+    ...         scale_factor=4,
+    ...         in_chans=2,
+    ...         out_chans=2,
+    ...         embed_dim=16,
+    ...         num_layers=4,
+    ...         use_mlp=True,)
+    >>> model(torch.randn(1, 2, 128, 256)).shape
+    torch.Size([1, 2, 128, 256])
+    """
+
+    def __init__(
+        self,
+        img_size=(128, 256),
+        grid="equiangular",
+        grid_internal="legendre-gauss",
+        in_chans=3,
+        out_chans=3,
+        embed_dims=[64, 128, 256, 512],
+        depths=[2, 2, 2, 2],
+        scale_factor=2,
+        activation_function="relu",
+        kernel_shape=(3, 3),
+        filter_basis_type="morlet",
+        transform_skip=False,
+        drop_conv_rate=0.1,
+        drop_path_rate=0.1,
+        drop_dense_rate=0.5,
+        downsampling_mode="bilinear",
+        upsampling_mode="bilinear",
+    ):
+        super().__init__()
+
+        self.img_size = img_size
+        self.grid = grid
+        self.grid_internal = grid_internal
+        self.in_chans = in_chans
+        self.out_chans = out_chans
+        self.embed_dims = embed_dims
+        self.num_blocks = len(self.embed_dims)
+        self.depths = depths
+        self.kernel_shape = kernel_shape
+
+        assert len(self.depths) == self.num_blocks
+
+        # activation function
+        if activation_function == "relu":
+            self.activation_function = nn.ReLU
+        elif activation_function == "gelu":
+            self.activation_function = nn.GELU
+        # for debugging purposes
+        elif activation_function == "identity":
+            self.activation_function = nn.Identity
+        else:
+            raise ValueError(f"Unknown activation function {activation_function}")
+
+        # set up drop path rates
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, self.num_blocks)]
+
+        self.dblocks = nn.ModuleList([])
+        out_shape = img_size
+        grid_in = grid
+        grid_out = grid_internal
+        in_channels = in_chans
+        for i in range(self.num_blocks):
+            out_shape_new = (out_shape[0] // scale_factor, out_shape[1] // scale_factor)
+            out_channels = self.embed_dims[i]
+            self.dblocks.append(
+                DownsamplingBlock(
+                    in_shape=out_shape,
+                    out_shape=out_shape_new,
+                    in_channels=in_channels,
+                    out_channels=out_channels,
+                    grid_in=grid_in,
+                    grid_out=grid_out,
+                    nrep=self.depths[i],
+                    kernel_shape=kernel_shape,
+                    basis_type=filter_basis_type,
+                    activation=self.activation_function,
+                    drop_conv_rate=drop_conv_rate,
+                    drop_path_rate=dpr[i],
+                    drop_dense_rate=drop_dense_rate,
+                    transform_skip=transform_skip,
+                    downsampling_mode=downsampling_mode,
+                )
+            )
+            out_shape = out_shape_new
+            grid_in = grid_internal
+            in_channels = out_channels
+
+        self.ublocks = nn.ModuleList([])
+        for i in range(self.num_blocks - 1, -1, -1):
+            in_shape = self.dblocks[i].out_shape
+            out_shape = self.dblocks[i].in_shape
+            in_channels = self.dblocks[i].out_channels
+            if i != self.num_blocks - 1:
+                in_channels = 2 * in_channels
+            out_channels = self.dblocks[i].in_channels
+            if i == 0:
+                out_channels = self.embed_dims[0]
+            grid_in = self.dblocks[i].grid_out
+            grid_out = self.dblocks[i].grid_in
+            self.ublocks.append(
+                UpsamplingBlock(
+                    in_shape=in_shape,
+                    out_shape=out_shape,
+                    in_channels=in_channels,
+                    out_channels=out_channels,
+                    grid_in=grid_in,
+                    grid_out=grid_out,
+                    kernel_shape=kernel_shape,
+                    basis_type=filter_basis_type,
+                    activation=self.activation_function,
+                    drop_conv_rate=drop_conv_rate,
+                    drop_path_rate=0.0,
+                    drop_dense_rate=drop_dense_rate,
+                    transform_skip=transform_skip,
+                    upsampling_mode=upsampling_mode,
+                )
+            )
+
+        self.head = nn.Conv2d(self.embed_dims[0], self.out_chans, kernel_size=1, bias=True)
+
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+
+        if isinstance(m, nn.Conv2d):
+            nn.init.trunc_normal_(m.weight, std=0.02)
+            if m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    def forward(self, x):
+
+        # encoder:
+        features = []
+        feat = x
+        for dblock in self.dblocks:
+            feat = dblock(feat)
+            features.append(feat)
+
+        # reverse list
+        features = features[::-1]
+
+        # perform upsample
+        ufeat = self.ublocks[0](features[0])
+        for feat, ublock in zip(features[1:], self.ublocks[1:]):
+            ufeat = ublock(torch.cat([feat, ufeat], dim=1))
+
+        # last layer
+        out = self.head(ufeat)
+
+        return out

.deps/torch_harmonics/examples/models/sfno.py

@@ -0,0 +1,423 @@
+# coding=utf-8
+
+# SPDX-FileCopyrightText: Copyright (c) 2022 The torch-harmonics Authors. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+#
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+#
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+#
+import math
+
+import torch
+import torch.nn as nn
+
+from torch_harmonics import RealSHT, InverseRealSHT
+
+from torch_harmonics.examples.models._layers import MLP, SpectralConvS2, SequencePositionEmbedding, SpectralPositionEmbedding, LearnablePositionEmbedding
+
+from functools import partial
+
+
+class SphericalFourierNeuralOperatorBlock(nn.Module):
+    """
+    Helper module for a single SFNO/FNO block. Can use both FFTs and SHTs to represent either FNO or SFNO blocks.
+
+    Parameters
+    ----------
+    forward_transform : torch.nn.Module
+        Forward transform to use for the block
+    inverse_transform : torch.nn.Module
+        Inverse transform to use for the block
+    input_dim : int
+        Input dimension
+    output_dim : int
+        Output dimension
+    mlp_ratio : float, optional
+        MLP expansion ratio, by default 2.0
+    drop_rate : float, optional
+        Dropout rate, by default 0.0
+    drop_path : float, optional
+        Drop path rate, by default 0.0
+    act_layer : torch.nn.Module, optional
+        Activation function to use, by default nn.GELU
+    norm_layer : str, optional
+        Type of normalization to use, by default "none"
+    inner_skip : str, optional
+        Type of inner skip connection to use, by default "none"
+    outer_skip : str, optional
+        Type of outer skip connection to use, by default "identity"
+    use_mlp : bool, optional
+        Whether to use MLP layers, by default True
+    bias : bool, optional
+        Whether to use bias, by default False
+
+    Returns
+    -------
+    torch.Tensor
+        Output tensor
+    """
+
+    def __init__(
+        self,
+        forward_transform,
+        inverse_transform,
+        input_dim,
+        output_dim,
+        mlp_ratio=2.0,
+        drop_rate=0.0,
+        drop_path=0.0,
+        act_layer=nn.GELU,
+        norm_layer="none",
+        inner_skip="none",
+        outer_skip="identity",
+        use_mlp=True,
+        bias=False,
+    ):
+        super().__init__()
+
+        if act_layer == nn.Identity:
+            gain_factor = 1.0
+        else:
+            gain_factor = 2.0
+
+        if inner_skip == "linear" or inner_skip == "identity":
+            gain_factor /= 2.0
+
+        self.global_conv = SpectralConvS2(forward_transform, inverse_transform, input_dim, output_dim, gain=gain_factor, bias=bias)
+
+        if inner_skip == "linear":
+            self.inner_skip = nn.Conv2d(input_dim, output_dim, 1, 1)
+            nn.init.normal_(self.inner_skip.weight, std=math.sqrt(gain_factor / input_dim))
+        elif inner_skip == "identity":
+            assert input_dim == output_dim
+            self.inner_skip = nn.Identity()
+        elif inner_skip == "none":
+            pass
+        else:
+            raise ValueError(f"Unknown skip connection type {inner_skip}")
+
+        # normalisation layer
+        if norm_layer == "layer_norm":
+            self.norm = nn.LayerNorm(normalized_shape=(inverse_transform.nlat, inverse_transform.nlon), eps=1e-6)
+        elif norm_layer == "instance_norm":
+            self.norm = nn.InstanceNorm2d(num_features=output_dim, eps=1e-6, affine=True, track_running_stats=False)
+        elif norm_layer == "none":
+            self.norm = nn.Identity()
+        else:
+            raise NotImplementedError(f"Error, normalization {self.norm_layer} not implemented.")
+
+        # dropout
+        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+
+        gain_factor = 1.0
+        if outer_skip == "linear" or inner_skip == "identity":
+            gain_factor /= 2.0
+
+        if use_mlp == True:
+            mlp_hidden_dim = int(output_dim * mlp_ratio)
+            self.mlp = MLP(
+                in_features=output_dim, out_features=input_dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop_rate=drop_rate, checkpointing=False, gain=gain_factor
+            )
+
+        if outer_skip == "linear":
+            self.outer_skip = nn.Conv2d(input_dim, input_dim, 1, 1)
+            torch.nn.init.normal_(self.outer_skip.weight, std=math.sqrt(gain_factor / input_dim))
+        elif outer_skip == "identity":
+            assert input_dim == output_dim
+            self.outer_skip = nn.Identity()
+        elif outer_skip == "none":
+            pass
+        else:
+            raise ValueError(f"Unknown skip connection type {outer_skip}")
+
+    def forward(self, x):
+
+        x, residual = self.global_conv(x)
+
+        x = self.norm(x)
+
+        if hasattr(self, "inner_skip"):
+            x = x + self.inner_skip(residual)
+
+        if hasattr(self, "mlp"):
+            x = self.mlp(x)
+
+        x = self.drop_path(x)
+
+        if hasattr(self, "outer_skip"):
+            x = x + self.outer_skip(residual)
+
+        return x
+
+
+class SphericalFourierNeuralOperator(nn.Module):
+    """
+    SphericalFourierNeuralOperator module. Implements the 'linear' variant of the Spherical Fourier Neural Operator
+    as presented in [1]. Spherical convolutions are applied via spectral transforms to apply a geometrically consistent
+    and approximately equivariant architecture.
+
+    Parameters
+    ----------
+    img_size : tuple, optional
+        Shape of the input channels, by default (128, 256)
+    grid : str, optional
+        Input grid type, by default "equiangular"
+    grid_internal : str, optional
+        Internal grid type for computations, by default "legendre-gauss"
+    scale_factor : int, optional
+        Scale factor to use, by default 3
+    in_chans : int, optional
+        Number of input channels, by default 3
+    out_chans : int, optional
+        Number of output channels, by default 3
+    embed_dim : int, optional
+        Dimension of the embeddings, by default 256
+    num_layers : int, optional
+        Number of layers in the network, by default 4
+    activation_function : str, optional
+        Activation function to use, by default "gelu"
+    encoder_layers : int, optional
+        Number of layers in the encoder, by default 1
+    use_mlp : int, optional
+        Whether to use MLPs in the SFNO blocks, by default True
+    mlp_ratio : int, optional
+        Ratio of MLP to use, by default 2.0
+    drop_rate : float, optional
+        Dropout rate, by default 0.0
+    drop_path_rate : float, optional
+        Dropout path rate, by default 0.0
+    normalization_layer : str, optional
+        Type of normalization layer to use ("layer_norm", "instance_norm", "none"), by default "none"
+    hard_thresholding_fraction : float, optional
+        Fraction of hard thresholding (frequency cutoff) to apply, by default 1.0
+    residual_prediction : bool, optional
+        Whether to add a single large skip connection, by default False
+    pos_embed : str, optional
+        Type of positional embedding to use, by default "none"
+    bias : bool, optional
+        Whether to use a bias, by default False
+
+    Example:
+    ----------
+    >>> model = SphericalFourierNeuralOperator(
+    ...         img_size=(128, 256),
+    ...         scale_factor=4,
+    ...         in_chans=2,
+    ...         out_chans=2,
+    ...         embed_dim=16,
+    ...         num_layers=4,
+    ...         use_mlp=True,)
+    >>> model(torch.randn(1, 2, 128, 256)).shape
+    torch.Size([1, 2, 128, 256])
+
+    References
+    ----------
+    .. [1] Bonev B., Kurth T., Hundt C., Pathak, J., Baust M., Kashinath K., Anandkumar A.;
+        "Spherical Fourier Neural Operators: Learning Stable Dynamics on the Sphere" (2023).
+        ICML 2023, https://arxiv.org/abs/2306.03838.
+    """
+
+    def __init__(
+        self,
+        img_size=(128, 256),
+        grid="equiangular",
+        grid_internal="legendre-gauss",
+        scale_factor=3,
+        in_chans=3,
+        out_chans=3,
+        embed_dim=256,
+        num_layers=4,
+        activation_function="gelu",
+        encoder_layers=1,
+        use_mlp=True,
+        mlp_ratio=2.0,
+        drop_rate=0.0,
+        drop_path_rate=0.0,
+        normalization_layer="none",
+        hard_thresholding_fraction=1.0,
+        residual_prediction=False,
+        pos_embed="none",
+        bias=False,
+    ):
+
+        super().__init__()
+
+        self.img_size = img_size
+        self.grid = grid
+        self.grid_internal = grid_internal
+        self.scale_factor = scale_factor
+        self.in_chans = in_chans
+        self.out_chans = out_chans
+        self.embed_dim = embed_dim
+        self.num_layers = num_layers
+        self.hard_thresholding_fraction = hard_thresholding_fraction
+        self.normalization_layer = normalization_layer
+        self.use_mlp = use_mlp
+        self.encoder_layers = encoder_layers
+        self.residual_prediction = residual_prediction
+
+        # activation function
+        if activation_function == "relu":
+            self.activation_function = nn.ReLU
+        elif activation_function == "gelu":
+            self.activation_function = nn.GELU
+        # for debugging purposes
+        elif activation_function == "identity":
+            self.activation_function = nn.Identity
+        else:
+            raise ValueError(f"Unknown activation function {activation_function}")
+
+        # compute downsampled image size. We assume that the latitude-grid includes both poles
+        self.h = (self.img_size[0] - 1) // scale_factor + 1
+        self.w = self.img_size[1] // scale_factor
+
+        # dropout
+        self.pos_drop = nn.Dropout(p=drop_rate) if drop_rate > 0.0 else nn.Identity()
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, self.num_layers)]
+
+        if pos_embed == "sequence":
+            self.pos_embed = SequencePositionEmbedding((self.h, self.w), num_chans=self.embed_dim, grid=grid_internal)
+        elif pos_embed == "spectral":
+            self.pos_embed = SpectralPositionEmbedding((self.h, self.w), num_chans=self.embed_dim, grid=grid_internal)
+        elif pos_embed == "learnable lat":
+            self.pos_embed = LearnablePositionEmbedding((self.h, self.w), num_chans=self.embed_dim, grid=grid_internal, embed_type="lat")
+        elif pos_embed == "learnable latlon":
+            self.pos_embed = LearnablePositionEmbedding((self.h, self.w), num_chans=self.embed_dim, grid=grid_internal, embed_type="latlon")
+        elif pos_embed == "none":
+            self.pos_embed = nn.Identity()
+        else:
+            raise ValueError(f"Unknown position embedding type {pos_embed}")
+
+        # construct an encoder with num_encoder_layers
+        num_encoder_layers = 1
+        encoder_hidden_dim = int(self.embed_dim * mlp_ratio)
+        current_dim = self.in_chans
+        encoder_layers = []
+        for l in range(num_encoder_layers - 1):
+            fc = nn.Conv2d(current_dim, encoder_hidden_dim, 1, bias=True)
+            # initialize the weights correctly
+            scale = math.sqrt(2.0 / current_dim)
+            nn.init.normal_(fc.weight, mean=0.0, std=scale)
+            if fc.bias is not None:
+                nn.init.constant_(fc.bias, 0.0)
+            encoder_layers.append(fc)
+            encoder_layers.append(self.activation_function())
+            current_dim = encoder_hidden_dim
+        fc = nn.Conv2d(current_dim, self.embed_dim, 1, bias=bias)
+        scale = math.sqrt(1.0 / current_dim)
+        nn.init.normal_(fc.weight, mean=0.0, std=scale)
+        if fc.bias is not None:
+            nn.init.constant_(fc.bias, 0.0)
+        encoder_layers.append(fc)
+        self.encoder = nn.Sequential(*encoder_layers)
+
+        # compute the modes for the sht
+        modes_lat = self.h
+        # due to some spectral artifacts with cufft, we substract one mode here
+        modes_lon = (self.w // 2 + 1) - 1
+
+        modes_lat = modes_lon = int(min(modes_lat, modes_lon) * self.hard_thresholding_fraction)
+
+        self.trans_down = RealSHT(*self.img_size, lmax=modes_lat, mmax=modes_lon, grid=self.grid).float()
+        self.itrans_up = InverseRealSHT(*self.img_size, lmax=modes_lat, mmax=modes_lon, grid=self.grid).float()
+        self.trans = RealSHT(self.h, self.w, lmax=modes_lat, mmax=modes_lon, grid=grid_internal).float()
+        self.itrans = InverseRealSHT(self.h, self.w, lmax=modes_lat, mmax=modes_lon, grid=grid_internal).float()
+
+        self.blocks = nn.ModuleList([])
+        for i in range(self.num_layers):
+
+            first_layer = i == 0
+            last_layer = i == self.num_layers - 1
+
+            block = SphericalFourierNeuralOperatorBlock(
+                self.trans_down if first_layer else self.trans,
+                self.itrans_up if last_layer else self.itrans,
+                self.embed_dim,
+                self.embed_dim,
+                mlp_ratio=mlp_ratio,
+                drop_rate=drop_rate,
+                drop_path=dpr[i],
+                act_layer=self.activation_function,
+                norm_layer=self.normalization_layer,
+                use_mlp=use_mlp,
+                bias=bias,
+            )
+
+            self.blocks.append(block)
+
+        # construct an decoder with num_decoder_layers
+        num_decoder_layers = 1
+        decoder_hidden_dim = int(self.embed_dim * mlp_ratio)
+        current_dim = self.embed_dim
+        decoder_layers = []
+        for l in range(num_decoder_layers - 1):
+            fc = nn.Conv2d(current_dim, decoder_hidden_dim, 1, bias=True)
+            # initialize the weights correctly
+            scale = math.sqrt(2.0 / current_dim)
+            nn.init.normal_(fc.weight, mean=0.0, std=scale)
+            if fc.bias is not None:
+                nn.init.constant_(fc.bias, 0.0)
+            decoder_layers.append(fc)
+            decoder_layers.append(self.activation_function())
+            current_dim = decoder_hidden_dim
+        fc = nn.Conv2d(current_dim, self.out_chans, 1, bias=bias)
+        scale = math.sqrt(1.0 / current_dim)
+        nn.init.normal_(fc.weight, mean=0.0, std=scale)
+        if fc.bias is not None:
+            nn.init.constant_(fc.bias, 0.0)
+        decoder_layers.append(fc)
+        self.decoder = nn.Sequential(*decoder_layers)
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        return {"pos_embed", "cls_token"}
+
+    def forward_features(self, x):
+
+        x = self.pos_drop(x)
+
+        for blk in self.blocks:
+            x = blk(x)
+
+        return x
+
+    def forward(self, x):
+
+        if self.residual_prediction:
+            residual = x
+
+        x = self.encoder(x)
+
+        if self.pos_embed is not None:
+            x = self.pos_embed(x)
+
+        x = self.forward_features(x)
+
+        x = self.decoder(x)
+
+        if self.residual_prediction:
+            x = x + residual
+
+        return x

.deps/torch_harmonics/examples/pde_dataset.py

@@ -0,0 +1,159 @@
+# coding=utf-8
+
+# SPDX-FileCopyrightText: Copyright (c) 2022 The torch-harmonics Authors. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+#
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+#
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+#
+
+import torch
+
+from math import ceil
+
+from .shallow_water_equations import ShallowWaterSolver
+
+
+class PdeDataset(torch.utils.data.Dataset):
+    """Custom Dataset class for PDE training data
+
+    Parameters
+    ----------
+    dt : float
+        Time step
+    nsteps : int
+        Number of solver steps
+    dims : tuple, optional
+        Number of latitude and longitude points, by default (384, 768)
+    grid : str, optional
+        Grid type, by default "equiangular"
+    pde : str, optional
+        PDE type, by default "shallow water equations"
+    initial_condition : str, optional
+        Initial condition type, by default "random"
+    num_examples : int, optional
+        Number of examples, by default 32
+    device : torch.device, optional
+        Device to use, by default torch.device("cpu")
+    normalize : bool, optional
+        Whether to normalize the input and target, by default True
+    stream : torch.cuda.Stream, optional
+        CUDA stream to use, by default None
+
+    Returns
+    -------
+    inp : torch.Tensor
+        Input tensor
+    tar : torch.Tensor
+        Target tensor
+    """
+
+    def __init__(
+        self,
+        dt,
+        nsteps,
+        dims=(384, 768),
+        grid="equiangular",
+        pde="shallow water equations",
+        initial_condition="random",
+        num_examples=32,
+        device=torch.device("cpu"),
+        normalize=True,
+        stream=None,
+    ):
+        self.num_examples = num_examples
+        self.device = device
+        self.stream = stream
+
+        self.nlat = dims[0]
+        self.nlon = dims[1]
+
+        # number of solver steps used to compute the target
+        self.nsteps = nsteps
+        self.normalize = normalize
+
+        if pde == "shallow water equations":
+            lmax = ceil(self.nlat / 3)
+            mmax = lmax
+            dt_solver = dt / float(self.nsteps)
+            self.solver = ShallowWaterSolver(self.nlat, self.nlon, dt_solver, lmax=lmax, mmax=mmax, grid=grid).to(self.device).float()
+        else:
+            raise NotImplementedError
+
+        self.set_initial_condition(ictype=initial_condition)
+
+        if self.normalize:
+            inp0, _ = self._get_sample()
+            self.inp_mean = torch.mean(inp0, dim=(-1, -2)).reshape(-1, 1, 1)
+            self.inp_var = torch.var(inp0, dim=(-1, -2)).reshape(-1, 1, 1)
+
+    def __len__(self):
+        length = self.num_examples if self.ictype == "random" else 1
+        return length
+
+    def set_initial_condition(self, ictype="random"):
+        self.ictype = ictype
+
+    def set_num_examples(self, num_examples=32):
+        self.num_examples = num_examples
+
+    def _get_sample(self):
+        if self.ictype == "random":
+            inp = self.solver.random_initial_condition(mach=0.2)
+        elif self.ictype == "galewsky":
+            inp = self.solver.galewsky_initial_condition()
+
+        # solve pde for n steps to return the target
+        tar = self.solver.timestep(inp, self.nsteps)
+        inp = self.solver.spec2grid(inp)
+        tar = self.solver.spec2grid(tar)
+
+        return inp, tar
+
+    def __getitem__(self, index):
+
+        # if self.stream is None:
+        #     self.stream = torch.cuda.Stream()
+
+        # with torch.cuda.stream(self.stream):
+        #     with torch.inference_mode():
+        #         with torch.no_grad():
+        #             inp, tar = self._get_sample()
+
+        #             if self.normalize:
+        #                 inp = (inp - self.inp_mean) / torch.sqrt(self.inp_var)
+        #                 tar = (tar - self.inp_mean) / torch.sqrt(self.inp_var)
+
+        # self.stream.synchronize()
+
+        with torch.inference_mode():
+            with torch.no_grad():
+                inp, tar = self._get_sample()
+
+                if self.normalize:
+                    inp = (inp - self.inp_mean) / torch.sqrt(self.inp_var)
+                    tar = (tar - self.inp_mean) / torch.sqrt(self.inp_var)
+
+        return inp.clone(), tar.clone()

.deps/torch_harmonics/examples/pde_sphere.py

@@ -0,0 +1,227 @@
+# coding=utf-8
+
+# SPDX-FileCopyrightText: Copyright (c) 2022 The torch-harmonics Authors. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+#
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+#
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+#
+
+
+import torch
+import torch.nn as nn
+import torch_harmonics as th
+from torch_harmonics.quadrature import _precompute_longitudes
+
+import math
+import numpy as np
+
+
+class SphereSolver(nn.Module):
+    """
+    Solver class on the sphere. Can solve the following PDEs:
+    - Allen-Cahn equation
+    - Ginzburg-Landau equation
+    
+    Parameters
+    -----------
+    nlat : int
+        Number of latitude points
+    nlon : int
+        Number of longitude points
+    dt : float
+        Time step size
+    lmax : int, optional
+        Maximum l mode for spherical harmonics, by default None
+    mmax : int, optional
+        Maximum m mode for spherical harmonics, by default None
+    grid : str, optional
+        Grid type ("equiangular", "legendre-gauss", "lobatto"), by default "equiangular"
+    radius : float, optional
+        Radius of the sphere, by default 1.0
+    coeff : float, optional
+        Coefficient for the PDE, by default 0.001
+    """
+
+    def __init__(self, nlat, nlon, dt, lmax=None, mmax=None, grid="equiangular", radius=1.0, coeff=0.001):
+        super().__init__()
+
+        # time stepping param
+        self.dt = dt
+
+        # grid parameters
+        self.nlat = nlat
+        self.nlon = nlon
+        self.grid = grid
+
+        # physical sonstants
+        self.register_buffer('radius', torch.as_tensor(radius, dtype=torch.float64))
+        self.register_buffer('coeff', torch.as_tensor(coeff, dtype=torch.float64))
+
+        # SHT
+        self.sht = th.RealSHT(nlat, nlon, lmax=lmax, mmax=mmax, grid=grid, csphase=False)
+        self.isht = th.InverseRealSHT(nlat, nlon, lmax=lmax, mmax=mmax, grid=grid, csphase=False)
+
+        self.lmax = lmax or self.sht.lmax
+        self.mmax = lmax or self.sht.mmax
+
+        # compute gridpoints
+        if self.grid == "legendre-gauss":
+            cost, _ = th.quadrature.legendre_gauss_weights(self.nlat, -1, 1)
+        elif self.grid == "lobatto":
+            cost, _ = th.quadrature.lobatto_weights(self.nlat, -1, 1)
+        elif self.grid == "equiangular":
+            cost, _ = th.quadrature.clenshaw_curtiss_weights(self.nlat, -1, 1)
+
+        # apply cosine transform and flip them
+        lats = -torch.arcsin(cost)
+        lons = _precompute_longitudes(self.nlon)
+
+        self.lmax = self.sht.lmax
+        self.mmax = self.sht.mmax
+
+        l = torch.arange(0, self.lmax).reshape(self.lmax, 1).cdouble()
+        l = l.expand(self.lmax, self.mmax)
+        # the laplace operator acting on the coefficients is given by l (l + 1)
+        lap = - l * (l + 1) / self.radius**2
+        invlap = - self.radius**2 / l / (l + 1)
+        invlap[0] = 0.
+
+        # register all
+        self.register_buffer('lats', lats)
+        self.register_buffer('lons', lons)
+        self.register_buffer('l', l)
+        self.register_buffer('lap', lap)
+        self.register_buffer('invlap', invlap)
+
+    def grid2spec(self, u):
+        return self.sht(u)
+
+    def spec2grid(self, uspec):
+        """Convert spectral coefficients to spatial data."""
+        return self.isht(uspec)
+
+    def dudtspec(self, uspec, pde='allen-cahn'):
+        """Compute the time derivative of spectral coefficients for different PDEs."""
+            
+        if pde == 'allen-cahn':
+            ugrid = self.spec2grid(uspec)
+            u3spec  = self.grid2spec(ugrid**3)
+            dudtspec = self.coeff*self.lap*uspec + uspec - u3spec
+        elif pde == 'ginzburg-landau':
+            ugrid = self.spec2grid(uspec)
+            u3spec  = self.grid2spec(ugrid**3)
+            dudtspec = uspec + (1. + 2.j)*self.coeff*self.lap*uspec - (1. + 2.j)*u3spec
+        else:
+            raise NotImplementedError(f"PDE type {pde} not implemented")
+
+        return dudtspec
+
+    def randspec(self):
+        """Generate random spectral data on the sphere."""
+        rspec = torch.randn_like(self.lap) / 4 / torch.pi
+        return rspec
+
+
+    def plot_griddata(self, data, fig, cmap='twilight_shifted', vmax=None, vmin=None, projection='3d', title=None, antialiased=False):
+        """
+        Plot data on the sphere grid. Requires cartopy for 3d plots.
+        
+        Parameters
+        -----------
+        data : torch.Tensor
+            Data to plot
+        fig : matplotlib.figure.Figure
+            Figure to plot on
+        cmap : str, optional
+            Colormap name, by default 'twilight_shifted'
+        vmax : float, optional
+            Maximum value for color scaling, by default None
+        vmin : float, optional
+            Minimum value for color scaling, by default None
+        projection : str, optional
+            Projection type ("mollweide", "3d"), by default "3d"
+        title : str, optional
+            Plot title, by default None
+        antialiased : bool, optional
+            Whether to use antialiasing, by default False
+            
+        Returns
+        -------
+        matplotlib.collections.QuadMesh
+            The plotted image object
+            
+        Raises
+        ------
+        NotImplementedError
+            If projection type is not supported
+        """
+        import matplotlib.pyplot as plt
+
+        lons = self.lons.squeeze() - torch.pi
+        lats = self.lats.squeeze()
+
+        if data.is_cuda:
+            data = data.cpu()
+            lons = lons.cpu()
+            lats = lats.cpu()
+
+        Lons, Lats = np.meshgrid(lons, lats)
+
+        if projection == 'mollweide':
+
+            #ax = plt.gca(projection=projection)
+            ax = fig.add_subplot(projection=projection)
+            im = ax.pcolormesh(Lons, Lats, data, cmap=cmap, vmax=vmax, vmin=vmin)
+            # ax.set_title("Elevation map of mars")
+            ax.grid(True)
+            ax.set_xticklabels([])
+            ax.set_yticklabels([])
+            plt.colorbar(im, orientation='horizontal')
+            plt.title(title)
+
+        elif projection == '3d':
+
+            import cartopy.crs as ccrs
+
+            proj = ccrs.Orthographic(central_longitude=0.0, central_latitude=25.0)
+
+            #ax = plt.gca(projection=proj, frameon=True)
+            ax = fig.add_subplot(projection=proj)
+            Lons = Lons*180/math.pi
+            Lats = Lats*180/math.pi
+
+            # contour data over the map.
+            im = ax.pcolormesh(Lons, Lats, data, cmap=cmap, transform=ccrs.PlateCarree(), antialiased=antialiased, vmax=vmax, vmin=vmin)
+            plt.title(title, y=1.05)
+
+        else:
+            raise NotImplementedError(f"Projection {projection} not implemented")
+
+        return im
+
+    def plot_specdata(self, data, fig, **kwargs):
+        """Plot spectral data by converting to spatial data first."""
+        return self.plot_griddata(self.isht(data), fig, **kwargs)

.deps/torch_harmonics/examples/shallow_water_equations.py

@@ -0,0 +1,394 @@
+# coding=utf-8
+
+# SPDX-FileCopyrightText: Copyright (c) 2022 The torch-harmonics Authors. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+#
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+#
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+#
+
+
+import torch
+import torch.nn as nn
+import torch_harmonics as th
+from torch_harmonics.quadrature import _precompute_longitudes
+
+import math
+import numpy as np
+
+
+class ShallowWaterSolver(nn.Module):
+    """
+    Shallow Water Equations (SWE) solver class for spherical geometry.
+    
+    Interface inspired by pyspharm and SHTns. Solves the shallow water equations
+    on a rotating sphere using spectral methods.
+    
+    Parameters
+    -----------
+    nlat : int
+        Number of latitude points
+    nlon : int
+        Number of longitude points
+    dt : float
+        Time step size
+    lmax : int, optional
+        Maximum l mode for spherical harmonics, by default None
+    mmax : int, optional
+        Maximum m mode for spherical harmonics, by default None
+    grid : str, optional
+        Grid type ("equiangular", "legendre-gauss", "lobatto"), by default "equiangular"
+    radius : float, optional
+        Radius of the sphere in meters, by default 6.37122E6 (Earth radius)
+    omega : float, optional
+        Angular velocity of rotation in rad/s, by default 7.292E-5 (Earth)
+    gravity : float, optional
+        Gravitational acceleration in m/s², by default 9.80616
+    havg : float, optional
+        Average height in meters, by default 10.e3
+    hamp : float, optional
+        Height amplitude in meters, by default 120.
+    """
+
+    def __init__(self, nlat, nlon, dt, lmax=None, mmax=None, grid="equiangular", radius=6.37122E6, \
+                 omega=7.292E-5, gravity=9.80616, havg=10.e3, hamp=120.):
+        super().__init__()
+
+        # time stepping param
+        self.dt = dt
+
+        # grid parameters
+        self.nlat = nlat
+        self.nlon = nlon
+        self.grid = grid
+
+        # physical sonstants
+        self.register_buffer('radius', torch.as_tensor(radius, dtype=torch.float64))
+        self.register_buffer('omega', torch.as_tensor(omega, dtype=torch.float64))
+        self.register_buffer('gravity', torch.as_tensor(gravity, dtype=torch.float64))
+        self.register_buffer('havg', torch.as_tensor(havg, dtype=torch.float64))
+        self.register_buffer('hamp', torch.as_tensor(hamp, dtype=torch.float64))
+
+        # SHT
+        self.sht = th.RealSHT(nlat, nlon, lmax=lmax, mmax=mmax, grid=grid, csphase=False)
+        self.isht = th.InverseRealSHT(nlat, nlon, lmax=lmax, mmax=mmax, grid=grid, csphase=False)
+        self.vsht = th.RealVectorSHT(nlat, nlon, lmax=lmax, mmax=mmax, grid=grid, csphase=False)
+        self.ivsht = th.InverseRealVectorSHT(nlat, nlon, lmax=lmax, mmax=mmax, grid=grid, csphase=False)
+
+        self.lmax = lmax or self.sht.lmax
+        self.mmax = lmax or self.sht.mmax
+
+        # compute gridpoints
+        if self.grid == "legendre-gauss":
+            cost, quad_weights = th.quadrature.legendre_gauss_weights(self.nlat, -1, 1)
+        elif self.grid == "lobatto":
+            cost, quad_weights = th.quadrature.lobatto_weights(self.nlat, -1, 1)
+        elif self.grid == "equiangular":
+            cost, quad_weights = th.quadrature.clenshaw_curtiss_weights(self.nlat, -1, 1)
+
+        quad_weights = quad_weights.reshape(-1, 1)
+
+        # apply cosine transform and flip them
+        lats = -torch.arcsin(cost)
+        lons = _precompute_longitudes(self.nlon)
+
+        self.lmax = self.sht.lmax
+        self.mmax = self.sht.mmax
+
+        # compute the laplace and inverse laplace operators
+        l = torch.arange(0, self.lmax).reshape(self.lmax, 1).double()
+        l = l.expand(self.lmax, self.mmax)
+        # the laplace operator acting on the coefficients is given by - l (l + 1)
+        lap = - l * (l + 1) / self.radius**2
+        invlap = - self.radius**2 / l / (l + 1)
+        invlap[0] = 0.
+
+        # compute coriolis force
+        coriolis = 2 * self.omega * torch.sin(lats).reshape(self.nlat, 1)
+
+        # hyperdiffusion
+        hyperdiff = torch.exp(torch.asarray((-self.dt / 2 / 3600.)*(lap / lap[-1, 0])**4))
+
+        # register all
+        self.register_buffer('lats', lats)
+        self.register_buffer('lons', lons)
+        self.register_buffer('l', l)
+        self.register_buffer('lap', lap)
+        self.register_buffer('invlap', invlap)
+        self.register_buffer('coriolis', coriolis)
+        self.register_buffer('hyperdiff', hyperdiff)
+        self.register_buffer('quad_weights', quad_weights)
+
+    def grid2spec(self, ugrid):
+        """Convert spatial data to spectral coefficients."""
+        return self.sht(ugrid)
+
+    def spec2grid(self, uspec):
+        """Convert spectral coefficients to spatial data."""
+        return self.isht(uspec)
+
+    def vrtdivspec(self, ugrid):
+        """Compute vorticity and divergence from velocity field."""
+        vrtdivspec = self.lap * self.radius * self.vsht(ugrid)
+        return vrtdivspec
+
+    def getuv(self, vrtdivspec):
+        """Compute wind vector from spectral coefficients of vorticity and divergence."""
+        return self.ivsht( self.invlap * vrtdivspec / self.radius)
+
+    def gethuv(self, uspec):
+        """Compute height and wind vector from spectral coefficients."""
+        hgrid = self.spec2grid(uspec[:1])
+        uvgrid = self.getuv(uspec[1:])
+        return torch.cat((hgrid, uvgrid), dim=-3)
+
+    def potential_vorticity(self, uspec):
+        """Compute potential vorticity from spectral coefficients."""
+        ugrid = self.spec2grid(uspec)
+        pvrt = (0.5 * self.havg * self.gravity / self.omega) * (ugrid[1] + self.coriolis) / ugrid[0]
+        return pvrt
+
+    def dimensionless(self, uspec):
+        """Remove dimensions from variables for dimensionless analysis."""
+        uspec[0] = (uspec[0] - self.havg * self.gravity) / self.hamp / self.gravity
+        # vorticity is measured in 1/s so we normalize using sqrt(g h) / r
+        uspec[1:] = uspec[1:] * self.radius / torch.sqrt(self.gravity * self.havg)
+        return uspec
+
+    def dudtspec(self, uspec):
+        """Compute time derivatives from solution represented in spectral coefficients."""
+        dudtspec = torch.zeros_like(uspec)
+
+        # compute the derivatives - this should be incorporated into the solver:
+        ugrid = self.spec2grid(uspec)
+        uvgrid = self.getuv(uspec[1:])
+
+        # phi = ugrid[0]
+        # vrtdiv = ugrid[1:]
+
+        tmp = uvgrid * (ugrid[1] + self.coriolis)
+        tmpspec = self.vrtdivspec(tmp)
+        dudtspec[2] = tmpspec[0]
+        dudtspec[1] = -1 * tmpspec[1]
+
+        tmp = uvgrid * ugrid[0]
+        tmp = self.vrtdivspec(tmp)
+        dudtspec[0] = -1 * tmp[1]
+
+        tmpspec = self.grid2spec(ugrid[0] + 0.5 * (uvgrid[0]**2 + uvgrid[1]**2))
+        dudtspec[2] = dudtspec[2] - self.lap * tmpspec
+
+        return dudtspec
+
+    def galewsky_initial_condition(self):
+        """Initialize non-linear barotropically unstable shallow water test case."""
+        device = self.lap.device
+
+        umax = 80.
+        phi0 = torch.asarray(torch.pi / 7., device=device)
+        phi1 = torch.asarray(0.5 * torch.pi - phi0, device=device)
+        phi2 = 0.25 * torch.pi
+        en = torch.exp(torch.asarray(-4.0 / (phi1 - phi0)**2, device=device))
+        alpha = 1. / 3.
+        beta = 1. / 15.
+
+        lats, lons = torch.meshgrid(self.lats, self.lons)
+
+        u1 = (umax/en)*torch.exp(1./((lats-phi0)*(lats-phi1)))
+        ugrid = torch.where(torch.logical_and(lats < phi1, lats > phi0), u1, torch.zeros(self.nlat, self.nlon, device=device))
+        vgrid = torch.zeros((self.nlat, self.nlon), device=device)
+        hbump = self.hamp * torch.cos(lats) * torch.exp(-((lons-torch.pi)/alpha)**2) * torch.exp(-(phi2-lats)**2/beta)
+
+        # intial velocity field
+        ugrid = torch.stack((ugrid, vgrid))
+        # intial vorticity/divergence field
+        vrtdivspec = self.vrtdivspec(ugrid)
+        vrtdivgrid = self.spec2grid(vrtdivspec)
+
+        # solve balance eqn to get initial zonal geopotential with a localized bump (not balanced).
+        tmp = ugrid * (vrtdivgrid + self.coriolis)
+        tmpspec = self.vrtdivspec(tmp)
+        tmpspec[1] = self.grid2spec(0.5 * torch.sum(ugrid**2, dim=0))
+        phispec = self.invlap*tmpspec[0] - tmpspec[1] + self.grid2spec(self.gravity*(self.havg + hbump))
+
+        # assemble solution
+        uspec = torch.zeros(3, self.lmax, self.mmax, dtype=vrtdivspec.dtype, device=device)
+        uspec[0] = phispec
+        uspec[1:] = vrtdivspec
+
+        return torch.tril(uspec)
+
+    def random_initial_condition(self, mach=0.1) -> torch.Tensor:
+        """Generate random initial condition on the sphere."""
+        device = self.lap.device
+        ctype = torch.complex128 if self.lap.dtype == torch.float64 else torch.complex64
+
+        # mach number relative to wave speed
+        llimit = mlimit = 120
+
+        # hgrid = self.havg + hamp * torch.randn(self.nlat, self.nlon, device=device, dtype=dtype)
+        # ugrid = uamp * torch.randn(self.nlat, self.nlon, device=device, dtype=dtype)
+        # vgrid = vamp * torch.randn(self.nlat, self.nlon, device=device, dtype=dtype)
+        # ugrid = torch.stack((ugrid, vgrid))
+
+        # initial geopotential
+        uspec = torch.zeros(3, self.lmax, self.mmax, dtype=ctype, device=self.lap.device)
+        uspec[:, :llimit, :mlimit] = torch.sqrt(torch.tensor(4 * torch.pi / llimit / (llimit+1), device=device, dtype=ctype)) * torch.randn_like(uspec[:, :llimit, :mlimit])
+
+        uspec[0] = self.gravity * self.hamp * uspec[0]
+        uspec[0, 0, 0] += torch.sqrt(torch.tensor(4 * torch.pi, device=device, dtype=ctype)) * self.havg * self.gravity
+        uspec[1:] = mach * uspec[1:] * torch.sqrt(self.gravity * self.havg) / self.radius
+        # uspec[1:] = self.vrtdivspec(self.spec2grid(uspec[1:]) * torch.cos(self.lats.reshape(-1, 1)))
+
+        # # intial velocity field
+        # ugrid = uamp * self.spec2grid(uspec[1])
+        # vgrid = vamp * self.spec2grid(uspec[2])
+        # ugrid = torch.stack((ugrid, vgrid))
+
+
+
+        # # intial vorticity/divergence field
+        # vrtdivspec = self.vrtdivspec(ugrid)
+        # vrtdivgrid = self.spec2grid(vrtdivspec)
+
+        # # solve balance eqn to get initial zonal geopotential with a localized bump (not balanced).
+        # tmp = ugrid * (vrtdivgrid + self.coriolis)
+        # tmpspec = self.vrtdivspec(tmp)
+        # tmpspec[1] = self.grid2spec(0.5 * torch.sum(ugrid**2, dim=0))
+        # phispec = self.invlap*tmpspec[0] - tmpspec[1] + self.grid2spec(self.gravity * hgrid)
+
+        # # assemble solution
+        # uspec = torch.zeros(3, self.lmax, self.mmax, dtype=phispec.dtype, device=device)
+        # uspec[0] = phispec
+        # uspec[1:] = vrtdivspec
+
+        return torch.tril(uspec)
+
+    def timestep(self, uspec: torch.Tensor, nsteps: int) -> torch.Tensor:
+        """Integrate the solution using Adams-Bashforth / forward Euler for nsteps steps."""
+        dudtspec = torch.zeros(3, 3, self.lmax, self.mmax, dtype=uspec.dtype, device=uspec.device)
+
+        # pointers to indicate the most current result
+        inew = 0
+        inow = 1
+        iold = 2
+
+        for iter in range(nsteps):
+            dudtspec[inew] = self.dudtspec(uspec)
+
+            # update vort,div,phiv with third-order adams-bashforth.
+            # forward euler, then 2nd-order adams-bashforth time steps to start.
+            if iter == 0:
+                dudtspec[inow] = dudtspec[inew]
+                dudtspec[iold] = dudtspec[inew]
+            elif iter == 1:
+                dudtspec[iold] = dudtspec[inew]
+
+            uspec = uspec + self.dt*( (23./12.) * dudtspec[inew] - (16./12.) * dudtspec[inow] + (5./12.) * dudtspec[iold] )
+
+            # implicit hyperdiffusion for vort and div.
+            uspec[1:] = self.hyperdiff * uspec[1:]
+
+            # cycle through the indices
+            inew = (inew - 1) % 3
+            inow = (inow - 1) % 3
+            iold = (iold - 1) % 3
+
+        return uspec
+
+    def integrate_grid(self, ugrid, dimensionless=False, polar_opt=0):
+        """Integrate the solution on the grid."""
+        dlon = 2 * torch.pi / self.nlon
+        radius = 1 if dimensionless else self.radius
+        if polar_opt > 0:
+            out = torch.sum(ugrid[..., polar_opt:-polar_opt, :] * self.quad_weights[polar_opt:-polar_opt] * dlon * radius**2, dim=(-2, -1))
+        else:
+            out = torch.sum(ugrid * self.quad_weights * dlon * radius**2, dim=(-2, -1))
+        return out
+
+
+    def plot_griddata(self, data, fig, cmap='twilight_shifted', vmax=None, vmin=None, projection='3d', title=None, antialiased=False):
+        """Plotting routine for data on the grid. Requires cartopy for 3d plots."""
+        import matplotlib.pyplot as plt
+
+        lons = self.lons.squeeze() - torch.pi
+        lats = self.lats.squeeze()
+
+        if data.is_cuda:
+            data = data.cpu()
+            lons = lons.cpu()
+            lats = lats.cpu()
+
+        Lons, Lats = np.meshgrid(lons, lats)
+
+        if projection == 'mollweide':
+
+            #ax = plt.gca(projection=projection)
+            ax = fig.add_subplot(projection=projection)
+            im = ax.pcolormesh(Lons, Lats, data, cmap=cmap, vmax=vmax, vmin=vmin)
+            # ax.set_title("Elevation map of mars")
+            ax.grid(True)
+            ax.set_xticklabels([])
+            ax.set_yticklabels([])
+            plt.colorbar(im, orientation='horizontal')
+            plt.title(title)
+
+        elif projection == '3d':
+
+            import cartopy.crs as ccrs
+
+            proj = ccrs.Orthographic(central_longitude=0.0, central_latitude=25.0)
+
+            #ax = plt.gca(projection=proj, frameon=True)
+            ax = fig.add_subplot(projection=proj)
+            Lons = Lons*180/math.pi
+            Lats = Lats*180/math.pi
+
+            # contour data over the map.
+            im = ax.pcolormesh(Lons, Lats, data, cmap=cmap, transform=ccrs.PlateCarree(), antialiased=antialiased, vmax=vmax, vmin=vmin)
+            plt.title(title, y=1.05)
+
+        elif projection == 'robinson':
+
+            import cartopy.crs as ccrs
+
+            proj = ccrs.Robinson(central_longitude=0.0)
+
+            #ax = plt.gca(projection=proj, frameon=True)
+            ax = fig.add_subplot(projection=proj)
+            Lons = Lons*180/math.pi
+            Lats = Lats*180/math.pi
+
+            # contour data over the map.
+            im = ax.pcolormesh(Lons, Lats, data, cmap=cmap, transform=ccrs.PlateCarree(), antialiased=antialiased, vmax=vmax, vmin=vmin)
+            plt.title(title, y=1.05)
+
+        else:
+            raise NotImplementedError
+
+        return im
+
+    def plot_specdata(self, data, fig, **kwargs):
+        return self.plot_griddata(self.isht(data), fig, **kwargs)

.deps/torch_harmonics/examples/stanford_2d3ds_dataset.py

@@ -0,0 +1,777 @@
+# coding=utf-8
+
+# SPDX-FileCopyrightText: Copyright (c) 2025 The torch-harmonics Authors. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+#
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+#
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+#
+
+import os
+import math
+
+import torch
+from torch.utils.data import Dataset, DataLoader, Subset
+
+import numpy as np
+
+from torch_harmonics.quadrature import _precompute_latitudes
+from torch_harmonics.examples.losses import get_quadrature_weights
+
+# some specifiers where to find the dataset
+DEFAULT_BASE_URL = "https://cvg-data.inf.ethz.ch/2d3ds/no_xyz/"
+DEFAULT_TAR_FILE_PAIRS = [
+    ("area_1_no_xyz.tar", "area_1"),
+    ("area_2_no_xyz.tar", "area_2"),
+    ("area_3_no_xyz.tar", "area_3"),
+    ("area_4_no_xyz.tar", "area_4"),
+    ("area_5a_no_xyz.tar", "area_5a"),
+    ("area_5b_no_xyz.tar", "area_5b"),
+    ("area_6_no_xyz.tar", "area_6"),
+]
+DEFAULT_LABELS_URL = "https://raw.githubusercontent.com/alexsax/2D-3D-Semantics/refs/heads/master/assets/semantic_labels.json"
+
+
+class Stanford2D3DSDownloader:
+    """
+    Convenience class for downloading the 2d3ds dataset [1].
+
+    Parameters
+    ----------
+    base_url : str, optional
+        Base URL for downloading the dataset, by default DEFAULT_BASE_URL
+    local_dir : str, optional
+        Local directory to store downloaded files, by default "data"
+
+    Returns
+    -------
+    data_folders : list
+        List of extracted directory names
+    class_labels : list
+        List of semantic class labels
+
+    References
+    ----------
+    .. [1] Armeni, I.,  Sax, S.,  Zamir, A. R.,  Savarese, S.;
+        "Joint 2D-3D-Semantic Data for Indoor Scene Understanding" (2017).
+        https://arxiv.org/abs/1702.01105.
+    """
+
+    def __init__(self, base_url: str = DEFAULT_BASE_URL, local_dir: str = "data"):
+       
+        self.base_url = base_url
+        self.local_dir = local_dir
+        os.makedirs(self.local_dir, exist_ok=True)
+
+    def _download_file(self, filename):
+        
+        import requests
+        from tqdm import tqdm
+
+        url = f"{self.base_url}/{filename}"
+        local_path = os.path.join(self.local_dir, filename)
+        if os.path.exists(local_path):
+            print(f"Note: Skipping download for {filename}, because it already exists")
+            return local_path
+
+        print(f"Downloading {filename}...")
+        temp_path = local_path.split(".")[0] + ".part"
+
+        # Resume logic
+        headers = {}
+        if os.path.exists(temp_path):
+            headers = {"Range": f"bytes={os.stat(temp_path).st_size}-"}
+
+        response = requests.get(url, headers=headers, stream=True, timeout=30)
+        if os.path.exists(temp_path):
+            total_size = int(response.headers.get("content-length", 0)) + os.stat(temp_path).st_size
+        else:
+            total_size = int(response.headers.get("content-length", 0))
+
+        with open(temp_path, "ab") as f, tqdm(desc=filename, total=total_size, unit="B", unit_scale=True, unit_divisor=1024, initial=os.stat(temp_path).st_size) as pbar:
+            for chunk in response.iter_content(chunk_size=1024):
+                if chunk:
+                    f.write(chunk)
+                    pbar.update(len(chunk))
+
+        os.rename(temp_path, local_path)
+        return local_path
+
+    def _extract_tar(self, tar_path):
+        
+        import tarfile
+
+        with tarfile.open(tar_path) as tar:
+            tar.extractall(path=self.local_dir)
+            tar_filenames = tar.getnames()
+            extracted_dir = tar_filenames[0]
+            os.remove(tar_path)
+            return extracted_dir
+
+    def download_dataset(self, file_extracted_directory_pairs=DEFAULT_TAR_FILE_PAIRS):
+        """
+        Download and extract the complete dataset.
+        
+        Parameters
+        -----------
+        file_extracted_directory_pairs : list, optional
+            List of (filename, extracted_folder_name) pairs, by default DEFAULT_TAR_FILE_PAIRS
+            
+        Returns
+        -------
+        tuple
+            (data_folders, class_labels) where data_folders is a list of extracted directory names
+            and class_labels is the semantic label mapping
+        """
+        import requests
+
+        data_folders = []
+        for file, extracted_folder_name in file_extracted_directory_pairs:
+            if not os.path.exists(os.path.join(self.local_dir, extracted_folder_name)):
+                downloaded_file = self._download_file(file)
+                data_folders.append(self._extract_tar(downloaded_file))
+            else:
+                print(f"Warning: Skipping D/L for '{file}' because folder '{extracted_folder_name}' already exists")
+                data_folders.append(extracted_folder_name)
+
+        labels_json_url = DEFAULT_LABELS_URL
+        class_labels = requests.get(labels_json_url).json()
+        return data_folders, class_labels
+
+    def _rgb_to_id(self, img, class_labels_map, class_labels_indices):
+        
+        # Convert to int32 first to avoid overflow
+        r = img[..., 0].astype(np.int32)
+        g = img[..., 1].astype(np.int32)
+        b = img[..., 2].astype(np.int32)
+        lookup_indices = r * 256 * 256 + g * 256 + b
+
+        def _convert(lookup: int) -> int:
+            # the dataset has a bad label for clutter, so we need to fix it
+            # clutter is 855309, but the labels file has it as 3341
+            # The original conversion used uint8, which overflowed the clutter label to 3341
+            # this is a fix to handle that accidental usage of undefined overflow behavior
+            if lookup == 855309:
+                label = class_labels_map[3341]  # clutter
+            else:
+                label = class_labels_map[lookup]
+            class_index = class_labels_indices.index(label)
+            return class_index
+
+        lookup_fn = np.vectorize(_convert)
+
+        return lookup_fn(lookup_indices)
+
+    def convert_dataset(
+        self,
+        data_folders,
+        class_labels,
+        rgb_path: str = "pano/rgb",
+        semantic_path: str = "pano/semantic",
+        depth_path: str = "pano/depth",
+        output_filename="semantic",
+        dataset_file: str = "stanford_2d3ds_dataset.h5",
+        downsampling_factor: int = 16,
+        remove_alpha_channel: bool = True,
+    ):
+        """
+        Convert the downloaded dataset to HDF5 format for efficient loading.
+        
+        Parameters
+        -----------
+        data_folders : list
+            List of extracted data folder names
+        class_labels : list
+            List of semantic class labels
+        rgb_path : str, optional
+            Relative path to RGB images within each data folder, by default "pano/rgb"
+        semantic_path : str, optional
+            Relative path to semantic labels within each data folder, by default "pano/semantic"
+        depth_path : str, optional
+            Relative path to depth images within each data folder, by default "pano/depth"
+        output_filename : str, optional
+            Suffix for semantic label files, by default "semantic"
+        dataset_file : str, optional
+            Output HDF5 filename, by default "stanford_2d3ds_dataset.h5"
+        downsampling_factor : int, optional
+            Factor by which to downsample images, by default 16
+        remove_alpha_channel : bool, optional
+            Whether to remove alpha channel from RGB images, by default True
+            
+        Returns
+        -------
+        str
+            Path to the created HDF5 dataset file
+        """
+        converted_dataset_path = os.path.join(self.local_dir, dataset_file)
+
+        from PIL import Image
+        from tqdm import tqdm
+        import h5py as h5
+
+        file_paths = []
+
+        min_vals = None
+        max_vals = None
+
+        # condition class labels first:
+        class_labels_map = [label.split("_")[0] for label in class_labels]
+        class_labels_indices = sorted(list(set(class_labels_map)))
+
+        # get all the file path input, output pairs
+        for base_path in data_folders:
+
+            rgb_dir = os.path.join(self.local_dir, base_path, rgb_path)
+            semantic_dir = os.path.join(self.local_dir, base_path, semantic_path)
+            depth_dir = os.path.join(self.local_dir, base_path, depth_path)
+
+            if os.path.exists(rgb_dir) and os.path.exists(semantic_dir) and os.path.exists(depth_dir):
+                for file_input in os.listdir(rgb_dir):
+                    if not file_input.endswith(".png"):
+                        continue
+                    rgb_filepath = os.path.join(rgb_dir, file_input)
+                    semantic_filepath = "_".join(os.path.splitext(os.path.basename(rgb_filepath))[0].split("_")[:-1]) + f"_{output_filename}.png"
+                    semantic_filepath = os.path.join(semantic_dir, semantic_filepath)
+                    depth_filepath = "_".join(os.path.splitext(os.path.basename(rgb_filepath))[0].split("_")[:-1]) + f"_depth.png"
+                    depth_filepath = os.path.join(depth_dir, depth_filepath)
+                    if not os.path.exists(semantic_filepath):
+                        print(f"Warning: Couldn't find output file in pair: ({rgb_filepath},{semantic_filepath})")
+                        continue
+
+                    if not os.path.exists(depth_filepath):
+                        print(f"Warning: Couldn't find depth file in pair: ({rgb_filepath},{depth_filepath})")
+                        continue
+
+                    file_paths.append((rgb_filepath, semantic_filepath, depth_filepath))
+            elif not os.path.exists(rgb_dir):
+                print("Warning: RGB dir doesn't exist: ", rgb_dir)
+                continue
+            elif not os.path.exists(semantic_dir):
+                print("Warning: Semantic dir doesn't exist: ", semantic_dir)
+                continue
+            elif not os.path.exists(depth_dir):
+                print("Warning: Depth dir doesn't exist: ", depth_dir)
+                continue
+
+        num_samples = len(file_paths)
+
+        if num_samples > 0:
+            first_rgb, first_semantic, first_depth = file_paths[0]
+            first_rgb = np.array(Image.open(first_rgb))
+            # first_semantic = np.array(Image.open(first_semantic))
+            # first_depth = np.array(Image.open(first_depth))
+
+            rgb_shape = first_rgb.shape
+            img_shape = (rgb_shape[0] // downsampling_factor, rgb_shape[1] // downsampling_factor)
+            rgb_channels = rgb_shape[2]
+
+            if remove_alpha_channel:
+                rgb_channels = 3
+        else:
+            raise ValueError(f"No samples found")
+
+        # create the dataset file
+        with h5.File(converted_dataset_path, "w") as h5file:
+            rgb_data = h5file.create_dataset("rgb", (num_samples, rgb_channels, *img_shape), "f4")
+            semantic_data = h5file.create_dataset("semantic", (num_samples, *img_shape), "i8")
+            depth_data = h5file.create_dataset("depth", (num_samples, *img_shape), "f4")
+            classes = h5file.create_dataset("class_labels", data=class_labels_indices)
+            num_classes = len(set(class_labels_indices))
+            data_source_path = h5file.create_dataset("data_source_path", (num_samples,), dtype=h5.string_dtype(encoding="utf-8"))
+            data_target_path = h5file.create_dataset("data_target_path", (num_samples,), dtype=h5.string_dtype(encoding="utf-8"))
+
+            # prepare computation of the class histogram
+            class_histogram = np.zeros(num_classes)
+            _, quad_weights = _precompute_latitudes(nlat=img_shape[0], grid="equiangular")
+            quad_weights = quad_weights.reshape(-1, 1) * 2 * torch.pi / float(img_shape[1])
+            quad_weights = quad_weights.tile(1, img_shape[1])
+            quad_weights /= torch.sum(quad_weights)
+            quad_weights = quad_weights.numpy()
+
+            for count in tqdm(range(num_samples), desc="preparing dataset"):
+                # open image
+                img = Image.open(file_paths[count][0])
+
+                # downsampling
+                if downsampling_factor != 1:
+                    # first width, then weight, weird
+                    img = img.resize(size=(img_shape[1], img_shape[0]), resample=Image.BILINEAR)
+
+                # remove alpha channel if requested
+                if remove_alpha_channel:
+                    img = img.convert("RGBA")
+                    background = Image.new("RGBA", img.size, (255, 255, 255))
+                    # compoe foreground and background and remove alpha channel
+                    img = np.array(Image.alpha_composite(background, img))
+                    r_data = img[:, :, :3]
+                else:
+                    r_data = np.array(img)
+
+                # transpose to channels first
+                r_data = np.transpose(r_data / 255.0, axes=(2, 0, 1))
+
+                # write to disk
+                rgb_data[count, ...] = r_data[...]
+                data_source_path[count] = file_paths[count][0]
+
+                # compute stats -> segmentation
+                # min/max
+                tmp_min = np.min(r_data, axis=(1, 2))
+                tmp_max = np.max(r_data, axis=(1, 2))
+                # mean/var
+                tmp_mean = np.sum(r_data * quad_weights[np.newaxis, :, :], axis=(1, 2))
+                tmp_m2 = np.sum(np.square(r_data - tmp_mean[:, np.newaxis, np.newaxis]) * quad_weights[np.newaxis, :, :])
+                if count == 0:
+                    # min/max
+                    min_vals = tmp_min
+                    max_vals = tmp_max
+                    # mean/var
+                    mean_vals = tmp_mean
+                    m2_vals = tmp_m2
+                else:
+                    # min/max
+                    min_vals = np.minimum(min_vals, tmp_min)
+                    max_vals = np.minimum(max_vals, tmp_max)
+                    # mean/var
+                    delta = tmp_mean - mean_vals
+                    mean_vals += delta / float(count + 1)
+                    m2_vals += tmp_m2 + delta * delta * float(count / (count + 1))
+
+                # get the target
+                sem = Image.open(file_paths[count][1])
+
+                # downsampling
+                if downsampling_factor != 1:
+                    sem = sem.resize(size=(img_shape[1], img_shape[0]), resample=Image.NEAREST)
+
+                sem_data = np.array(sem, dtype=np.uint32)
+
+                # map to classes
+                sem_data = self._rgb_to_id(sem_data, class_labels_map, class_labels_indices)
+
+                # write to file
+                semantic_data[count, ...] = sem_data[...]
+                data_target_path[count] = file_paths[count][1]
+
+                # Here we want depth
+                dep = Image.open(file_paths[count][2])
+
+                if downsampling_factor != 1:
+                    dep = dep.resize(size=(img_shape[1], img_shape[0]), resample=Image.NEAREST)
+                dep_data = np.array(dep)
+
+                depth_data[count, ...] = dep_data[...] / 65536.0
+
+                # compute stats -> depth
+                # min/max
+                tmp_min_depth = np.min(dep_data, axis=(0, 1))
+                tmp_max_depth = np.max(dep_data, axis=(0, 1))
+                # mean/var
+                tmp_mean_depth = np.sum(dep_data * quad_weights[:, :])
+                tmp_m2_depth = np.sum(np.square(dep_data - tmp_mean_depth) * quad_weights[:, :])
+                if count == 0:
+                    min_vals_depth = tmp_min_depth
+                    max_vals_depth = tmp_max_depth
+                    mean_vals_depth = tmp_mean_depth
+                    m2_vals_depth = tmp_m2_depth
+                else:
+                    min_vals_depth = np.minimum(min_vals_depth, tmp_min_depth)
+                    max_vals_depth = np.minimum(max_vals_depth, tmp_max_depth)
+                    delta = tmp_mean_depth - mean_vals_depth
+                    mean_vals_depth += delta / float(count + 1)
+                    m2_vals_depth += tmp_m2_depth + delta * delta * float(count / (count + 1))
+
+                # update the class histogram
+                for c in range(num_classes):
+                    class_histogram[c] += quad_weights[sem_data == c].sum()
+
+            # record min/max
+            h5file.create_dataset("min_rgb", data=min_vals.astype(np.float32))
+            h5file.create_dataset("max_rgb", data=max_vals.astype(np.float32))
+            h5file.create_dataset("mean_rgb", data=mean_vals.astype(np.float32))
+            std_vals = np.sqrt(m2_vals / float(num_samples - 1))
+            h5file.create_dataset("std_rgb", data=std_vals.astype(np.float32))
+
+            # record min/max
+            h5file.create_dataset("min_depth", data=min_vals_depth.astype(np.float32))
+            h5file.create_dataset("max_depth", data=max_vals_depth.astype(np.float32))
+            h5file.create_dataset("mean_depth", data=mean_vals_depth.astype(np.float32))
+            std_vals_depth = np.sqrt(m2_vals_depth / float(num_samples - 1))
+            h5file.create_dataset("std_depth", data=std_vals_depth.astype(np.float32))
+
+            # record class histogram
+            class_histogram = class_histogram / num_samples
+            h5file.create_dataset("class_histogram", data=class_histogram.astype(np.float32))
+
+        return converted_dataset_path
+
+    def prepare_dataset(self, file_extracted_directory_pairs=DEFAULT_TAR_FILE_PAIRS, dataset_file: str = "stanford_2d3ds_dataset.h5", downsampling_factor: int = 16):
+
+        converted_dataset_path = os.path.join(self.local_dir, dataset_file)
+        if os.path.exists(converted_dataset_path):
+            print(
+                f"Dataset file at {converted_dataset_path} already exists. Skipping download and conversion. If you want to create a new dataset file, delete or rename the existing file."
+            )
+            return converted_dataset_path
+
+        data_folders, class_labels = self.download_dataset(file_extracted_directory_pairs=file_extracted_directory_pairs)
+        converted_dataset_path = self.convert_dataset(data_folders=data_folders, class_labels=class_labels, dataset_file=dataset_file, downsampling_factor=downsampling_factor)
+
+        self.converted_dataset_path = converted_dataset_path
+
+        return self.converted_dataset_path
+
+
+class StanfordSegmentationDataset(Dataset):
+    """
+    Spherical segmentation dataset from [1].
+
+    Parameters
+    ----------
+    dataset_file : str
+        Path to the HDF5 dataset file
+    ignore_alpha_channel : bool, optional
+        Whether to ignore the alpha channel in the RGB images, by default True
+    log_depth : bool, optional
+        Whether to log the depth values, by default False
+    exclude_polar_fraction : float, optional
+        Fraction of polar points to exclude, by default 0.0
+
+    Returns
+    -------
+    StanfordSegmentationDataset
+        Dataset object
+
+    References
+    ----------
+    .. [1] Armeni, I.,  Sax, S.,  Zamir, A. R.,  Savarese, S.;
+        "Joint 2D-3D-Semantic Data for Indoor Scene Understanding" (2017).
+        https://arxiv.org/abs/1702.01105.
+    """
+
+    def __init__(
+        self,
+        dataset_file,
+        ignore_alpha_channel=True,
+        exclude_polar_fraction=0,
+    ):
+
+        import h5py as h5
+
+        self.dataset_file = dataset_file
+        self.exclude_polar_fraction = exclude_polar_fraction
+
+        with h5.File(self.dataset_file, "r") as h5file:
+            self.img_rgb = h5file["rgb"][0].shape
+            self.img_seg = h5file["semantic"][0].shape
+            self.num_samples = h5file["rgb"].shape[0]
+            self.num_classes = h5file["class_labels"].shape[0]
+
+            self.class_labels = [class_name.decode("utf-8") for class_name in h5file["class_labels"][...].tolist()]
+            self.class_histogram = np.array(h5file["class_histogram"][...])
+            self.class_histogram = self.class_histogram / self.class_histogram.sum()
+
+            self.mean = h5file["mean_rgb"][...]
+            self.std = h5file["std_rgb"][...]
+            self.min = h5file["min_rgb"][...]
+            self.max = h5file["max_rgb"][...]
+
+            self.img_filepath = h5file["data_source_path"][...]
+            self.tar_filepath = h5file["data_target_path"][...]
+
+        if ignore_alpha_channel:
+            self.img_rgb = (3, self.img_rgb[1], self.img_rgb[2])
+
+        # open file and check for
+        self.h5file = None
+        self.rgb = None
+        self.semantic = None
+
+        # return index set to false by default
+        # when true, the __getitem__ method will return the index of the input,target pair
+        self.return_index = False
+
+    @property
+    def target_shape(self):
+        return self.img_seg
+
+    @property
+    def input_shape(self):
+        return self.img_rgb
+
+    def set_return_index(self, return_index: bool):
+        self.return_index = return_index
+
+    def get_img_filepath(self, idx: int):
+        return self.img_filepath[idx]
+
+    def get_tar_filepath(self, idx: int):
+        return self.tar_filepath[idx]
+
+    def _id_to_class(self, class_id):
+        if class_id > self.num_classes:
+            print("WARNING: ID > number of classes!")
+            return None
+        return self.segmentation_classes[class_id]
+
+    def _mask_invalid(self, tar):
+        return np.where(tar >= self.num_classes, -100, tar)
+
+    def __len__(self):
+        return self.num_samples
+
+    def _init_files(self):
+        import h5py as h5
+
+        self.h5file = h5.File(self.dataset_file, "r")
+        self.rgb = self.h5file["rgb"]
+        self.semantic = self.h5file["semantic"]
+
+    def reset(self):
+        self.rgb = None
+        self.semantic = None
+        if self.h5file is not None:
+            self.h5file.close()
+            del self.h5file
+        self.h5file = None
+
+    def __getitem__(self, idx, mask_invalid=True):
+
+        if self.h5file is None:
+            # init files
+            self._init_files()
+
+        rgb = self.rgb[idx, 0 : self.img_rgb[0], 0 : self.img_rgb[1], 0 : self.img_rgb[2]]
+        sem = self.semantic[idx, 0 : self.img_seg[0], 0 : self.img_seg[1]]
+        if mask_invalid:
+            sem = self._mask_invalid(sem)
+
+        if self.exclude_polar_fraction > 0:
+            hcut = int(self.exclude_polar_fraction * sem.shape[0])
+            if hcut > 0:
+                sem[0:hcut, :] = -100
+                sem[-hcut:, :] = -100
+
+        return rgb, sem
+
+
+class StanfordDatasetSubset(Subset):
+    def __init__(self, dataset, indices, return_index=False):
+        super().__init__(dataset, indices)
+        self.return_index = return_index
+        self.dataset = dataset
+
+    def set_return_index(self, value):
+        self.return_index = value
+
+    def __getitem__(self, index):
+        real_index = self.indices[index]
+        data = self.dataset[real_index]
+
+        if self.return_index:
+            return data[0], data[1], real_index
+        else:
+            # Otherwise, return only (data, target)
+            return data[0], data[1]
+
+
+class StanfordDepthDataset(Dataset):
+    """
+    Spherical segmentation dataset from [1].
+
+    Parameters
+    ----------
+    dataset_file : str
+        Path to the HDF5 dataset file
+    ignore_alpha_channel : bool, optional
+        Whether to ignore the alpha channel in the RGB images, by default True
+    log_depth : bool, optional
+        Whether to log the depth values, by default False
+    exclude_polar_fraction : float, optional
+        Fraction of polar points to exclude, by default 0.0
+
+    References
+    ----------
+    .. [1] Armeni, I.,  Sax, S.,  Zamir, A. R.,  Savarese, S.;
+        "Joint 2D-3D-Semantic Data for Indoor Scene Understanding" (2017).
+        https://arxiv.org/abs/1702.01105.
+    """
+
+    def __init__(self, dataset_file, ignore_alpha_channel=True, log_depth=False, exclude_polar_fraction=0.0):
+
+        import h5py as h5
+
+        self.dataset_file = dataset_file
+        self.log_depth = log_depth
+        self.exclude_polar_fraction = exclude_polar_fraction
+        with h5.File(self.dataset_file, "r") as h5file:
+            self.img_rgb = h5file["rgb"][0].shape
+            self.img_depth = h5file["depth"][0].shape
+            self.num_samples = h5file["rgb"].shape[0]
+
+            self.mean_in = h5file["mean_rgb"][...]
+            self.std_in = h5file["std_rgb"][...]
+            self.min_in = h5file["min_rgb"][...]
+            self.max_in = h5file["max_rgb"][...]
+
+            self.mean_out = h5file["mean_depth"][...]
+            self.std_out = h5file["std_depth"][...]
+            self.min_out = h5file["min_depth"][...]
+            self.max_out = h5file["max_depth"][...]
+
+        if ignore_alpha_channel:
+            self.img_rgb = (3, self.img_rgb[1], self.img_rgb[2])
+
+        # open file and check for
+        self.h5file = None
+        self.rgb = None
+        self.depth = None
+
+    @property
+    def target_shape(self):
+        return self.img_depth
+
+    @property
+    def input_shape(self):
+        return self.img_rgb
+
+    def __len__(self):
+        return self.num_samples
+
+    def _init_files(self):
+        import h5py as h5
+
+        self.h5file = h5.File(self.dataset_file, "r")
+        self.rgb = self.h5file["rgb"]
+        self.depth = self.h5file["depth"]
+
+    def reset(self):
+        self.rgb = None
+        self.depth = None
+        if self.h5file is not None:
+            self.h5file.close()
+            del self.h5file
+        self.h5file = None
+
+    def _mask_invalid(self, tar):
+        return tar * np.where(tar == tar.max(), 0, 1)
+
+    def __getitem__(self, idx, mask_invalid=True):
+
+        if self.h5file is None:
+            # init files
+            self._init_files()
+
+        rgb = self.rgb[idx, 0 : self.img_rgb[0], 0 : self.img_rgb[1], 0 : self.img_rgb[2]]
+
+        depth = self.depth[idx, 0 : self.img_depth[0], 0 : self.img_depth[1]]
+        if mask_invalid:
+            depth = self._mask_invalid(depth)
+
+        if self.exclude_polar_fraction > 0:
+            hcut = int(self.exclude_polar_fraction * depth.shape[0])
+            if hcut > 0:
+                depth[0:hcut, :] = 0
+                depth[-hcut:, :] = 0
+
+        if self.log_depth:
+            depth = np.log(1 + depth)
+
+        return rgb, depth
+
+
+def compute_stats_s2(dataset: Dataset, normalize_target: bool = False):
+    """
+    Compute stats using parallel welford reduction and quadrature on the sphere. 
+    The parallel welford reduction follows this article (parallel algorithm): https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance
+    """
+
+    nexamples = len(dataset)
+    count = 0
+    for isample in range(nexamples):
+        token = dataset[isample]
+
+        # dimension of inp and tar are (3, nlat, nlon)
+        inp, tar = token
+        nlat = tar.shape[-2]
+        nlon = tar.shape[-1]
+
+        # pre-compute quadrature weights
+        if isample == 0:
+            quad_weights = get_quadrature_weights(nlat=inp.shape[1], nlon=inp.shape[2], grid="equiangular", tile=True).numpy().astype(np.float64)
+
+        # this is a special case for the depth dataset
+        # TODO: maybe make this an argument
+        if normalize_target:
+            mask = np.where(tar == 0, 0, 1)
+            masked_area = np.sum(mask * quad_weights[np.newaxis, :, :], axis=(-2, -1))
+
+        # get initial welford values
+        if isample == 0:
+            # input
+            inp_means = np.sum(inp * quad_weights[np.newaxis, :, :], axis=(-2, -1))
+            inp_m2s = np.sum(np.square(inp - inp_means[:, np.newaxis, np.newaxis]) * quad_weights[np.newaxis, :, :], axis=(-2, -1))
+
+            # target
+            if normalize_target:
+                tar_means = np.sum(mask * tar * quad_weights[np.newaxis, :, :], axis=(-2, -1)) / masked_area
+                tar_m2s = np.sum(mask * np.square(tar - tar_means[:, np.newaxis, np.newaxis]) * quad_weights[np.newaxis, :, :], axis=(-2, -1)) / masked_area
+
+            # update count
+            count = 1
+
+        # do welford update
+        else:
+            # input
+            # get new mean and m2
+            inp_mean = np.sum(inp * quad_weights[np.newaxis, :, :], axis=(-2, -1))
+            inp_m2 = np.sum(np.square(inp - inp_mean[:, np.newaxis, np.newaxis]) * quad_weights[np.newaxis, :, :], axis=(-2, -1))
+            # update welford values
+            inp_delta = inp_mean - inp_means
+            inp_m2s = inp_m2s + inp_m2 + inp_delta**2 * count / float(count + 1)
+            inp_means = inp_means + inp_delta / float(count + 1)
+
+            # target
+            if normalize_target:
+                # get new mean and m2
+                tar_mean = np.sum(mask * tar * quad_weights[np.newaxis, :, :], axis=(-2, -1)) / masked_area
+                tar_m2 = np.sum(mask * np.square(tar - tar_mean[:, np.newaxis, np.newaxis]) * quad_weights[np.newaxis, :, :], axis=(-2, -1)) / masked_area
+                # update welford values
+                tar_delta = tar_mean - tar_means
+                tar_m2s = tar_m2s + tar_m2 + tar_delta**2 * count / float(count + 1)
+                tar_means = tar_means + tar_delta / float(count + 1)
+
+            # update count
+            count += 1
+
+    # finalize
+    inp_stds = np.sqrt(inp_m2s / float(count))
+    result = (inp_means.astype(np.float32), inp_stds.astype(np.float32))
+
+    if normalize_target:
+        tar_stds = np.sqrt(tar_m2s / float(count))
+        result += (tar_means.astype(np.float32), tar_stds.astype(np.float32))
+
+    return result

.deps/torch_harmonics/filter_basis.py

@@ -0,0 +1,326 @@
+# coding=utf-8
+
+# SPDX-FileCopyrightText: Copyright (c) 2022 The torch-harmonics Authors. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+#
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+#
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+#
+
+import abc
+from typing import Tuple, Union, Optional
+import math
+
+import torch
+
+from torch_harmonics.cache import lru_cache
+
+
+def _circle_dist(x1: torch.Tensor, x2: torch.Tensor):
+    return torch.minimum(torch.abs(x1 - x2), torch.abs(2 * math.pi - torch.abs(x1 - x2)))
+
+
+def _log_factorial(x: torch.Tensor):
+    return torch.lgamma(x + 1)
+
+
+def _factorial(x: torch.Tensor):
+    return torch.exp(_log_factorial(x))
+
+
+class FilterBasis(metaclass=abc.ABCMeta):
+    """Abstract base class for a filter basis"""
+
+    def __init__(
+        self,
+        kernel_shape: Union[int, Tuple[int], Tuple[int, int]],
+    ):
+
+        self.kernel_shape = kernel_shape
+
+    @property
+    @abc.abstractmethod
+    def kernel_size(self):
+
+        raise NotImplementedError
+
+    # @abc.abstractmethod
+    # def compute_vals(self, r: torch.Tensor, phi: torch.Tensor, r_cutoff: float):
+    #     """
+    #     Computes the values of the filter basis
+    #     """
+    #     raise NotImplementedError
+
+    @abc.abstractmethod
+    def compute_support_vals(self, r: torch.Tensor, phi: torch.Tensor, r_cutoff: float):
+
+        raise NotImplementedError
+
+
+@lru_cache(typed=True, copy=False)
+def get_filter_basis(kernel_shape: Union[int, Tuple[int], Tuple[int, int]], basis_type: str) -> FilterBasis:
+    """Factory function to generate the appropriate filter basis"""
+
+    if basis_type == "piecewise linear":
+        return PiecewiseLinearFilterBasis(kernel_shape=kernel_shape)
+    elif basis_type == "morlet":
+        return MorletFilterBasis(kernel_shape=kernel_shape)
+    elif basis_type == "zernike":
+        return ZernikeFilterBasis(kernel_shape=kernel_shape)
+    else:
+        raise ValueError(f"Unknown basis_type {basis_type}")
+
+
+class PiecewiseLinearFilterBasis(FilterBasis):
+    """Tensor-product basis on a disk constructed from piecewise linear basis functions."""
+
+    def __init__(
+        self,
+        kernel_shape: Union[int, Tuple[int], Tuple[int, int]],
+    ):
+
+        if isinstance(kernel_shape, int):
+            kernel_shape = [kernel_shape]
+        if len(kernel_shape) == 1:
+            kernel_shape = [kernel_shape[0], 1]
+        elif len(kernel_shape) != 2:
+            raise ValueError(f"expected kernel_shape to be a list or tuple of length 1 or 2 but got {kernel_shape} instead.")
+
+        super().__init__(kernel_shape=kernel_shape)
+
+    @property
+    def kernel_size(self):
+        """Compute the number of basis functions in the kernel."""
+        return (self.kernel_shape[0] // 2) * self.kernel_shape[1] + self.kernel_shape[0] % 2
+
+    def _compute_support_vals_isotropic(self, r: torch.Tensor, phi: torch.Tensor, r_cutoff: float):
+
+        # enumerator for basis function
+        ikernel = torch.arange(self.kernel_size, device=r.device).reshape(-1, 1, 1)
+
+        # collocation points
+        nr = self.kernel_shape[0]
+        dr = 2 * r_cutoff / (nr + 1)
+
+        # compute the support
+        if nr % 2 == 1:
+            ir = ikernel * dr
+        else:
+            ir = (ikernel + 0.5) * dr
+
+        # find the indices where the rotated position falls into the support of the kernel
+        iidx = torch.argwhere(((r - ir).abs() <= dr) & (r <= r_cutoff))
+        vals = 1 - (r[iidx[:, 1], iidx[:, 2]] - ir[iidx[:, 0], 0, 0]).abs() / dr
+
+        return iidx, vals
+
+    def _compute_support_vals_anisotropic(self, r: torch.Tensor, phi: torch.Tensor, r_cutoff: float):
+
+        # enumerator for basis function
+        ikernel = torch.arange(self.kernel_size, device=r.device).reshape(-1, 1, 1)
+
+        # collocation points
+        nr = self.kernel_shape[0]
+        nphi = self.kernel_shape[1]
+        dr = 2 * r_cutoff / (nr + 1)
+        dphi = 2.0 * math.pi / nphi
+
+        # disambiguate even and uneven cases and compute the support
+        if nr % 2 == 1:
+            ir = ((ikernel - 1) // nphi + 1) * dr
+            iphi = ((ikernel - 1) % nphi) * dphi - math.pi
+        else:
+            ir = (ikernel // nphi + 0.5) * dr
+            iphi = (ikernel % nphi) * dphi - math.pi
+
+        # find the indices where the rotated position falls into the support of the kernel
+        if nr % 2 == 1:
+            # find the support
+            cond_r = ((r - ir).abs() <= dr) & (r <= r_cutoff)
+            cond_phi = (ikernel == 0) | (_circle_dist(phi, iphi).abs() <= dphi)
+            # find indices where conditions are met
+            iidx = torch.argwhere(cond_r & cond_phi)
+            # compute the distance to the collocation points
+            dist_r = (r[iidx[:, 1], iidx[:, 2]] - ir[iidx[:, 0], 0, 0]).abs()
+            dist_phi = _circle_dist(phi[iidx[:, 1], iidx[:, 2]], iphi[iidx[:, 0], 0, 0])
+            # compute the value of the basis functions
+            vals  = 1 - dist_r / dr
+            vals *= torch.where((iidx[:, 0] > 0), (1 - dist_phi / dphi), 1.0)
+
+        else:
+            # in the even case, the inner basis functions overlap into areas with a negative areas
+            rn = -r
+            phin = torch.where(phi + math.pi >= math.pi, phi - math.pi, phi + math.pi)
+            # find the support
+            cond_r = ((r - ir).abs() <= dr) & (r <= r_cutoff)
+            cond_phi = _circle_dist(phi, iphi).abs() <= dphi
+            cond_rn = ((rn - ir).abs() <= dr) & (rn <= r_cutoff)
+            cond_phin = _circle_dist(phin, iphi) <= dphi
+            # find indices where conditions are met
+            iidx = torch.argwhere((cond_r & cond_phi) | (cond_rn & cond_phin))
+
+            dist_r = (r[iidx[:, 1], iidx[:, 2]] - ir[iidx[:, 0], 0, 0]).abs()
+            dist_phi = _circle_dist(phi[iidx[:, 1], iidx[:, 2]], iphi[iidx[:, 0], 0, 0])
+            dist_rn = (rn[iidx[:, 1], iidx[:, 2]] - ir[iidx[:, 0], 0, 0]).abs()
+            dist_phin = _circle_dist(phin[iidx[:, 1], iidx[:, 2]], iphi[iidx[:, 0], 0, 0])
+            # compute the value of the basis functions
+            vals = cond_r[iidx[:, 0], iidx[:, 1], iidx[:, 2]] * (1 - dist_r / dr)
+            vals *= cond_phi[iidx[:, 0], iidx[:, 1], iidx[:, 2]] * (1 - dist_phi / dphi)
+            valsn = cond_rn[iidx[:, 0], iidx[:, 1], iidx[:, 2]] * (1 - dist_rn / dr)
+            valsn *= cond_phin[iidx[:, 0], iidx[:, 1], iidx[:, 2]] * (1 - dist_phin / dphi)
+            vals += valsn
+
+        return iidx, vals
+
+    def compute_support_vals(self, r: torch.Tensor, phi: torch.Tensor, r_cutoff: float):
+        """Computes the index set that falls into the kernel's support and returns both indices and values."""
+
+        if self.kernel_shape[1] > 1:
+            return self._compute_support_vals_anisotropic(r, phi, r_cutoff=r_cutoff)
+        else:
+            return self._compute_support_vals_isotropic(r, phi, r_cutoff=r_cutoff)
+
+
+class MorletFilterBasis(FilterBasis):
+    """Morlet-style filter basis on the disk. A Gaussian is multiplied with a Fourier basis in x and y directions."""
+
+    def __init__(
+        self,
+        kernel_shape: Union[int, Tuple[int], Tuple[int, int]],
+    ):
+
+        if isinstance(kernel_shape, int):
+            kernel_shape = [kernel_shape, kernel_shape]
+        if len(kernel_shape) != 2:
+            raise ValueError(f"expected kernel_shape to be a list or tuple of 2 but got {kernel_shape} instead.")
+
+        super().__init__(kernel_shape=kernel_shape)
+
+    @property
+    def kernel_size(self):
+
+        return self.kernel_shape[0] * self.kernel_shape[1]
+
+    def gaussian_window(self, r: torch.Tensor, width: float = 1.0):
+
+        return 1 / (2 * math.pi * width**2) * torch.exp(-0.5 * r**2 / (width**2))
+
+    def hann_window(self, r: torch.Tensor, width: float = 1.0):
+
+        return torch.cos(0.5 * torch.pi * r / width) ** 2
+
+    def compute_support_vals(self, r: torch.Tensor, phi: torch.Tensor, r_cutoff: float, width: float = 1.0):
+
+        # enumerator for basis function
+        ikernel = torch.arange(self.kernel_size, device=r.device).reshape(-1, 1, 1)
+        nkernel = ikernel % self.kernel_shape[1]
+        mkernel = ikernel // self.kernel_shape[1]
+
+        # get relevant indices
+        iidx = torch.argwhere((r <= r_cutoff) & torch.full_like(ikernel, True, dtype=torch.bool, device=r.device))
+
+        # get corresponding r, phi, x and y coordinates
+        r = r[iidx[:, 1], iidx[:, 2]] / r_cutoff
+        phi = phi[iidx[:, 1], iidx[:, 2]]
+        x = r * torch.sin(phi)
+        y = r * torch.cos(phi)
+        n = nkernel[iidx[:, 0], 0, 0]
+        m = mkernel[iidx[:, 0], 0, 0]
+
+        harmonic = torch.where(n % 2 == 1, torch.sin(torch.ceil(n / 2) * math.pi * x / width), torch.cos(torch.ceil(n / 2) * math.pi * x / width))
+        harmonic *= torch.where(m % 2 == 1, torch.sin(torch.ceil(m / 2) * math.pi * y / width), torch.cos(torch.ceil(m / 2) * math.pi * y / width))
+
+        # computes the envelope. To ensure that the curve is roughly 0 at the boundary, we rescale the Gaussian by 0.25
+        vals = self.hann_window(r, width=width) * harmonic
+
+        return iidx, vals
+
+
+class ZernikeFilterBasis(FilterBasis):
+    """Zernike polynomials which are defined on the disk. See https://en.wikipedia.org/wiki/Zernike_polynomials"""
+
+    def __init__(
+        self,
+        kernel_shape: Union[int, Tuple[int]],
+    ):
+
+        if isinstance(kernel_shape, tuple) or isinstance(kernel_shape, list):
+            kernel_shape = kernel_shape[0]
+        if not isinstance(kernel_shape, int):
+            raise ValueError(f"expected kernel_shape to be an integer but got {kernel_shape} instead.")
+
+        super().__init__(kernel_shape=kernel_shape)
+
+    @property
+    def kernel_size(self):
+
+        return (self.kernel_shape * (self.kernel_shape + 1)) // 2
+
+    def zernikeradial(self, r: torch.Tensor, n: torch.Tensor, m: torch.Tensor):
+
+        out = torch.zeros_like(r)
+        bound = (n - m) // 2 + 1
+        max_bound = bound.max().item()
+
+        for k in range(max_bound):
+
+            inc = (-1) ** k * _factorial(n - k) * r ** (n - 2 * k) / (math.factorial(k) * _factorial((n + m) // 2 - k) * _factorial((n - m) // 2 - k))
+            out += torch.where(k < bound, inc, 0.0)
+
+        return out
+
+    def zernikepoly(self, r: torch.Tensor, phi: torch.Tensor, n: torch.Tensor, l: torch.Tensor):
+
+        m = 2 * l - n
+        return torch.where(m < 0, self.zernikeradial(r, n, -m) * torch.sin(m * phi), self.zernikeradial(r, n, m) * torch.cos(m * phi))
+
+    def compute_support_vals(self, r: torch.Tensor, phi: torch.Tensor, r_cutoff: float, width: float = 0.25):
+
+        # enumerator for basis function
+        ikernel = torch.arange(self.kernel_size, device=r.device).reshape(-1, 1, 1)
+
+        # get relevant indices
+        iidx = torch.argwhere((r <= r_cutoff) & torch.full_like(ikernel, True, dtype=torch.bool, device=r.device))
+
+        # indexing logic for zernike polynomials
+        # the total index is given by (n * (n + 2) + l ) // 2 which needs to be reversed
+        # precompute shifts in the level of the "pyramid"
+        nshifts = torch.arange(self.kernel_shape, device=r.device)
+        nshifts = (nshifts + 1) * nshifts // 2
+        # find the level and position within the pyramid
+        nkernel = torch.searchsorted(nshifts, ikernel, right=True) - 1
+        lkernel = ikernel - nshifts[nkernel]
+        # mkernel = 2 * ikernel - nkernel * (nkernel + 2)
+
+        # get corresponding coordinates and n and l indices
+        r = r[iidx[:, 1], iidx[:, 2]] / r_cutoff
+        phi = phi[iidx[:, 1], iidx[:, 2]]
+        n = nkernel[iidx[:, 0], 0, 0]
+        l = lkernel[iidx[:, 0], 0, 0]
+
+        # computes the Zernike polynomials using helper functions
+        vals = self.zernikepoly(r, phi, n, l)
+
+        return iidx, vals

.deps/torch_harmonics/legendre.py

@@ -0,0 +1,215 @@
+# coding=utf-8
+
+# SPDX-FileCopyrightText: Copyright (c) 2022 The torch-harmonics Authors. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+# 
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+#
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+#
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+#
+
+from typing import Optional
+import math
+import torch
+
+from torch_harmonics.cache import lru_cache
+
+
+def clm(l: int, m: int) -> float:
+    """Defines the normalization factor to orthonormalize the Spherical Harmonics."""
+    return math.sqrt((2*l + 1) / 4 / math.pi) * math.sqrt(math.factorial(l-m) / math.factorial(l+m))
+
+def legpoly(mmax: int, lmax: int, x: torch.Tensor, norm: Optional[str]="ortho", inverse: Optional[bool]=False, csphase: Optional[bool]=True) -> torch.Tensor:
+    """
+    Computes the values of (-1)^m c^l_m P^l_m(x) at the positions specified by x.
+    The resulting tensor has shape (mmax, lmax, len(x)). The Condon-Shortley Phase (-1)^m
+    can be turned off optionally.
+
+    Parameters
+    -----------
+    mmax: int
+        Maximum order of the spherical harmonics
+    lmax: int
+        Maximum degree of the spherical harmonics
+    x: torch.Tensor
+        Tensor of positions at which to evaluate the Legendre polynomials
+    norm: Optional[str]
+        Normalization of the Legendre polynomials
+    inverse: Optional[bool]
+        Whether to compute the inverse Legendre polynomials
+    csphase: Optional[bool]
+        Whether to apply the Condon-Shortley phase (-1)^m
+
+    Returns
+    -------
+    out: torch.Tensor
+        Tensor of Legendre polynomial values
+
+    References
+    ----------
+    [1] Schaeffer, N.; Efficient spherical harmonic transforms aimed at pseudospectral numerical simulations, G3: Geochemistry, Geophysics, Geosystems.
+    [2] Rapp, R.H.; A Fortran Program for the Computation of Gravimetric Quantities from High Degree Spherical Harmonic Expansions, Ohio State University Columbus; report; 1982;
+        https://apps.dtic.mil/sti/citations/ADA123406
+    [3] Schrama, E.; Orbit integration based upon interpolated gravitational gradients
+    """
+
+    # compute the tensor P^m_n:
+    nmax = max(mmax,lmax)
+    vdm = torch.zeros((nmax, nmax, len(x)), dtype=torch.float64, device=x.device, requires_grad=False)
+        
+    norm_factor = 1.0 if norm == "ortho" else math.sqrt(4 * math.pi)
+    norm_factor = 1.0 / norm_factor if inverse else norm_factor
+
+    # initial values to start the recursion
+    vdm[0,0,:] = norm_factor / math.sqrt(4 * math.pi)
+
+    # fill the diagonal and the lower diagonal
+    for l in range(1, nmax):
+        vdm[l-1, l, :] = math.sqrt(2*l + 1) * x * vdm[l-1, l-1, :]
+        vdm[l, l, :] = torch.sqrt( (2*l + 1) * (1 + x) * (1 - x) / 2 / l ) * vdm[l-1, l-1, :]
+
+    # fill the remaining values on the upper triangle and multiply b
+    for l in range(2, nmax):
+        for m in range(0, l-1):
+            vdm[m, l, :] = x * math.sqrt((2*l - 1) / (l - m) * (2*l + 1) / (l + m)) * vdm[m, l-1, :] \
+                            - math.sqrt((l + m - 1) / (l - m) * (2*l + 1) / (2*l - 3) * (l - m - 1) / (l + m)) * vdm[m, l-2, :]
+
+    if norm == "schmidt":
+        for l in range(0, nmax):
+            if inverse:
+                vdm[:, l, : ] = vdm[:, l, : ] * math.sqrt(2*l + 1)
+            else:
+                vdm[:, l, : ] = vdm[:, l, : ] / math.sqrt(2*l + 1)
+
+    vdm = vdm[:mmax, :lmax]
+
+    if csphase:
+        for m in range(1, mmax, 2):
+            vdm[m] *= -1
+
+    return vdm
+
+@lru_cache(typed=True, copy=True)
+def _precompute_legpoly(mmax: int , lmax: int, t: torch.Tensor,
+                        norm: Optional[str]="ortho", inverse: Optional[bool]=False, csphase: Optional[bool]=True) -> torch.Tensor:
+    """
+    Computes the values of (-1)^m c^l_m P^l_m(\cos \theta) at the positions specified by t (theta).
+    The resulting tensor has shape (mmax, lmax, len(x)). The Condon-Shortley Phase (-1)^m
+    can be turned off optionally.
+
+    Parameters
+    -----------
+    mmax: int
+        Maximum order of the spherical harmonics
+    lmax: int
+        Maximum degree of the spherical harmonics
+    t: torch.Tensor
+        Tensor of positions at which to evaluate the Legendre polynomials
+    norm: Optional[str]
+        Normalization of the Legendre polynomials
+    inverse: Optional[bool]
+        Whether to compute the inverse Legendre polynomials
+    csphase: Optional[bool]
+        Whether to apply the Condon-Shortley phase (-1)^m
+
+    Returns
+    -------
+    out: torch.Tensor
+        Tensor of Legendre polynomial values
+
+    References
+    ----------
+    [1] Schaeffer, N.; Efficient spherical harmonic transforms aimed at pseudospectral numerical simulations, G3: Geochemistry, Geophysics, Geosystems.
+    [2] Rapp, R.H.; A Fortran Program for the Computation of Gravimetric Quantities from High Degree Spherical Harmonic Expansions, Ohio State University Columbus; report; 1982;
+        https://apps.dtic.mil/sti/citations/ADA123406
+    [3] Schrama, E.; Orbit integration based upon interpolated gravitational gradients
+    """
+
+    return legpoly(mmax, lmax, torch.cos(t), norm=norm, inverse=inverse, csphase=csphase)
+
+@lru_cache(typed=True, copy=True)
+def _precompute_dlegpoly(mmax: int, lmax: int, t: torch.Tensor,
+                         norm: Optional[str]="ortho", inverse: Optional[bool]=False, csphase: Optional[bool]=True) -> torch.Tensor:
+    """
+    Computes the values of the derivatives $\frac{d}{d \theta} P^m_l(\cos \theta)$
+    at the positions specified by t (theta), as well as $\frac{1}{\sin \theta} P^m_l(\cos \theta)$,
+    needed for the computation of the vector spherical harmonics. The resulting tensor has shape
+    (2, mmax, lmax, len(t)).
+
+    Parameters
+    -----------
+    mmax: int
+        Maximum order of the spherical harmonics
+    lmax: int
+        Maximum degree of the spherical harmonics
+    t: torch.Tensor
+        Tensor of positions at which to evaluate the Legendre polynomials
+    norm: Optional[str]
+        Normalization of the Legendre polynomials
+    inverse: Optional[bool]
+        Whether to compute the inverse Legendre polynomials
+    csphase: Optional[bool]
+        Whether to apply the Condon-Shortley phase (-1)^m
+
+    Returns
+    -------
+    out: torch.Tensor
+        Tensor of Legendre polynomial values
+
+    References
+    ----------
+    [2] Wang, B., Wang, L., Xie, Z.; Accurate calculation of spherical and vector spherical harmonic expansions via spectral element grids; Adv Comput Math.
+    """
+
+    pct = _precompute_legpoly(mmax+1, lmax+1, t, norm=norm, inverse=inverse, csphase=False)
+
+    dpct = torch.zeros((2, mmax, lmax, len(t)), dtype=torch.float64, device=t.device, requires_grad=False)
+
+    # fill the derivative terms wrt theta
+    for l in range(0, lmax):
+
+        # m = 0
+        dpct[0, 0, l] = - math.sqrt(l*(l+1)) * pct[1, l]
+
+        # 0 < m < l
+        for m in range(1, min(l, mmax)):
+            dpct[0, m, l] = 0.5 * ( math.sqrt((l+m)*(l-m+1)) * pct[m-1, l] - math.sqrt((l-m)*(l+m+1)) * pct[m+1, l] )
+
+        # m == l
+        if mmax > l:
+            dpct[0, l, l] = math.sqrt(l/2) * pct[l-1, l]
+
+        # fill the - 1j m P^m_l / sin(phi). as this component is purely imaginary,
+        # we won't store it explicitly in a complex array
+        for m in range(1, min(l+1, mmax)):
+            # this component is implicitly complex
+            # we do not divide by m here as this cancels with the derivative of the exponential
+            dpct[1, m, l] = 0.5 * math.sqrt((2*l+1)/(2*l+3)) * \
+                ( math.sqrt((l-m+1)*(l-m+2)) * pct[m-1, l+1] + math.sqrt((l+m+1)*(l+m+2)) * pct[m+1, l+1] )
+
+    if csphase:
+        for m in range(1, mmax, 2):
+            dpct[:, m] *= -1
+
+    return dpct

.deps/torch_harmonics/plotting.py

@@ -0,0 +1,265 @@
+# coding=utf-8
+
+# SPDX-FileCopyrightText: Copyright (c) 2025 The torch-harmonics Authors. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+#
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+#
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+#
+
+import numpy as np
+import os
+
+# guarded imports
+try:
+    import matplotlib.pyplot as plt
+except ImportError as err:
+    plt = None
+
+try:
+    import cartopy
+    import cartopy.crs as ccrs
+except ImportError as err:
+    cartopy = None
+    ccrs = None
+
+
+def check_plotting_dependencies():
+    if plt is None:
+        raise ImportError("matplotlib is required for plotting functions. Install it with 'pip install matplotlib'")
+    if cartopy is None:
+        raise ImportError("cartopy is required for map plotting. Install it with 'pip install cartopy'")
+
+
+def get_projection(
+    projection,
+    central_latitude=0,
+    central_longitude=0,
+):
+    """
+    Get a cartopy projection object for map plotting.
+    
+    Parameters
+    -----------
+    projection : str
+        Projection type ("orthographic", "robinson", "platecarree", "mollweide")
+    central_latitude : float, optional
+        Central latitude for the projection, by default 0
+    central_longitude : float, optional
+        Central longitude for the projection, by default 0
+        
+    Returns
+    -------
+    cartopy.crs.Projection
+        Cartopy projection object
+        
+    Raises
+    ------
+    ValueError
+        If projection type is not supported
+    """
+    if projection == "orthographic":
+        proj = ccrs.Orthographic(central_latitude=central_latitude, central_longitude=central_longitude)
+    elif projection == "robinson":
+        proj = ccrs.Robinson(central_longitude=central_longitude)
+    elif projection == "platecarree":
+        proj = ccrs.PlateCarree(central_longitude=central_longitude)
+    elif projection == "mollweide":
+        proj = ccrs.Mollweide(central_longitude=central_longitude)
+    else:
+        raise ValueError(f"Unknown projection mode {projection}")
+
+    return proj
+
+
+def plot_sphere(
+    data, fig=None, projection="robinson", cmap="RdBu", title=None, colorbar=False, coastlines=False, gridlines=False,  central_latitude=0, central_longitude=0, lon=None, lat=None, **kwargs
+):
+    """
+    Plots a function defined on the sphere using pcolormesh
+    
+    Parameters
+    -----------
+    data : numpy.ndarray or torch.Tensor
+        Data to plot with shape (nlat, nlon)
+    fig : matplotlib.figure.Figure, optional
+        Figure to plot on, by default None (creates new figure)
+    projection : str, optional
+        Map projection type, by default "robinson"
+    cmap : str, optional
+        Colormap name, by default "RdBu"
+    title : str, optional
+        Plot title, by default None
+    colorbar : bool, optional
+        Whether to add a colorbar, by default False
+    coastlines : bool, optional
+        Whether to add coastlines, by default False
+    gridlines : bool, optional
+        Whether to add gridlines, by default False
+    central_latitude : float, optional
+        Central latitude for projection, by default 0
+    central_longitude : float, optional
+        Central longitude for projection, by default 0
+    lon : numpy.ndarray, optional
+        Longitude coordinates, by default None (auto-generated)
+    lat : numpy.ndarray, optional
+        Latitude coordinates, by default None (auto-generated)
+    **kwargs
+        Additional arguments passed to pcolormesh
+        
+    Returns
+    -------
+    matplotlib.collections.QuadMesh
+        The plotted image object
+    """
+
+    # make sure cartopy exist
+    check_plotting_dependencies()
+
+    if fig == None:
+        fig = plt.figure()
+
+    nlat = data.shape[-2]
+    nlon = data.shape[-1]
+    if lon is None:
+        lon = np.linspace(0, 2 * np.pi, nlon + 1)[:-1]
+    if lat is None:
+        lat = np.linspace(np.pi / 2.0, -np.pi / 2.0, nlat)
+    Lon, Lat = np.meshgrid(lon, lat)
+
+    # convert radians to degrees
+    Lon = Lon * 180 / np.pi
+    Lat = Lat * 180 / np.pi
+
+    # get the projection. Latitude is shifted to match plot_sphere
+    proj = get_projection(projection, central_latitude=central_latitude, central_longitude=central_longitude)
+
+    ax = fig.add_subplot(projection=proj)
+
+    # contour data over the map.
+    im = ax.pcolormesh(Lon, Lat, data, cmap=cmap, transform=ccrs.PlateCarree(), antialiased=False, **kwargs)
+
+    # add features if requested
+    if coastlines:
+        ax.add_feature(cartopy.feature.COASTLINE, edgecolor="white", facecolor="none", linewidth=1.5)
+
+    # add colorbar if requested
+    if colorbar:
+        plt.colorbar(im)
+
+    # add gridlines
+    if gridlines:
+        gl = ax.gridlines(crs=ccrs.PlateCarree(), draw_labels=False, linewidth=1, color="gray", alpha=0.6, linestyle="--")
+
+    # add title with smaller font
+    plt.title(title, y=1.05, fontsize=8)
+
+    return im
+
+
+def imshow_sphere(data, fig=None, projection="robinson", title=None, central_latitude=0, central_longitude=0, **kwargs):
+    """
+    Displays an image on the sphere
+    
+    Parameters
+    -----------
+    data : numpy.ndarray or torch.Tensor
+        Data to display with shape (nlat, nlon)
+    fig : matplotlib.figure.Figure, optional
+        Figure to plot on, by default None (creates new figure)
+    projection : str, optional
+        Map projection type, by default "robinson"
+    title : str, optional
+        Plot title, by default None
+    central_latitude : float, optional
+        Central latitude for projection, by default 0
+    central_longitude : float, optional
+        Central longitude for projection, by default 0
+    **kwargs
+        Additional arguments passed to imshow
+        
+    Returns
+    -------
+    matplotlib.image.AxesImage
+        The displayed image object
+    """
+
+    # make sure cartopy exist
+    check_plotting_dependencies()
+
+    if fig == None:
+        fig = plt.figure()
+
+    # get the projection. Latitude is shifted to match plot_sphere
+    proj = get_projection(projection, central_latitude=central_latitude, central_longitude=central_longitude + 180)
+
+    ax = fig.add_subplot(projection=proj)
+
+    # contour data over the map.
+    im = ax.imshow(data, transform=ccrs.PlateCarree(), **kwargs)
+
+    # add title
+    plt.title(title, y=1.05)
+
+    return im
+
+
+# def plot_data(data,
+#                 fig=None,
+#                 cmap="RdBu",
+#                 title=None,
+#                 colorbar=False,
+#                 coastlines=False,
+#                 central_longitude=0,
+#                 lon=None,
+#                 lat=None,
+#                 **kwargs):
+#     if fig == None:
+#         fig = plt.figure()
+
+#     nlat = data.shape[-2]
+#     nlon = data.shape[-1]
+#     if lon is None:
+#         lon = np.linspace(0, 2*np.pi, nlon+1)[:-1]
+#     if lat is None:
+#         lat = np.linspace(np.pi/2., -np.pi/2., nlat)
+#     Lon, Lat = np.meshgrid(lon, lat)
+
+#     proj = ccrs.Robinson(central_longitude=central_longitude)
+#     # proj = ccrs.Mollweide(central_longitude=central_longitude)
+
+#     ax = fig.add_subplot(projection=proj)
+#     Lon = Lon*180/np.pi
+#     Lat = Lat*180/np.pi
+
+#     # contour data over the map.
+#     im = ax.pcolormesh(Lon, Lat, data, cmap=cmap, transform=ccrs.PlateCarree(), antialiased=False, **kwargs)
+#     if coastlines:
+#         ax.add_feature(cartopy.feature.COASTLINE, edgecolor='white', facecolor='none', linewidth=1.5)
+#     if colorbar:
+#         plt.colorbar(im)
+#     plt.title(title, y=1.05)
+
+#     return im