ilyes319 / e3nn-jax

jax library for E3 Equivariant Neural Networks

e3nn-jax

🚀 44% faster than pytorch*

*Speed comparison done with a full model (MACE) during training (revMD-17) on a GPU (NVIDIA RTX A5000)

Documentation

💥 Warning 💥

Please always check the ChangeLog for breaking changes.

Installation

To install the latest released version:

pip install --upgrade e3nn-jax

To install the latest GitHub version:

pip install git+https://github.com/e3nn/e3nn-jax.git

To install from a local copy for development, we recommend creating a virtual enviroment:

python3 -m venv .venv
source .venv/bin/activate
pip install -r requirements.txt

To check that the tests are running:

pip install pytest
pytest e3nn_jax/_src/tensor_products_test.py

What is different from the PyTorch version?

• No more shared_weights and internal_weights in TensorProduct. Extensive use of jax.vmap instead (see example below)
• Support of python structure IrrepsArray that contains a contiguous version of the data and a list of jnp.ndarray for the data. This allows to avoid unnecessary jnp.concatenante followed by indexing to reverse the concatenation (even that jax.jit is probably able to unroll the concatenations)
• Support of None in the list of jnp.ndarray to avoid unnecessary computation with zeros (basically imposing 0 * x = 0, which is not simplified by default by jax because 0 * nan = nan)

Examples

The examples are moved in the documentation.

Citing

@misc{e3nn_paper,
    doi = {10.48550/ARXIV.2207.09453},
    url = {https://arxiv.org/abs/2207.09453},
    author = {Geiger, Mario and Smidt, Tess},
    keywords = {Machine Learning (cs.LG), Artificial Intelligence (cs.AI), Neural and Evolutionary Computing (cs.NE), FOS: Computer and information sciences, FOS: Computer and information sciences}, 
    title = {e3nn: Euclidean Neural Networks},
    publisher = {arXiv},
    year = {2022},
    copyright = {Creative Commons Attribution 4.0 International}
}