Metadata-Version: 2.5
Name: e3j
Version: 0.1.0
Summary: Fast Euclid equivariant operations for GPU and TPU
Author-email: Olivier Peltre <o.peltre@instadeep.com>, Armand Picard <a.picard@instadeep.com>, Adrien Pichard <a.pichard@instadeep.com>, Miguel Bragança <m.braganca@instadeep.com>
License: Apache-2.0
License-File: LICENSE
Requires-Python: <3.15,>=3.11
Requires-Dist: e3nn-jax>=0.20.6
Requires-Dist: flax~=0.12.0
Requires-Dist: gitpython~=3.1.43
Requires-Dist: jax>=0.6.2
Requires-Dist: matplotlib>=3.9.0
Requires-Dist: numpy>=1.26.4
Requires-Dist: pandas~=2.2.2
Requires-Dist: pydantic>=2.13.3
Requires-Dist: setuptools<78.2,>=75.8
Provides-Extra: ops
Requires-Dist: e3j-ops==0.1.0; extra == 'ops'
Requires-Dist: pybind11>=2.13.6; extra == 'ops'
Description-Content-Type: text/markdown

# 🌐 e3j

A fast, full-featured and platform-agnostic equivariance backend for GPU and TPU!

This library can be used in place of [e3nn] and [e3x] to replace slow operations
in Machine Learned Interatomic Potentials (MLIPs), with carefully optimized
and open-source CUDA and Pallas kernels.

The equivariance backend of our MLIP library is [e3j] as of [mlip] 0.2.0.

[e3nn]: https://github.com/e3nn/e3nn-jax
[e3x]: https://github.com/google-research/e3x
[mlip]: https://github.com/instadeepai/mlip

## Installation

### Pulling from PyPI

The [`e3j`][e3j-pypi] package is available on PyPI.
It consists of a thin JAX-based Python API which can run on CPU, GPU and TPU,
supporting  Python versions from 3.11 to 3.14 included.

For efficiency on GPU, our CUDA binaries are bundled as the [`e3j_ops`][e3j-ops-pypi]
package on PyPI. The compatible version of the binaries should
be pulled by requiring the `"e3j[ops]"` extra:

```sh
# requirements.txt
e3j[ops] >= 0.1.0b0
jax[cuda13_local] ~= 0.8.0
```
See [JAX installation](https://docs.jax.dev/en/latest/installation.html) instructions for more information on JAX versions and their CUDA support. We recommend using a version of JAX above 0.7.0 and CUDA 13.

### Building from source

Our dependencies are managed with uv. After cloning the repository, you can
build from source by running run one of:

```sh
# Existing CUDA 13 install with `e3j_ops` kernels:
uv sync --group cuda13_local --extra ops
# Install CUDA 13 via pip and the `exp` group for benchmarks:
uv sync --group cuda13 --extra ops
```

The Python build internally relies on CMake, [scikit-build] and pybind11. You can also look at the [Makefile](Makefile) for alternate recipes to build kernels,
C++ tests and the Python bindings.

The [e3j_ops](lib/e3j_ops/README.md)
Python package only contains our CUDA binaries and bindings
to their associated XLA handlers. It is not meant to be used as standalone until its
ABI is reported stable.

## Features

`e3j` provides a platform-agnostic API for GPU and TPU:

- 🖥️ The same Python API on CPU, GPU and TPU, with a portable JAX fallback
when kernels don't apply
- 🔗 Interoperability made easy via plain `jax.Array` semantics
- 🌐 All equivariant building blocks: spherical harmonics (`Harmonics`),
tensor products (`TensorProduct`, `Bigotimes`), message-passing convolution
(`Convolution`) and learnable linear maps (`Linear`, `LinearIndexwise`)
as [flax.linen.Module] with weight initializations matching [e3nn]
- 🏎️ Fused CUDA kernels for GPU (tensor product, message-passing convolution,
scatter-add), shipped as the standalone [`e3j_ops`][e3j-ops-pypi] wheel
and dispatched through XLA-FFI
- 🧮 Fused Pallas Mosaic-TPU kernels, computing gather, tensor product,
scalar mixing and scatter in a single kernel
- 🔁 Custom VJP rules for every fused kernel, so they differentiate
under `jax.grad` like any other JAX primitive
- 🎛️ SPMD support for muliple-device MLIP training
- 🧱 Multiple memory layouts (leading channels, trailing channels, and a flat
[e3nn]-compatible layout) to trade coalescing off against interoperability
- 📐 Representation utilities: O(3) and SO(3) spaces, irreps filtering and
irrep filtering, permutations and generalized Clebsch-Gordan coefficients
- 🔌 Full coverage of the [e3nn] and [e3x] layers used by an MLIP, kernel-backed
or not, so an existing model can be ported over entirely — to train, simulate
and benchmark end to end, see [mlip]


<div align="center">
  <a href="https://instadeepai.github.io/e3j/animation.html"><img src="docs/animations/kernel_thumbnail.png" alt="Watch: E3J's Message Passing Convolution kernel on TPU" width="540"></a>
</div>

[flax.linen.Module]: https://flax-linen.readthedocs.io/en/latest/api_reference/flax.linen/module.html


## Contributing

Bug reports, questions and feature requests are welcome via [GitHub][e3j] issues and discussions.


## Citing
If you use [e3j] within your work, we kindly ask you to cite the following preprint:

```
@article{Peltre26-e3j,
    title   = {{E3J}: an Efficient and Open-Source Euclidean Equivariance Backend},
    author  = {Peltre, Olivier and Picard, Armand and Pichard, Adrien and Giacomoni, Luca and Braganca, Miguel and Heyraud, Valentin and Brunken, Christoph and Tilly, Jules},
    journal = {preprint},
    year    = {2026},
    url     = {(preprint)}
  }
}
```
[e3j]: https://github.com/instadeepai/e3j
[e3j-pypi]: https://pypi.org/projects/e3j
[e3j-ops-pypi]: https://pypi.org/projects/e3j_ops
[scikit-build]: https://scikit-build.readthedocs.io/en/latest/
