e3tools

Building Blocks for Equivariant Neural Networks in e3nn and PyTorch 2.0

https://github.com/prescient-design/e3tools

Science Score: 44.0%

This score indicates how likely this project is to be science-related based on various indicators:

  • CITATION.cff file
    Found CITATION.cff file
  • codemeta.json file
    Found codemeta.json file
  • .zenodo.json file
    Found .zenodo.json file
  • DOI references
  • Academic publication links
  • Academic email domains
  • Institutional organization owner
  • JOSS paper metadata
  • Scientific vocabulary similarity
    Low similarity (8.2%) to scientific vocabulary
Last synced: 6 months ago · JSON representation ·

Repository

Building Blocks for Equivariant Neural Networks in e3nn and PyTorch 2.0

Basic Info
Statistics
  • Stars: 15
  • Watchers: 3
  • Forks: 2
  • Open Issues: 0
  • Releases: 8
Created 12 months ago · Last pushed 7 months ago
Metadata Files
Readme License Citation

README.md

e3tools

A repository of building blocks in PyTorch 2.0 for E(3)/SE(3)-equivariant neural networks, built on top of e3nn: - Equivariant Linear Layers: e3tools.nn.Linear - Equivariant Convolution: e3tools.nn.Conv and e3tools.nn.SeparableConv - Equivariant Multi-Layer Perceptrons (MLPs): e3tools.nn.EquivariantMLP - Equivariant Layer Norm: e3tools.nn.LayerNorm - Equivariant Activations: e3tools.nn.Gate, e3tools.nn.GateWrapper and e3tools.nn.Gated - Separable Equivariant Tensor Products: e3tools.nn.SeparableTensorProduct - Extracting Irreps: e3tools.nn.ExtractIrreps - Self-Interactions: e3tools.nn.LinearSelfInteraction

All modules are compatible with torch.compile for JIT compilation.

Note that you may need to turn off the old torch JIT compiler for some e3nn modules, at the top of your script (example): python import e3nn e3nn.set_optimization_defaults(jit_script_fx=False)

Installation

Install from PyPI:

bash pip install e3tools

or get the latest development version from GitHub: bash pip install git+https://github.com/prescient-design/e3tools.git

Examples

We provide examples of a convolution-based and attention-based E(3)-equivariant message passing networks built with e3tools. We also provide an example training script on QM9: bash python examples/train_qm9.py --model conv

We see an approximate 2.5x improvement in training speed with torch.compile.

Owner

  • Name: Prescient Design
  • Login: prescient-design
  • Kind: organization
  • Email: prescient@gene.com

A Genentech Accelerator

Citation (CITATION.cff)

cff-version: 1.2.0
message: "If you use this software, please cite it as below."
authors:
  - family-names: Kleinhenz
    given-names: Joseph
    orcid: https://orcid.org/0000-0003-3670-0431
  - family-names: Daigavane
    given-names: Ameya
    orcid: https://orcid.org/0000-0002-5116-3075
title: "e3tools"
version: 0.1.1
date-released: 2025-04-04

GitHub Events

Total
  • Create event: 11
  • Commit comment event: 1
  • Issues event: 3
  • Release event: 8
  • Watch event: 15
  • Delete event: 1
  • Issue comment event: 6
  • Push event: 61
  • Public event: 1
  • Pull request event: 1
  • Fork event: 1
Last Year
  • Create event: 11
  • Commit comment event: 1
  • Issues event: 3
  • Release event: 8
  • Watch event: 15
  • Delete event: 1
  • Issue comment event: 6
  • Push event: 61
  • Public event: 1
  • Pull request event: 1
  • Fork event: 1

Dependencies

.github/workflows/pre-commit.yml actions
  • actions/checkout v3 composite
  • actions/setup-python v3 composite
  • pre-commit/action v3.0.1 composite
.github/workflows/pypi.yml actions
  • actions/checkout v4 composite
  • astral-sh/setup-uv v5 composite
.github/workflows/testpypi.yml actions
  • actions/checkout v4 composite
  • astral-sh/setup-uv v5 composite
.github/workflows/unit.yml actions
  • actions/checkout v4 composite
  • astral-sh/setup-uv v5 composite
pyproject.toml pypi
  • e3nn >=0.5.5
  • jaxtyping >=0.2.38
  • torch >=2.4.1
uv.lock pypi
  • anyio 4.8.0
  • backports-tarfile 1.2.0
  • certifi 2025.1.31
  • cffi 1.17.1
  • click 8.1.8
  • colorama 0.4.6
  • cryptography 44.0.2
  • distlib 0.3.9
  • e3nn 0.5.5
  • e3tools *
  • filelock 3.17.0
  • fsspec 2025.2.0
  • h11 0.14.0
  • hatch 1.14.0
  • hatchling 1.27.0
  • httpcore 1.0.7
  • httpx 0.28.1
  • hyperlink 21.0.0
  • idna 3.10
  • importlib-metadata 8.6.1
  • iniconfig 2.0.0
  • jaraco-classes 3.4.0
  • jaraco-context 6.0.1
  • jaraco-functools 4.1.0
  • jaxtyping 0.2.38
  • jeepney 0.9.0
  • jinja2 3.1.5
  • keyring 25.6.0
  • markdown-it-py 3.0.0
  • markupsafe 3.0.2
  • mdurl 0.1.2
  • more-itertools 10.6.0
  • mpmath 1.3.0
  • networkx 3.4.2
  • numpy 2.2.3
  • nvidia-cublas-cu12 12.1.3.1
  • nvidia-cuda-cupti-cu12 12.1.105
  • nvidia-cuda-nvrtc-cu12 12.1.105
  • nvidia-cuda-runtime-cu12 12.1.105
  • nvidia-cudnn-cu12 9.1.0.70
  • nvidia-cufft-cu12 11.0.2.54
  • nvidia-curand-cu12 10.3.2.106
  • nvidia-cusolver-cu12 11.4.5.107
  • nvidia-cusparse-cu12 12.1.0.106
  • nvidia-nccl-cu12 2.20.5
  • nvidia-nvjitlink-cu12 12.8.61
  • nvidia-nvtx-cu12 12.1.105
  • opt-einsum 3.4.0
  • opt-einsum-fx 0.1.4
  • packaging 24.2
  • pathspec 0.12.1
  • pexpect 4.9.0
  • platformdirs 4.3.6
  • pluggy 1.5.0
  • ptyprocess 0.7.0
  • pycparser 2.22
  • pygments 2.19.1
  • pytest 8.3.4
  • pywin32-ctypes 0.2.3
  • rich 13.9.4
  • scipy 1.15.2
  • secretstorage 3.3.3
  • shellingham 1.5.4
  • sniffio 1.3.1
  • sympy 1.13.3
  • tomli-w 1.2.0
  • tomlkit 0.13.2
  • torch 2.4.1
  • triton 3.0.0
  • trove-classifiers 2025.3.3.18
  • typing-extensions 4.12.2
  • userpath 1.9.2
  • uv 0.6.4
  • virtualenv 20.29.2
  • wadler-lindig 0.1.3
  • zipp 3.21.0
  • zstandard 0.23.0
.github/workflows/static.yml actions
  • actions/checkout v4 composite
  • actions/configure-pages v5 composite
  • actions/deploy-pages v4 composite
  • actions/upload-pages-artifact v3 composite
  • astral-sh/setup-uv v5 composite
examples/requirements.txt pypi
  • e3tools *
  • numpy *
  • torch *
  • torch_geometric *
  • tqdm *