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

We present e3j, a fast Euclid-equivariance backend for geometric deep learning applications with JAX bindings for GPU and TPU.

By Peltre, Picard, Pichard +6Sep 28, 2026

Score████░░░░░░3.6

Key numbers

34% speed-up over cuEquivariance on
80% efficiency over the H100
80% of a TPUv6e memory

VerdictWorth a reader's time today.

Read the original PDF

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.

Olivier Peltre, Armand Picard, Adrien Pichard, Miguel Bragança, Luca Giacomoni, Valentin Heyraud, Zachary Weller-Davies, Christoph Brunken, Jules Tilly

The editor's rubric

Heuristic review

Dimension	Level	Weight	What that level means
Leverage	███░░ 3	24%	A method or resource many groups across the field will adopt within a year.
Magnitude	████░ 4	18%	A qualitative jump: a capability or regime that did not exist before.
Evidence	███░░ 3	14%	Solid: multiple benchmarks or cohorts, ablations, fair baselines, released code or data.
Novelty	██░░░ 2	16%	A new combination of known ideas.
Trajectory	███░░ 3	18%	A clear path to scale.
Stakes	██░░░ 2	10%	Benefits a professional community (practitioners, clinicians, engineers).

Editor’s rationale

Heuristic triage from title and abstract text only, not a reading of the paper. Cues found: method (we propose); gains (orders of magnitude, state of the art, outperforms); verification (code released); scale (scalable, orders of magnitude, efficient).

How the score was computed

rank-2026-09-29

Score████░░░░░░3.6

Score = 10 × (65% × adjusted merit / 10 + 25% × attention + 10% × freshness)

Merit: 5.8 / 10; Weighted rubric, evidence-gated.
Adjusted merit: 4.8 / 10; Shrunk toward the desk prior by editor confidence (44%).
Attention: 0%; Citations, upvotes, points, mentions.
Freshness: 49%; Half-life decay since publication.

Citations0 (reference 15, via semantic-scholar, Oct 1, 2026, 02:16 UTC)
Influential citations0 (reference 3, via semantic-scholar, Oct 1, 2026, 02:16 UTC)

The record

Reviewed by heuristic-v2 on Sep 29, 2026, 23:53 UTC. Paper type: method.
Categories: cs.LG, cs.DC, cs.MS, physics.comp-ph
BRIEF, No.2 in the Physics edition of September 30, 2026.
TOP, No.6 in the Front page edition of September 29, 2026.
TOP, No.3 in the Physics edition of September 29, 2026.