OpenEquivariance

0.6.6 verified Sat May 09 auth: no python

A fast GPU JIT kernel generator for the Clebsch-Gordon Tensor Product, currently at version 0.6.6. Released about monthly. Supports PyTorch and JAX backends.

pip install openequivariance

Common errors

error AttributeError: module 'openequivariance' has no attribute 'oe' ↓

cause Incorrect import: `import openequivariance` does not expose `oe`. Need to use `from openequivariance import oe`.

fix

Replace import openequivariance with from openequivariance import oe.

error RuntimeError: Unsupported dtype ↓

cause Did not specify `dtype` argument or passed a non-supported type like torch.int.

fix

Pass dtype=torch.float32 (or torch.float64) in the TPProblem constructor.

Warnings

breaking Import path changed: in early versions (pre-0.6), the module was importable as `import openequivariance`. Since 0.6, use `from openequivariance import oe`. ↓

fix Change `import openequivariance` to `from openequivariance import oe`.

gotcha Irrep specifications changed from tuples of (multiplicity, (l, p)) to (multiplicity, (l, p)) where p is parity (1 or -1). Old code using integers may break. ↓

fix Ensure irreps are specified as list of (mult, (l, parity)) tuples.

deprecated The `TPProblem` constructor's `dtype` argument requires explicit float types; implicit type promotion may be removed in future. ↓

fix Always pass `dtype=torch.float32` or `jnp.float32` explicitly.

Imports

oe
wrong
```
import openequivariance
```
correct
```
from openequivariance import oe
```
Top-level module is 'oe', not package name. Use 'from openequivariance import oe'.
TPProblem
```
from openequivariance import TPProblem
```

Quickstart

Minimal tensor product using Torch backend on CPU.

import torch
from openequivariance import oe

irreps_in1 = [(1, (1, 1)), (1, (1, 2)), (1, (2, 1)), (1, (2, 2))]
irreps_in2 = [(1, (1, 1)), (1, (1, 2)), (1, (2, 1)), (1, (2, 2))]
irreps_out = [(1, (1, 1)), (1, (1, 2)), (1, (2, 1)), (1, (2, 2))]
instruction = [(0, 0, 0, "uvu", True), (1, 1, 1, "uvu", True)]

# Create TP problem
tp = TPProblem(irreps_in1, irreps_in2, irreps_out, instruction, dtype=torch.float32, device='cpu')

# Input tensors
N = 8
x1 = torch.randn(N, 16)
x2 = torch.randn(N, 16)

# Forward pass
result = tp(x1, x2)
print(result.shape)