JAX Support

.github/workflows/docs.yaml

@@ -23,7 +23,7 @@ jobs:
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
-          pip install sphinx furo
+          pip install -r docs/requirements.txt 
       - name: Build website
         run: |
           sphinx-build -M dirhtml docs docs/_build

.github/workflows/release.yaml

@@ -6,7 +6,7 @@ on:
 # ref: https://packaging.python.org/en/latest/guides/publishing-package-distribution-releases-using-github-actions-ci-cd-workflows/
 
 jobs:
-  build:
+  build-oeq:
     name: Build distribution
     runs-on: ubuntu-latest
     steps:
@@ -16,21 +16,20 @@ jobs:
       with:
         python-version: '3.10'
     - name: install dependencies, then build source tarball
-      run: | 
+      run: |
+        cd openequivariance 
         python3 -m pip install build --user
         python3 -m build --sdist 
     - name: store the distribution packages 
       uses: actions/upload-artifact@v4
       with:
         name: python-package-distributions
-        path: dist/
+        path: openequivariance/dist/
 
   pypi-publish:
     name: Upload release to PyPI
     runs-on: ubuntu-latest
-    # build task to be completed first
-    needs: build
-    # Specifying a GitHub environment is optional, but strongly encouraged
+    needs: build-oeq
     environment: 
         name: pypi
         url: https://pypi.org/p/openequivariance
@@ -42,6 +41,47 @@ jobs:
         uses: actions/download-artifact@v4
         with:
           name: python-package-distributions
-          path: dist/
+          path: openequivariance/dist/
+      - name: publish package distributions to PyPI
+        uses: pypa/gh-action-pypi-publish@release/v1
+
+  # ------------------------------------
+
+  build-oeq-extjax:
+    name: Build distribution
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v4
+    - name: set up Python
+      uses: actions/setup-python@v5
+      with:
+        python-version: '3.10'
+    - name: install dependencies, then build source tarball
+      run: |
+        cd openequivariance_extjax 
+        python3 -m pip install build --user
+        python3 -m build --sdist 
+    - name: store the distribution packages 
+      uses: actions/upload-artifact@v4
+      with:
+        name: python-package-distributions
+        path: openequivariance_extjax/dist/
+
+  pypi-publish-extjax:
+    name: Upload release to PyPI
+    runs-on: ubuntu-latest
+    needs: build-oeq-extjax
+    environment: 
+        name: pypi
+        url: https://pypi.org/p/openequivariance_extjax
+    permissions:
+      # IMPORTANT: this permission is mandatory for Trusted Publishing
+      id-token: write
+    steps:
+      - name: download the distributions
+        uses: actions/download-artifact@v4
+        with:
+          name: python-package-distributions
+          path: openequivariance_extjax/dist/
       - name: publish package distributions to PyPI
         uses: pypa/gh-action-pypi-publish@release/v1

.github/workflows/requirements_cuda_ci.txt

@@ -1,4 +1,6 @@
 numpy==2.2.5
 torch==2.7.0 --index-url https://download.pytorch.org/whl/cu128
 pytest==8.3.5
-ninja==1.11.1.4
+ninja==1.11.1.4
+nanobind==2.10.2
+scikit-build-core==0.11.6

.github/workflows/verify_extension_build.yml

@@ -1,4 +1,4 @@
-name: OEQ CUDA C++ Extension Build Verification
+name: OEQ C++ Extension Build Verification
 
 on:
   push:
@@ -29,10 +29,14 @@ jobs:
         sudo apt-get update
         sudo apt install nvidia-cuda-toolkit
         pip install -r .github/workflows/requirements_cuda_ci.txt
-        pip install -e .
+        pip install -e "./openequivariance"
 
-    - name: Test extension build via import 
+    - name: Test CUDA extension build via import 
       run: |
         pytest \
           tests/import_test.py::test_extension_built \
-          tests/import_test.py::test_torch_extension_built 
+          tests/import_test.py::test_torch_extension_built 
+
+    - name: Test JAX extension build
+      run: |
+        XLA_DIRECT_DOWNLOAD=1 pip install -e "./openequivariance_extjax" --no-build-isolation 

.gitignore

@@ -38,7 +38,6 @@ triton_autotuning
 paper_benchmarks
 paper_benchmarks_v2
 paper_benchmarks_v3
