atomicarchitects / e3nn-pt2

e3nn + PyTorch 2.0

e3nn-pt2

Caution

API is experimental and likely to break

e3nn + PyTorch 2.0

Installation

> python -m pip install .

Amuze-Bouche

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]))

Acknowledgement

• e3nn
• e3nn-jax
• e3x