A repository of building blocks in PyTorch 2.0 for E(3)/SE(3)-equivariant neural networks, built on top of e3nn:
- Equivariant Convolution:
e3tools.nn.Conv - Equivariant Multi-Layer Perceptrons (MLPs):
e3tools.nn.EquivariantMLP - Equivariant Layer Norm:
e3tools.nn.LayerNorm - Equivariant Activations:
e3tools.nn.Gateande3tools.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.
Install from PyPI:
pip install e3toolsor get the latest development version from GitHub:
pip install git+https://github.com/prescient-design/e3tools.gitWe 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:
python examples/train_qm9.py --model convWe see an approximate 2.5x improvement in training speed with torch.compile.