e3nn-mlx¶

e3nn-mlx is an MLX-native implementation of Euclidean neural-network building blocks for Apple silicon. It keeps the familiar e3nn.o3 and e3nn.nn organization while using MLX arrays, compilation, automatic differentiation, and generated Metal kernels.

Important

This project is an independent MLX port and is currently beta software. It is not full drop-in binary replacement for PyTorch e3nn, although high-level compatibility api exists. Consult the compatibility and numerical conventions before depending on an operation that is not covered by the public API below.

Real spherical harmonics through degree five generated by e3nn-mlx

This animation contains all 36 real spherical-harmonic components through degree \(l=5\). The surfaces are evaluated by e3nn_mlx.o3.spherical_harmonics(); the visualization converts the modern torch-compatible basis into the legacy display convention used by the original e3nn animation.

This animation can be rebuilt with

.venv/bin/python tutorials/spherical_harmonics_animation.py \
    --output docs/_static/sphharm_mlx.gif --transparent

--transparent leaves the page showing through instead of painting the figure white, so the animation reads correctly on both light and dark documentation themes. Drop the flag for a white background.

Where to start¶

New to equivariance? Read Irreducible representations, then work through the first equivariant operation.
Prefer notebook-based introductions? Explore the adapted historical e3nn tutorials on tensor types, spherical-tensor operations, and invariant atomic descriptors.
Coming from e3nn/PyTorch? Start with the migration guide.
Building a model? See tensor products, the convolution example, and point models.
Optimizing Apple-silicon workloads? Read performance and differentiation.

A minimal example¶

import mlx.core as mx
from e3nn_mlx import o3

irreps = o3.Irreps("4x0e + 2x1o")
linear = o3.Linear(irreps, "8x0e + 4x1o")
x = mx.random.normal((128, irreps.dim))
y = linear(x)
mx.eval(y)

assert y.shape == (128, linear.irreps_out.dim)

The scalar (0e) and vector (1o) channels are mixed only in ways compatible with their transformation laws. The same operation accepts an IrrepsArray when representation-aware values are useful.