e3tools

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

Installation

Install from PyPI:

pip install e3tools

or get the latest development version from GitHub:

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:

python examples/train_qm9.py --model conv

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