instadeepai/e3j

★ 13PythonAudience · researcherComplexity · 5/5Setup · hard

Mindmap

mindmap
  root((e3j))
    What It Does
      Equivariant operations
      JAX library
      Molecular AI backend
    Tech Stack
      Python
      JAX
      CUDA kernels
    Use Cases
      Interatomic potentials
      Molecular dynamics
      Neural force fields
    Features
      GPU acceleration
      CPU fallback
      E3nn replacement
    Audience
      ML researchers
      Computational chemists

mindmap root((e3j)) What It Does Equivariant operations JAX library Molecular AI backend Tech Stack Python JAX CUDA kernels Use Cases Interatomic potentials Molecular dynamics Neural force fields Features GPU acceleration CPU fallback E3nn replacement Audience ML researchers Computational chemists

Click or tap to explore — scroll the page freely

Things people build with this

USE CASE 1

Train machine-learned interatomic potential models for molecular dynamics simulations using fast equivariant operations on GPU.

USE CASE 2

Replace e3nn or e3x equivariant layers in an existing JAX model with e3j to get faster GPU throughput on the same hardware.

USE CASE 3

Run equivariant neural network operations on CPU using the Python-only API without installing CUDA binaries.

Tech stack

PythonJAXCUDA

Getting it running

Difficulty · hard Time to first run · 1h+

GPU acceleration requires a separate CUDA binary package, a CUDA-compatible GPU is needed for full performance, though a CPU-only mode is available.

In plain English

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.

Copy-paste prompts

Prompt 1

I'm building a machine-learned interatomic potential in JAX and want to use e3j for equivariant operations. Show me how to import e3j and apply a basic tensor product operation to atom feature embeddings.

Prompt 2

How do I install the e3j CUDA binary package alongside the Python API to enable GPU acceleration? Walk me through the setup steps including what version of JAX is required.

Prompt 3

What is the practical difference between e3j and e3nn for equivariant neural networks? When should I switch to e3j for a molecular simulation project?

Prompt 4

Show me how to verify that an e3j operation is rotationally equivariant by rotating input atom positions and checking the output transforms correctly.

Open on GitHub → Explain another repo

← instadeepai on gitmyhub — every repo by this author, as a profile.

Verify against the repo before relying on details.