Skip to main content

Equivariant convolutional neural networks for the group E(3) of 3 dimensional rotations, translations, and mirrors.

Project description

e3nn-jax Coverage Status

Documentation Documentation Status

:boom: Warning :boom:

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}
}

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

e3nn_jax-0.12.0.tar.gz (88.9 kB view hashes)

Uploaded Source

Built Distribution

e3nn_jax-0.12.0-py3-none-any.whl (108.8 kB view hashes)

Uploaded Python 3

Supported by

AWS AWS Cloud computing and Security Sponsor Datadog Datadog Monitoring Fastly Fastly CDN Google Google Download Analytics Microsoft Microsoft PSF Sponsor Pingdom Pingdom Monitoring Sentry Sentry Error logging StatusPage StatusPage Status page