Metadata-Version: 2.4
Name: e3tools
Version: 0.1.2
Summary: A repository of building blocks in PyTorch for E(3)/SE(3)-equivariant neural networks
Author-email: Joseph Kleinhenz <kleinhej@gene.com>, Ameya Daigavane <ameyad@mit.edu>
License-File: LICENSE
Requires-Python: >=3.10
Requires-Dist: e3nn>=0.5.5
Requires-Dist: jaxtyping>=0.2.38
Requires-Dist: setuptools>=78.1.0
Requires-Dist: torch>=2.5.1
Provides-Extra: dev
Requires-Dist: hatch>=1.14.0; extra == 'dev'
Requires-Dist: ipython>=8.34.0; extra == 'dev'
Requires-Dist: pytest>=8.3.4; extra == 'dev'
Requires-Dist: ruff; extra == 'dev'
Requires-Dist: sphinx; extra == 'dev'
Description-Content-Type: text/markdown

# e3tools

A repository of building blocks in PyTorch 2.0 for E(3)/SE(3)-equivariant neural networks, built on top of [e3nn](https://github.com/e3nn/e3nn):
- Equivariant Convolution: [`e3tools.nn.Conv`](https://github.com/prescient-design/e3tools/blob/main/src/e3tools/nn/_conv.py#L16)
- Equivariant Multi-Layer Perceptrons (MLPs): [`e3tools.nn.EquivariantMLP`](https://github.com/prescient-design/e3tools/blob/main/src/e3tools/nn/_mlp.py#L86)
- Equivariant Layer Norm: [`e3tools.nn.LayerNorm`](https://github.com/prescient-design/e3tools/blob/main/src/e3tools/nn/_layer_norm.py#L9)
- Equivariant Activations: [`e3tools.nn.Gate`](https://github.com/prescient-design/e3tools/blob/main/src/e3tools/nn/_gate.py#L10) and [`e3tools.nn.Gated`](https://github.com/prescient-design/e3tools/blob/main/src/e3tools/nn/_gate.py#L68)
- Separable Equivariant Tensor Products: [`e3tools.nn.SeparableTensorProduct`](https://github.com/prescient-design/e3tools/blob/main/src/e3tools/nn/_tensor_product.py#L8)
- Extracting Irreps: [`e3tools.nn.ExtractIrreps`](https://github.com/prescient-design/e3tools/blob/main/src/e3tools/nn/_extract_irreps.py#L5)
- Self-Interactions: [`e3tools.nn.LinearSelfInteraction`](https://github.com/prescient-design/e3tools/blob/main/src/e3tools/nn/_interaction.py#L5)

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

## 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](https://github.com/prescient-design/e3tools/blob/main/examples/models/conv.py) and [attention-based](https://github.com/prescient-design/e3tools/blob/main/examples/models/transformer.py) E(3)-equivariant message passing networks built with `e3tools`. We also provide an [example training script on QM9](https://github.com/prescient-design/e3tools/blob/main/examples/train_qm9.py):
```bash
python examples/train_qm9.py --model conv
```

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