Metadata-Version: 2.4
Name: e3nn-mlx
Version: 0.1.1
Summary: Euclidean neural networks for MLX
Author: Rostislav Fedorov
License-Expression: MIT
Project-URL: Homepage, https://github.com/lamalab-org/e3nn_mlx
Project-URL: Changelog, https://github.com/lamalab-org/e3nn_mlx/blob/main/CHANGELOG.md
Project-URL: Documentation, https://lamalab-org.github.io/e3nn_mlx/
Project-URL: Issues, https://github.com/lamalab-org/e3nn_mlx/issues
Project-URL: Repository, https://github.com/lamalab-org/e3nn_mlx.git
Keywords: equivariance,e3nn,machine-learning,mlx,molecular-modeling
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Science/Research
Classifier: Operating System :: MacOS
Classifier: Operating System :: POSIX :: Linux
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3 :: Only
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Programming Language :: Python :: 3.14
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: Scientific/Engineering :: Physics
Requires-Python: >=3.11
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: mlx<0.33,>=0.31; sys_platform == "darwin" and platform_machine == "arm64"
Requires-Dist: mlx[cpu]<0.33,>=0.31; sys_platform == "linux"
Requires-Dist: numpy<3,>=2
Provides-Extra: test
Requires-Dist: pytest<10,>=9; extra == "test"
Requires-Dist: numpy<3,>=2.3; extra == "test"
Provides-Extra: reference
Requires-Dist: e3nn==0.5.8; extra == "reference"
Requires-Dist: numpy==2.3.1; extra == "reference"
Requires-Dist: torch==2.7.1; extra == "reference"
Provides-Extra: release
Requires-Dist: build<2,>=1.2; extra == "release"
Requires-Dist: check-wheel-contents<1,>=0.6; extra == "release"
Requires-Dist: twine<7,>=6; extra == "release"
Provides-Extra: docs
Requires-Dist: furo>=2024.8.6; extra == "docs"
Requires-Dist: myst-parser<5,>=4; extra == "docs"
Requires-Dist: sphinx<10,>=8; extra == "docs"
Requires-Dist: sphinx-copybutton<1,>=0.5; extra == "docs"
Provides-Extra: tutorials
Requires-Dist: ase<4,>=3.26; extra == "tutorials"
Requires-Dist: ipywidgets<9,>=8; extra == "tutorials"
Requires-Dist: matplotlib<4,>=3.9; extra == "tutorials"
Requires-Dist: plotly<7,>=6; extra == "tutorials"
Requires-Dist: scipy<2,>=1.14; extra == "tutorials"
Dynamic: license-file

# e3nn_mlx

Incremental refactor of e3nn into a backend-neutral core plus an MLX-native runtime.

The current implementation provides:

- backend-agnostic O(3) metadata in `e3nn_core`
- e3nn-compatible real Wigner-3j and Clebsch--Gordan coefficients
- batched real Wigner matrices and O(3) parity transforms through arbitrary `l`
- spherical harmonics generated recursively from the canonical Wigner basis
- weighted and unweighted MLX tensor products and standard wrappers
- lightweight MLX Linear, Gate, Norm, and reduction operations
- graph/radial utilities, the gated `gate_points_2102` network, and the complete
  modular `v2106` point-model family
- pinned upstream e3nn numerical fixtures and Apple-Silicon CI coverage

