Migrating from e3nn/PyTorch¶

The high-level namespaces deliberately resemble upstream e3nn:

e3nn/PyTorch	e3nn-mlx
`import torch`	`import mlx.core as mx`
`from e3nn import o3, nn`	`from e3nn_mlx import o3, nn`
`torch.Tensor`	`mx.array`
`loss.backward()`	`mx.value_and_grad(...)`
`torch.jit`	`mx.compile`
`torch_scatter.scatter`	`e3nn_mlx.scatter_sum`

Most constructors preserve the upstream names and representation arguments:

import mlx.core as mx
from e3nn_mlx import nn, o3

linear = o3.Linear("16x0e + 8x1o", "32x0e + 8x1o")
gate = nn.Gate(
    "16x0e", [mx.tanh],
    "8x0e", [mx.sigmoid],
    "8x1o",
)

Deliberate differences¶

MLX evaluates lazily. Include mx.eval(...) before reading results and when timing work.
mx.compile traces fixed array computation. Dynamic neighbor discovery is eager; construct edges once and use a model’s forward_with_edges method in compiled training or inference.
IrrepsArray is optional in e3nn_mlx.o3 and e3nn_mlx.nn, but the original flat API continues to use it for explicit representation checking.
Generated Metal kernels are selected automatically and remain subject to dtype, shape, and autodiff-transform boundaries.
TorchScript-specific APIs do not apply to MLX.

For an operation-by-operation statement of supported behavior, see the compatibility and numerical conventions and high-level API notes.