Custom Linear CUDA kernel

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

[mariogeiger]

import cuequivariance as cue
import numpy as np

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

e0 = cue.descriptors.fully_connected_tensor_product(
    cue.Irreps(cue.O3, "3x0e"),
    cue.Irreps(cue.O3, "12x0e + 32x1o"),
    cue.Irreps(cue.O3, "32x0e + 48x1o"),
)

print(f"old ordering: {e0.d.operands[0].segments}")
print(f"new ordering: 3x ({e.d.operands[0].segments})")

# convert old weights (w0) to new weights (w)
w0 = np.random.randn(e0.inputs[0].dim)
num_elements = e0.inputs[1].dim
w = []
for s in e0.d.operands[0].segment_slices():
    w.append(w0[s].reshape(num_elements, -1))
w = np.concatenate(w, axis=1).flatten()