Numerical compatibility is guaranteed through `l=6`; reference and structural
tests exercise selected operations through `l=8`.  See
[`docs/COMPATIBILITY.md`](https://github.com/lamalab-org/e3nn_mlx/blob/main/docs/COMPATIBILITY.md) for conventions, shapes,
normalization, and the release gate.

## Installation

Install the published package with:

```bash
python -m pip install e3nn-mlx
```

Python 3.11 or newer is required. Apple Silicon uses MLX's Metal backend;
Linux installs the MLX CPU backend. Generated Metal kernels automatically fall
back to general MLX operations when Metal is unavailable.

## e3nn-style API

New code can use the familiar upstream namespace layout with ordinary MLX
arrays:

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

linear = o3.Linear("16x0e + 16x1o", "32x0e + 16x1o")
x = mx.random.normal((128, linear.irreps_in.dim))
y = linear(x)
```

`e3nn_mlx.o3`, `e3nn_mlx.nn`, `e3nn_mlx.math`, and
`e3nn_mlx.nn.models.v2106` mirror the implemented PyTorch/e3nn organization.
Raw inputs return raw MLX arrays; `IrrepsArray` inputs preserve typed outputs.
The original flat API remains backward-compatible. See
[`docs/HIGH_LEVEL_API.md`](https://github.com/lamalab-org/e3nn_mlx/blob/main/docs/HIGH_LEVEL_API.md) for migration examples,
model imports, compilation, deliberate differences, and no-overhead evidence.

## Development

Install the test dependencies and run the suite on Apple Silicon:

```bash
python -m pip install -e '.[test]'
python -m pytest
```

MLX tests may skip when Metal is unavailable locally.  The required macOS CI job
sets `E3NN_MLX_REQUIRE_RUNTIME=1`, which turns such skips into failures.

## Gate-points model

`e3nn_mlx.models.gate_points_2102.Network` ports the upstream gated point-cloud
network. It accepts a dictionary containing `pos`, `x`, `z`, and an optional
integer `batch` array and returns an `IrrepsArray`. `GatePointsNetwork` and
`GatePointsConvolution` are also exported from `e3nn_mlx`.

Neighbor discovery with `radius_graph` is eager because MLX 0.31 does not have
a device-side dynamic nonzero operation. For compiled training or inference,
construct the topology once and call `Network.forward_with_edges`; the radial
embedding, spherical harmonics, message passing, gating, and graph reduction
then remain inside the compiled MLX graph.

The complete modular June 2021 family is available from
`e3nn_mlx.models.v2106`. It includes `Convolution`, `MessagePassing`,
`SimpleNetwork`, and `NetworkForAGraphWithAttributes`; both networks support
eager dictionary input and compiled fixed-edge execution. See
[`docs/V2106_POINT_MODELS.md`](https://github.com/lamalab-org/e3nn_mlx/blob/main/docs/V2106_POINT_MODELS.md) for their APIs,
compilation boundary, upstream semantics, and verification coverage.

## Performance evaluation

The compact cross-framework harness in
[`evals/`](https://github.com/lamalab-org/e3nn_mlx/blob/main/evals/README.md)
always compares three isolated workers: upstream e3nn on Torch CPU using all
available CPU threads, compiled general MLX, and compiled MLX with generated
kernels. It covers spherical harmonics, tensor products, Linear, and scatter
aggregation in forward and training modes. Smoke and full presets produce
synchronized JSON, CSV, SVG, and HTML reports. Separate seeded randomized
harnesses qualify numerical outputs and VJPs against upstream e3nn.

### JVP and generated-kernel boundary

A Jacobian-vector product (JVP) propagates a chosen input perturbation through
a function without constructing its complete Jacobian. Most users do not call
JVP directly: ordinary inference and standard training—including MACE energy,
force, and parameter-gradient training—use forward evaluation and reverse-mode
gradients.

JVP is useful for more specialized atomistic workflows, including:

- Hessian-vector products and directional force-constant calculations;
- phonon, vibrational-response, and stability algorithms that propagate a
  displacement direction;
- mixed position/parameter response calculations;
- tangent dynamics, sensitivity analysis, and forward-mode Jacobian APIs;
- debugging equivariance by differentiating along an infinitesimal rotation.

MLX 0.31 cannot currently apply JVP directly to a `CustomKernel` primitive.
For these workflows, select the fully differentiable MLX implementation:

```python
import e3nn_mlx

# Spherical harmonics and fixed-topology graph reduction
y = e3nn_mlx.spherical_harmonics(
    degrees, vectors, use_custom_kernel=False
)
summed = e3nn_mlx.scatter_sum(
    messages,
    edge_dst,
    num_nodes,
    use_custom_kernel=False,
    jvp_safe=True,
)

# TensorProduct: use this callable inside mx.jvp
y = tensor_product.differentiable_arrays(left, right, weights)
```

Because MLX 0.31's indexed-add primitive does not implement JVP, the scatter
fallback requires eager, fixed indices. It uses a sparse sorted prefix sum with
linear memory rather than a dense node-by-edge incidence matrix. Ordinary
scatter calls keep the faster indexed-add implementation.

This changes execution strategy, not mathematical conventions or accuracy.
Reverse-mode gradients and reverse-over-reverse second derivatives remain
supported by the generated kernels. See the
[`evals/` guide](https://github.com/lamalab-org/e3nn_mlx/blob/main/evals/README.md)
for the three-way performance comparison and randomized parity commands.
