See the section below for example usage and 
our API page for more details.

PyTorch Examples 
Here's a CG tensor product implemented by e3nn:

gen = torch.Generator(device='cuda')

batch_size = 1000 X_ir, Y_ir, Z_ir = o3.Irreps("1x2e"), o3.Irreps("1x3e"), o3.Irreps("1x2e") X = torch.rand(batch_size, X_ir.dim, device='cuda', generator=gen) Y = torch.rand(batch_size, Y_ir.dim, device='cuda', generator=gen)

instructions=[(0, 0, 0, "uvu", True)]

tp_e3nn = o3.TensorProduct(X_ir, Y_ir, Z_ir, instructions, shared_weights=False, internal_weights=False).to('cuda') W = torch.rand(batch_size, tp_e3nn.weight_numel, device='cuda', generator=gen)

Z = tp_e3nn(X, Y, W) print(torch.norm(Z))

And here's the same tensor product using openequivariance. We require that your
tensors are stored on a CUDA device for this to work:

problem = oeq.TPProblem(X_ir, Y_ir, Z_ir, instructions, shared_weights=False, internal_weights=False) tp_fast = oeq.TensorProduct(problem, torch_op=True)

Z = tp_fast(X, Y, W) # Reuse X, Y, W from earlier print(torch.norm(Z))

Our interface for oeq.TPProblem is almost a strict superset of 
o3.TensorProduct (two key differences: we 
impose internal_weights=False and add support for multiple datatypes). 
You can pass e3nn Irreps instances directly or 
use oeq.Irreps, which is identical.

We recommend reading the e3nn documentation and API reference first, then using our kernels 
as drop-in replacements. We support most "uvu" and "uvw" tensor products; 
see this section for an up-to-date list of supported configurations.
Important: For many configurations, our kernels return results identical to
e3nn up to floating point roundoff (this includes all "uvu" problems with
multiplicity 1 for all irreps in the second input). For other configurations 
(e.g. any "uvw" connection modes), we return identical 
results up to a well-defined reordering of the weights relative to e3nn.
If you're executing tensor products as part of a message passing graph
neural network, we offer fused kernels that save both memory and compute time:

node_ct, nonzero_ct = 3, 4

Receiver, sender indices for message passing GNN

X = torch.rand(node_ct, X_ir.dim, device='cuda', generator=gen) Y = torch.rand(nonzero_ct, Y_ir.dim, device='cuda', generator=gen) W = torch.rand(nonzero_ct, problem.weight_numel, device='cuda', generator=gen)

tp_conv = oeq.TensorProductConv(problem, torch_op=True, deterministic=False) # Reuse problem from earlier Z = tp_conv.forward(X, Y, W, edge_index[0], edge_index[1]) # Z has shape [node_ct, z_ir.dim] print(torch.norm(Z))

If you can guarantee EdgeIndex is sorted by receiver index and supply the transpose
permutation, we can provide even greater speedup (and deterministic results) 
by avoiding atomics:

Now we can use the faster deterministic algorithm

Note: you don't need Pytorch geometric to use our kernels. When
deterministic=False, the sender and receiver indices can have
arbitrary order.
JAX Examples
After installation, use the library
as follows. Set OEQ_NOTORCH=1
in your environment to avoid the PyTorch import in
the regular openequivariance package.

os.environ["OEQ_NOTORCH"] = "1" import openequivariance as oeq

seed = 42 key = jax.random.PRNGKey(seed)

batch_size = 1000 X_ir, Y_ir, Z_ir = oeq.Irreps("1x2e"), oeq.Irreps("1x3e"), oeq.Irreps("1x2e") problem = oeq.TPProblem(X_ir, Y_ir, Z_ir, [(0, 0, 0, "uvu", True)], shared_weights=False, internal_weights=False)

node_ct, nonzero_ct = 3, 4 edge_index = jax.numpy.array( [ [0, 1, 1, 2], [1, 0, 2, 1], ], dtype=jax.numpy.int32, # NOTE: This int32, not int64 )

X = jax.random.uniform(key, shape=(node_ct, X_ir.dim), minval=0.0, maxval=1.0, dtype=jax.numpy.float32) Y = jax.random.uniform(key, shape=(nonzero_ct, Y_ir.dim), minval=0.0, maxval=1.0, dtype=jax.numpy.float32) W = jax.random.uniform(key, shape=(nonzero_ct, problem.weight_numel), minval=0.0, maxval=1.0, dtype=jax.numpy.float32)

tp_conv = oeq.jax.TensorProductConv(problem, deterministic=False) Z = tp_conv.forward( X, Y, W, edge_index[0], edge_index[1] ) print(jax.numpy.linalg.norm(Z))

Test JAX JIT

## Citation and Acknowledgements
If you find this code useful, please cite our paper: