# E3J: An Efficient and Open-Source Backend for Euclidean Equivariant Operations on GPU and TPU

arXiv: https://arxiv.org/abs/2609.35099 · HTML: https://mlsanity.com/paper/2609.35099/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 19 of its 31 references that are in the map.
- First posted on arXiv 28 Sep 2026.

## Builds on (most cited first, 19 of 19)

- [Tensor field networks: Rotation- and translation-equivariant neural networks for 3D point clouds](https://mlsanity.com/paper/1802.08219/) (2018)
- [E(3)-Equivariant Graph Neural Networks for Data-Efficient and Accurate Interatomic Potentials](https://mlsanity.com/paper/2101.03164/) (2021)
- [MACE: Higher Order Equivariant Message Passing Neural Networks for Fast and Accurate Force Fields](https://mlsanity.com/paper/2206.07697/) (2022)
- [Equiformer: Equivariant Graph Attention Transformer for 3D Atomistic Graphs](https://mlsanity.com/paper/2206.11990/) (2022)
- [EquiformerV2: Improved Equivariant Transformer for Scaling to Higher-Degree Representations](https://mlsanity.com/paper/2306.12059/) (2023)
- [Reducing SO(3) Convolutions to SO(2) for Efficient Equivariant GNNs](https://mlsanity.com/paper/2302.03655/) (2023)
- [Learning Smooth and Expressive Interatomic Potentials for Physical Property Prediction](https://mlsanity.com/paper/2502.12147/) (2025)
- [SE(3)-equivariant prediction of molecular wavefunctions and electronic densities](https://mlsanity.com/paper/2106.02347/) (2021)
- [Enabling Efficient Equivariant Operations in the Fourier Basis via Gaunt Tensor Products](https://mlsanity.com/paper/2401.10216/) (2024)
- [An Efficient Sparse Kernel Generator for O(3)-Equivariant Deep Networks](https://mlsanity.com/paper/2501.13986/) (2025)
- [The Price of Freedom: Exploring Expressivity and Runtime Tradeoffs in Equivariant Tensor Products](https://mlsanity.com/paper/2506.13523/) (2025)
- [E2Former: An Efficient and Equivariant Transformer with Linear-Scaling Tensor Products](https://mlsanity.com/paper/2501.19216/) (2025)
- [EquiformerV3: Scaling Efficient, Expressive, and General SE(3)-Equivariant Graph Attention Transformers](https://mlsanity.com/paper/2604.09130/) (2026)
- [A recipe for scalable attention-based MLIPs: unlocking long-range accuracy with all-to-all node attention](https://mlsanity.com/paper/2603.06567/) (2026)
- [Importance of equivariant and invariant symmetries for fluid flow modeling](https://mlsanity.com/paper/2307.05486/) (2023)
- [Learning Inter-Atomic Potentials without Explicit Equivariance](https://mlsanity.com/paper/2510.00027/) (2025)
- [Asymptotically Fast Clebsch-Gordan Tensor Products with Vector Spherical Harmonics](https://mlsanity.com/paper/2602.21466/) (2026)
- [E2Former-V2: On-the-Fly Equivariant Attention with Linear Activation Memory](https://mlsanity.com/paper/2601.16622/) (2026)
- [Complete and Efficient Covariants for 3D Point Configurations with Application to Learning Molecular Quantum Properties](https://mlsanity.com/paper/2409.02730/) (2024)

## Abstract (arXiv)

We present e3j, a fast Euclid-equivariance backend for geometric deep learning applications with JAX bindings for GPU and TPU. Leveraging both optimized CUDA and Pallas kernels and algorithmic improvements, the library achieves state-of-the-art throughput and runtime on both forward and backward paths. On a machine learning interatomic potential (MLIP) use case, it outperforms established backends, measuring up to 34% speed-up over cuEquivariance on water box NPT simulation using MACE, while remaining fully open source. E3j achieves over 80% efficiency over the H100 maximum memory bandwidth on tensor product operations, and in many cases more than doubles throughput of message passing convolutions forward compared to previously available backends. In addition, with the release of dedicated Pallas TPU kernel, e3j opens the possibility of large scale equivariant deep learning workloads on TPU architectures, which has so far been difficult to achieve. Our benchmarks show that e3j also achieves over 80% of a TPUv6e memory bandwidth, up to one order of magnitude more than e3nn-jax. The library is available on GitHub, PyPI and is released under an open source Apache 2.0 license.