Indexed Weights Linear

[mariogeiger]

import torch
import cuequivariance_torch as cuet
import cuequivariance as cue

e = cue.descriptors.linear(
    cue.Irreps(cue.O3, "12x0e + 32x1o"),
    cue.Irreps(cue.O3, "32x0e + 48x1o"),
).flatten_modes("i")
d = e.d

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
x0 = torch.randn(3, d.operands[0].size, device=device)
x1 = torch.randn(100, d.operands[1].size, device=device)
i = torch.randint(0, 3, (100,), device=device, dtype=torch.int32)

m = cuet.EquivariantTensorProduct(
    e,
    device=device,
    math_dtype=torch.float32,
    layout=cue.ir_mul,
    index_first_input=True,
)
y2 = m(x0, x1, indices=i)