Install & Compatibility
Where this runs
tested against v0.21.0 · pip install
no network on importno background threads
Install × environment matrix
Each cell = how many times install + import succeeded across repeated harness runs. Partial = flaky.
glibc = Debian/Ubuntu slim · musl = Alpine Linux
muslpy 3.10–3.920 runs
build_error
glibcpy 3.10–3.920 runs
installs and imports cleanly · install 24.8s · import 5.536s · 647MB
679MB installed
● package 679MB
Code
Verified usage
Verified import paths — ran on the pinned version, not inferred.
Irreps
✓ from e3nn_jax import Irreps
spherical_harmonics
✓ from e3nn_jax import spherical_harmonics
rand_irreps
✓ from e3nn_jax import rand_irreps
flax.Linear
✓ from e3nn_jax.flax import Linear
✗ from e3nn_jax.linear import Linear
Flax modules are in the `e3nn_jax.flax` submodule, not directly under `e3nn_jax` or a generic `linear`.
This quickstart demonstrates how to define Irreducible Representations (Irreps), generate random equivariant features, and compute spherical harmonics from 3D positions, which are core operations in e3nn-jax. It uses `jax.random` for reproducibility and illustrates basic data generation and transformation.
import jax
import jax.numpy as jnp
from e3nn_jax import Irreps, rand_irreps, spherical_harmonics
key = jax.random.PRNGKey(0)
# Define input and output Irreps
irreps_in = Irreps("1x0e + 2x1o")
irreps_sh = Irreps("0e + 1o + 2e") # Spherical harmonics up to l=2
# Create random input features and positions
features = rand_irreps(key, irreps_in, (10,)).array # 10 samples
positions = jax.random.normal(key, (10, 3)) # 10 samples, 3D coordinates
# Compute spherical harmonics
sh = spherical_harmonics(irreps_sh, positions, normalize=True, normalization='component')
print(f"Input features irreps: {irreps_in}")
print(f"Input features shape: {features.shape}")
print(f"Spherical harmonics irreps: {irreps_sh}")
print(f"Spherical harmonics shape: {sh.shape}")
Debug
Known issues
breakingJAX version compatibility is critical. `e3nn-jax` closely tracks JAX's development. Upgrading JAX (especially `jaxlib`) often requires a corresponding `e3nn-jax` upgrade to avoid cryptic JIT errors, `AttributeError`s, or unexpected behavior.fixAlways check the `e3nn-jax` release notes and `pyproject.toml` for supported JAX versions. Update both `jax` and `e3nn-jax` simultaneously, ensuring `jaxlib` matches your hardware and `jax` version (e.g., `pip install --upgrade jax jaxlib e3nn-jax`).
affects: All versions, especially when crossing major JAX versions.
breakingThe `Irreps._repr_html_` method was removed in version 0.20.0, which means `Irreps` objects no longer render as rich HTML in Jupyter notebooks by default.fixIf you relied on rich HTML output in notebooks, you will need to manually format the `Irreps` object for display, e.g., using `str(irreps_object)` or custom display logic.
affects: >=0.20.0
breakingThe argument order for `ir_in` and `ir_out` in `e3nn_jax.flax.Linear` was swapped around version 0.17.0 to align with more natural tensor flow.fixReview the signature of `e3nn_jax.flax.Linear` in your code. Explicitly pass `ir_in=...` and `ir_out=...` to avoid ambiguity and ensure correctness.
affects: >=0.17.0
gotchaFor optimal performance (especially with GPU), `jaxlib` should often be installed manually *before* `e3nn-jax` to ensure the correct hardware-specific version is picked up. `pip install e3nn-jax` alone might install a CPU-only `jaxlib` or an incompatible version.fixFollow the official JAX installation guide (`https://github.com/google/jax#installation`) to install the correct `jax` and `jaxlib` version for your specific CPU/CUDA setup first, then install `e3nn-jax`.
affects: All versions
Errors
Common errors & fixes
TypeError: 'Irreps' object is not callable
Attempting to call an `Irreps` object as if it were a function. `Irreps` is a class used to define a representation, not a function to convert a string.
fixInstantiate the `Irreps` class by passing the string to its constructor, e.g., `my_irreps = Irreps("1x0e + 2x1o")` instead of `my_irreps = Irreps("...")()`. RuntimeError: JAX is not installed correctly. Please follow the instructions at https://github.com/google/jax#installation to install JAX.
This typically indicates a mismatch between your installed `jax`, `jaxlib`, and potentially CUDA versions (if on GPU), or a corrupted installation.
fixReinstall `jax` and `jaxlib` carefully, ensuring the `jaxlib` version matches your hardware environment (CPU/CUDA) as specified in the official JAX installation guides. Often, this means explicitly installing `jaxlib` for your CUDA version (e.g., `pip install --upgrade "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html`) before installing `e3nn-jax`.
ValueError: Arguments '...' and '...' have different shapes. Expected equal shapes, but got ... and ...
Shape mismatch errors frequently occur in JAX's JITted functions when input array shapes unexpectedly change, or `e3nn-jax` operations receive inputs that do not conform to the expected Irreps or batch dimensions.
fixDebug the shapes of all inputs to the `e3nn-jax` function causing the error. Ensure `Irreps` are correctly defined for both input and output, and that batch dimensions are consistent. For dynamic batching in JIT, consider using `jax.vmap`.
TypeError: rand_irreps() missing 1 required positional argument: 'irreps'
Incorrect usage of `e3nn_jax.rand_irreps` (or similar functions). For example, `rand_irreps` requires both a JAX `PRNGKey` and an `Irreps` object.
fixConsult the latest `e3nn-jax` documentation or source code for the correct function signature. Ensure all required positional and keyword arguments are provided with correct types. For `rand_irreps`, it should be `rand_irreps(key, irreps_object, shape_tuple)`.
Upgrade
Version history
0.21.0latest on PyPI · released Apr 1, 2026
Audit
Dependencies
jaxrequiredCore dependency for array manipulation and automatic differentiation.
jaxlibrequiredJAX's backend. Often requires careful manual installation for GPU support to ensure the correct hardware-specific version is used. `e3nn-jax` will install a CPU-compatible `jaxlib` by default if not present.