e3j
Euclid-equivariant operations for JAX
A high-performance JAX library of equivariant mathematical operations for training AI models that simulate how atoms and molecules behave, with custom CUDA kernels that run faster than the earlier e3nn and e3x alternatives.
e3j is a Python library built for researchers who train AI models to simulate how atoms and molecules behave. It provides a specific kind of mathematical building block called Euclidean equivariant operations. In plain terms, equivariance means that if you rotate or mirror a molecule in space and then run it through the model, the output rotates or mirrors in the same predictable way, rather than producing a completely different result. Building models with this property requires specialized mathematical operations that are more complex than the arithmetic used in standard deep learning.
The library is designed to run on JAX, a numerical computing framework developed by Google that is popular in scientific machine learning for its ability to run efficiently on GPUs and TPUs. The main purpose of e3j is to replace the equivalent operations in two earlier libraries called e3nn and e3x, which handle the same mathematics but more slowly. The performance improvement comes from custom GPU code written at a lower level than Python, using CUDA kernels that the authors optimized specifically for this kind of computation.
The primary use case described in the README is training machine-learned interatomic potentials, which are AI models that predict the energy and forces in a system of atoms. These models are used in chemistry and materials science research to simulate molecular dynamics without running expensive physics calculations from scratch. The library is already the equivariance backend for InstaDeep's own interatomic potential library as of version 0.2.0 of that project.
The library was in pre-release at the time of the README, with version 0.1.0 planned for early June 2026. The GPU performance features require installing an additional package containing the compiled CUDA binaries. A CPU-only version with the Python API works without those extra binaries and can run on standard hardware. The project is open source and the source code for both the Python layer and the low-level GPU code is included in the repository.
Where it fits
- Train machine-learned interatomic potential models for molecular dynamics simulations using fast equivariant operations on GPU.
- Replace e3nn or e3x equivariant layers in an existing JAX model with e3j to get faster GPU throughput on the same hardware.
- Run equivariant neural network operations on CPU using the Python-only API without installing CUDA binaries.