Caution
API is experimental and likely to break
> python -m pip install .
import e3nn_pt2
x1 = e3nn_pt2.IrrepsArray("1o", Array([1.0, 0.0, 0.0]).reshape(1,1,3).to(device='cuda'))
x2 = e3nn_pt2.IrrepsArray("1o", Array([0.0, 1.0, 0.0]).reshape(1,1,3).to(device='cuda'))
tp = e3nn_pt2.TensorProduct(x1.irreps, x2.irreps)
x3 = tp(x1.array, x2.array)
# Indexing the "1o" irrep.
np.testing.assert_allclose(x3.array[0,0,1:4].cpu().numpy(), Array([0.0, 0.0, 1 / 2**0.5]))
- e3nn
- e3nn-jax
- e3x