E3J: An Efficient and Open-Source Backend for Euclidean Equivariant Operations on GPU and TPU
Abstract
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.
Related
- Fault-Tolerant Logical Operations and Efficient State Preparation in Modular Quantum Architectures with Noisy Interfaces
- Realization of a Primary Irradiance Scale for Solar irradiance Measurements based on a Predictable Quantum Efficient Detector and a Quasi-Monochromatic Portable Tunable Source.
- SmoQyElPhQMC.jl: An open-source Julia package for efficient and scalable quantum Monte Carlo simulations of electron-phonon coupled models
- QwaveMPS: An efficient open-source Python package for simulating non-Markovian waveguide-QED using matrix product states
- Efficient Quantum Simulation of Variable-Coefficient Transport with Continuous Source Injection
BibTeX
@misc{peltre2026efficient,
title = {{E3J: An Efficient and Open-Source Backend for Euclidean Equivariant Operations on GPU and TPU}},
author = {Olivier Peltre and Armand Picard and Adrien Pichard and Miguel Bragança and Luca Giacomoni and Valentin Heyraud and Zachary Weller-Davies and Christoph Brunken and Jules Tilly},
howpublished = {arXiv},
year = {2026},
month = sep,
doi = {10.48550/arxiv.2609.35099},
eprint = {2609.35099},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2609.35099},
note = {CC BY 4.0},
}
RIS
TY - UNPB
TI - E3J: An Efficient and Open-Source Backend for Euclidean Equivariant Operations on GPU and TPU
AU - Olivier Peltre
AU - Armand Picard
AU - Adrien Pichard
AU - Miguel Bragança
AU - Luca Giacomoni
AU - Valentin Heyraud
AU - Zachary Weller-Davies
AU - Christoph Brunken
AU - Jules Tilly
PY - 2026
DA - 2026/09/28
PB - arXiv
DO - 10.48550/arxiv.2609.35099
UR - https://arxiv.org/abs/2609.35099
AB - 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.
ER -