mlsanityData to 14 Oct 2026

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

arXiv 2609.35099 ↗First posted 28 Sep 2026cs.LG cs.DC cs.MScited by 0

Not on the map. Why?

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.

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

19 of 31 references.

References · 19

  1. 2018Tensor field networks: Rotation- and translation-equivariant neural networks for 3D point cloudscited by 746
  2. 2021E(3)-Equivariant Graph Neural Networks for Data-Efficient and Accurate Interatomic Potentialscited by 708
  3. 2022MACE: Higher Order Equivariant Message Passing Neural Networks for Fast and Accurate Force Fieldscited by 474
  4. 2022Equiformer: Equivariant Graph Attention Transformer for 3D Atomistic Graphscited by 264
  5. 2023EquiformerV2: Improved Equivariant Transformer for Scaling to Higher-Degree Representationscited by 204
  6. 2023Reducing SO(3) Convolutions to SO(2) for Efficient Equivariant GNNscited by 117
  7. 2025Learning Smooth and Expressive Interatomic Potentials for Physical Property Predictioncited by 84
  8. 2021SE(3)-equivariant prediction of molecular wavefunctions and electronic densitiescited by 67
  9. 2024Enabling Efficient Equivariant Operations in the Fourier Basis via Gaunt Tensor Productscited by 36
  10. 2025An Efficient Sparse Kernel Generator for O(3)-Equivariant Deep Networkscited by 22
  11. 2025The Price of Freedom: Exploring Expressivity and Runtime Tradeoffs in Equivariant Tensor Productscited by 15
  12. 2025E2Former: An Efficient and Equivariant Transformer with Linear-Scaling Tensor Productscited by 12
7 more
  1. 2026EquiformerV3: Scaling Efficient, Expressive, and General SE(3)-Equivariant Graph Attention Transformerscited by 11
  2. 2026A recipe for scalable attention-based MLIPs: unlocking long-range accuracy with all-to-all node attentioncited by 8
  3. 2023Importance of equivariant and invariant symmetries for fluid flow modelingcited by 5
  4. 2025Learning Inter-Atomic Potentials without Explicit Equivariancecited by 5
  5. 2026Asymptotically Fast Clebsch-Gordan Tensor Products with Vector Spherical Harmonicscited by 4
  6. 2026E2Former-V2: On-the-Fly Equivariant Attention with Linear Activation Memorycited by 3
  7. 2024Complete and Efficient Covariants for 3D Point Configurations with Application to Learning Molecular Quantum Propertiescited by 2

Built on by

Not cited yet.