E3x: E(3)-Equivariant Deep Learning Made Easy

E3x is a JAX library for constructing efficient E(3)-equivariant deep learning architectures built on top of Flax.

The goal is to provide common neural network building blocks for E(3)-equivariant architectures to make the development of models operating on three-dimensional data (point clouds, polygon meshes, etc.) easier.

This is not an officially supported Google product.

Installation

The easiest way to install E3x is via the Python Package Index (PyPI). Simply run

> python -m pip install --upgrade e3x

and you should be good to go.

Alternatively, you can clone this repository, enter the directory and run:

> python -m pip install .

If you are a developer, you might want to also install the optional development dependencies by running

> python -m pip install .[dev]

instead.

Documentation

Documentation for E3x, including usage examples and tutorials can be found here. For a more detailed overview over the mathematical theory behind E3x, please refer to this paper.

Citing E3x

If you find E3x useful and use it in your work, please cite:

@article{unke2024e3x,
  title={\texttt{E3x}: $\mathrm{E}(3)$-Equivariant Deep Learning Made Easy},
  author={Unke, Oliver T. and Maennel, Hartmut},
  journal={arXiv preprint arXiv:2401.07595},
  year={2024}
}