-openequivariance/extlib/*.so
 
 get_node.sh
 *.egg-info

CHANGELOG.md

@@ -1,5 +1,17 @@
 ## Latest Changes
 
+### v0.5.0 (2025-12-25) 
+JAX support is now available in 
+OpenEquivariance for BOTH NVIDIA and
+AMD GPUs! See the 
+[documentation](https://passionlab.github.io/OpenEquivariance/)
+and README.md for instructions on installation
+and usage. 
+
+Minor changes:
+- Defer error reporting when CUDA is not available
+  to the first library usage in code, not library load. 
+
 ### v0.4.1 (2025-09-04)
 Minor update, fixes a bug loading JIT-compiled modules
 with PyTorch 2.9.

README.md

@@ -1,8 +1,9 @@
 # OpenEquivariance
-[![OEQ CUDA C++ Extension Build Verification](https://github.com/PASSIONLab/OpenEquivariance/actions/workflows/verify_extension_build.yml/badge.svg?event=push)](https://github.com/PASSIONLab/OpenEquivariance/actions/workflows/verify_extension_build.yml)
+[![OEQ C++ Extension Build Verification](https://github.com/PASSIONLab/OpenEquivariance/actions/workflows/verify_extension_build.yml/badge.svg?event=push)](https://github.com/PASSIONLab/OpenEquivariance/actions/workflows/verify_extension_build.yml)
 [![License](https://img.shields.io/badge/License-BSD_3--Clause-blue.svg)](https://opensource.org/licenses/BSD-3-Clause)
 
-[[Examples]](#show-me-some-examples) 
+[[PyTorch Examples]](#pytorch-examples) 
+[[JAX Examples]](#jax-examples)
 [[Citation and Acknowledgements]](#citation-and-acknowledgements)
 
 OpenEquivariance is a CUDA and HIP kernel generator for the Clebsch-Gordon tensor product, 
@@ -12,8 +13,8 @@ that [e3nn](https://e3nn.org/) supports
 commonly found in graph neural networks 
 (e.g. [Nequip](https://github.com/mir-group/nequip) or
 [MACE](https://github.com/ACEsuit/mace)). To get 
-started, ensure that you have GCC 9+ on your system 
-and install our package via
+started with PyTorch, ensure that you have PyTorch 
+and GCC 9+ available before installing our package via 
 
 ```bash
 pip install openequivariance
@@ -29,11 +30,26 @@ computation and memory consumption significantly.
 For detailed instructions on tests, benchmarks, MACE / Nequip, and our API,
 check out the [documentation](https://passionlab.github.io/OpenEquivariance).
 
-📣 📣 OpenEquivariance was accepted to the 2025 SIAM Conference on Applied and 
-Computational Discrete Algorithms (Proceedings Track)! Catch the talk in 
-Montréal and check out the [camera-ready copy on Arxiv](https://arxiv.org/abs/2501.13986) (available May 12, 2025).
+⭐️ **JAX**: Our latest update brings
+support for JAX. For NVIDIA GPUs, 
+install it (after installing JAX) 
+with the following two commands strictly in order:
 
-## Show me some examples
+``` bash
+pip install openequivariance[jax]
+pip install openequivariance_extjax --no-build-isolation
+```
+
+For AMD GPUs:
+``` bash
+pip install openequivariance[jax]
+JAX_HIP=1 pip install openequivariance_extjax --no-build-isolation
+```
+
+See the section below for example usage and 
+our [API page](https://passionlab.github.io/OpenEquivariance/api/) for more details.
+
+## PyTorch Examples 
 Here's a CG tensor product implemented by e3nn: 
 
 ```python
@@ -127,6 +143,48 @@ print(torch.norm(Z))
 `deterministic=False`, the `sender` and `receiver` indices can have
 arbitrary order.
 
+## JAX Examples
+After installation, use the library
+as follows. Set `OEQ_NOTORCH=1`
+in your environment to avoid the PyTorch import in
+the regular `openequivariance` package.
+```python
+import jax
+import os
+
+os.environ["OEQ_NOTORCH"] = "1"
+import openequivariance as oeq
+
+seed = 42
+key = jax.random.PRNGKey(seed)
+
+batch_size = 1000
+X_ir, Y_ir, Z_ir = oeq.Irreps("1x2e"), oeq.Irreps("1x3e"), oeq.Irreps("1x2e")
+problem = oeq.TPProblem(X_ir, Y_ir, Z_ir, [(0, 0, 0, "uvu", True)], shared_weights=False, internal_weights=False)
+
+
+node_ct, nonzero_ct = 3, 4
+edge_index = jax.numpy.array(
+    [
+        [0, 1, 1, 2],
+        [1, 0, 2, 1],
+    ],
+    dtype=jax.numpy.int32, # NOTE: This int32, not int64
+)
+
+X = jax.random.uniform(key, shape=(node_ct, X_ir.dim), minval=0.0, maxval=1.0, dtype=jax.numpy.float32)
+Y = jax.random.uniform(key, shape=(nonzero_ct, Y_ir.dim),
+                        minval=0.0, maxval=1.0, dtype=jax.numpy.float32)
+W = jax.random.uniform(key, shape=(nonzero_ct, problem.weight_numel),
+                        minval=0.0, maxval=1.0, dtype=jax.numpy.float32)
+
+tp_conv = oeq.jax.TensorProductConv(problem, deterministic=False)
+Z = tp_conv.forward(
+    X, Y, W, edge_index[0], edge_index[1]
+)
+print(jax.numpy.linalg.norm(Z))
+```
+
 ## Citation and Acknowledgements
 If you find this code useful, please cite our paper:
 

docs/api.rst

@@ -8,7 +8,7 @@ OpenEquivariance API
 OpenEquivariance exposes two key classes: :py:class:`openequivariance.TensorProduct`, which replaces
 ``o3.TensorProduct`` from e3nn, and :py:class:`openequivariance.TensorProductConv`, which fuses
 the CG tensor product with a subsequent graph convolution. Initializing either class triggers
-JIT compilation of a custom kernel, which can take a few seconds. 
+JIT compilation of a custom kernel, which can take a few seconds.
 
 Both classes require a configuration object specified 
 by :py:class:`openequivariance.TPProblem`, which has a constructor
@@ -17,6 +17,9 @@ We recommend reading the `e3nn documentation <https://docs.e3nn.org/en/latest/>`
 trying our code. OpenEquivariance cannot accelerate all tensor products; see 
 :doc:`this page </supported_ops>` for a list of supported configurations.
 
+PyTorch API
+------------------------
+
 .. autoclass:: openequivariance.TensorProduct
     :members: forward, reorder_weights_from_e3nn, reorder_weights_to_e3nn, to
     :undoc-members:
@@ -27,14 +30,39 @@ trying our code. OpenEquivariance cannot accelerate all tensor products; see
     :undoc-members:
     :exclude-members: name
 
-.. autoclass:: openequivariance.TPProblem
-    :members:
-    :undoc-members:
-
 .. autofunction:: openequivariance.torch_to_oeq_dtype
 
 .. autofunction:: openequivariance.torch_ext_so_path
 
+JAX API
+------------------------
+The JAX API consists of ``TensorProduct`` and ``TensorProductConv``
+classes that behave identically to their PyTorch counterparts. These classes
+do not conform exactly to the e3nn-jax API, but perform the same computation.
+
+If you plan to use ``oeq.jax`` without PyTorch installed, 
+you need to set ``OEQ_NOTORCH=1`` in your local environment (within Python,
+``os.environ["OEQ_NOTORCH"] = 1``). For the moment, we require this to avoid 
+breaking the PyTorch version of OpenEquivariance.
+
+
+.. autoclass:: openequivariance.jax.TensorProduct
+    :members: forward, reorder_weights_from_e3nn, reorder_weights_to_e3nn
+    :undoc-members:
+    :exclude-members:
+
+.. autoclass:: openequivariance.jax.TensorProductConv
+    :members: forward, reorder_weights_from_e3nn, reorder_weights_to_e3nn
+    :undoc-members:
+    :exclude-members: 
+
+Common API
+---------------------
+
+.. autoclass:: openequivariance.TPProblem
+    :members:
+    :undoc-members:
+
 API Identical to e3nn
 ---------------------
 

docs/conf.py

@@ -28,11 +28,17 @@
 html_theme = "furo"
 # html_static_path = ["_static"]
 
-extensions = [
-    "sphinx.ext.autodoc",
+extensions = ["sphinx.ext.autodoc", "sphinx_inline_tabs"]
+
+sys.path.insert(0, str(Path("../openequivariance").resolve()))
+
+autodoc_mock_imports = [
+    "torch",
+    "jax",
+    "openequivariance._torch.extlib",
+    "openequivariance.jax.extlib",
+    "openequivariance_extjax",
+    "jinja2",
+    "numpy",
 ]
-
-sys.path.insert(0, str(Path("..").resolve()))
-
-autodoc_mock_imports = ["torch", "openequivariance.extlib", "jinja2", "numpy"]
 autodoc_typehints = "description"