Caution
API is experimental and likely to break
> python -m pip install .import tinye3nn
# Create Irreps as per e3nn convention (currently not supporting other input formats except strings)
x_irreps = tinye3nn.so3.Irreps("32x0e + 32x1o + 32x2e")
y_irreps = tinye3nn.so3.Irreps("0e + 1o + 2e")
# Initialize arrays and move them to the GPU
x = x_irreps.randn()
y = y_irreps.randn()
# Initialize TP + Linear module with the option to specify the device
mod = tinye3nn.nn.TensorProductLinear(x_irreps, y_irreps)
# Run the model :)
out = mod(x, y)- e3nn
- e3nn-jax
- e3x
- tinygrad