🌐 e3j

A fast, full-featured and platform-agnostic equivariance backend for GPU and TPU!

This library can be used in place of e3nn and e3x to replace slow operations in Machine Learned Interatomic Potentials (MLIPs), with carefully optimized and open-source CUDA and Pallas kernels.

The equivariance backend of our MLIP library is e3j as of mlip 0.2.0.

Installation

Pulling from PyPI

The e3j package is available on PyPI. It consists of a thin JAX-based Python API which can run on CPU, GPU and TPU, supporting Python versions from 3.11 to 3.14 included.

For efficiency on GPU, our CUDA binaries are bundled as the e3j_ops package on PyPI. The compatible version of the binaries should be pulled by requiring the "e3j[ops]" extra:

# requirements.txt
e3j[ops] >= 0.1.0b0
jax[cuda13_local] ~= 0.8.0

See JAX installation instructions for more information on JAX versions and their CUDA support. We recommend using a version of JAX above 0.7.0 and CUDA 13.

Building from source

Our dependencies are managed with uv. After cloning the repository, you can build from source by running run one of:

# Existing CUDA 13 install with `e3j_ops` kernels:
uv sync --group cuda13_local --extra ops
# Install CUDA 13 via pip:
uv sync --group cuda13 --extra ops

The Python build internally relies on CMake, scikit-build and pybind11. You can also look at the Makefile for alternate recipes to build kernels, C++ tests and the Python bindings.

The e3j_ops Python package only contains our CUDA binaries and bindings to their associated XLA handlers. It is not meant to be used as standalone until its ABI is reported stable.

Features

e3j provides a platform-agnostic API for GPU and TPU:

• 🖥️ The same Python API on CPU, GPU and TPU, with a portable JAX fallback when kernels don't apply
• 🔗 Interoperability made easy via plain jax.Array semantics
• 🌐 All equivariant building blocks: spherical harmonics (Harmonics), tensor products (TensorProduct, Bigotimes), message-passing convolution (Convolution) and learnable linear maps (Linear, LinearIndexwise) as flax.linen.Module with weight initializations matching e3nn
• 🏎️ Fused CUDA kernels for GPU (tensor product, message-passing convolution, scatter-add), shipped as the standalone e3j_ops wheel and dispatched through XLA-FFI
• 🧮 Fused Pallas Mosaic-TPU kernels, computing gather, tensor product, scalar mixing and scatter in a single kernel
• 🔁 Custom VJP rules for every fused kernel, so they differentiate under jax.grad like any other JAX primitive
• 🎛️ SPMD support for muliple-device MLIP training
• 🧱 Multiple memory layouts (leading channels, trailing channels, and a flat e3nn-compatible layout) to trade coalescing off against interoperability
• 📐 Representation utilities: O(3) and SO(3) spaces, irreps filtering and irrep filtering, permutations and generalized Clebsch-Gordan coefficients
• 🔌 Full coverage of the e3nn and e3x layers used by an MLIP, kernel-backed or not, so an existing model can be ported over entirely — to train, simulate and benchmark end to end, see mlip

Contributing

Bug reports, questions and feature requests are welcome via GitHub issues and discussions.

Citing

If you use e3j within your work, we kindly ask you to cite the following preprint:

@article{Peltre26-e3j,
    title   = {{E3J}: an Efficient and Open-Source Euclidean Equivariance Backend},
    author  = {Peltre, Olivier and Picard, Armand and Pichard, Adrien and Giacomoni, Luca and Braganca, Miguel and Heyraud, Valentin and Brunken, Christoph and Tilly, Jules},
    journal = {preprint},
    year    = {2026},
    url     = {(preprint)}
  }
}