Are multiplicities > 1 in irreps_in2 supported in uvw TensorProduct?

[mjrs33]

Is support for irreps_in2 multiplicity > 1 in uvw mode (FullyConnectedTensorProduct) currently incomplete, or am I overlooking a usage detail?
Any guidance (or a confirmation that this is a bug) would be greatly appreciated.

import e3nn.o3 as o3
import openequivariance as oeq
import torch

def create_tps(irreps_in1, irreps_in2, irreps_out):
    instructions = [
        (i_1, i_2, i_out, "uvw", True, 1.0)
        for i_1, (_, ir_1) in enumerate(irreps_in1)
        for i_2, (_, ir_2) in enumerate(irreps_in2)
        for i_out, (_, ir_out) in enumerate(irreps_out)
        if ir_out in ir_1 * ir_2
    ]
    kwargs = {
        "irreps_in1": irreps_in1,
        "irreps_in2": irreps_in2,
        "irreps_out": irreps_out,
        "instructions": instructions,
        "internal_weights": False,
        "shared_weights": False,
        "path_normalization": "element",
        "irrep_normalization": "component"
    }
    tp_e3nn = o3.TensorProduct(**kwargs).cuda()
    problem = oeq.TPProblem(**kwargs)
    tp_oeq = oeq.TensorProduct(problem)
    
    x1 = irreps_in1.randn(1, -1, device="cuda")
    x2 = irreps_in2.randn(1, -1, device="cuda")
    w = torch.randn(1, tp_e3nn.weight_numel, device="cuda")
    return tp_e3nn, tp_oeq, x1, x2, w


# multiplicity=1
torch.manual_seed(0)
irreps_in1 = o3.Irreps("8x0e+4x1e")
irreps_in2 = o3.Irreps("1x0e")
irreps_out = o3.Irreps("8x0e+4x1e")
tp_e3nn, tp_oeq, x1, x2, w = create_tps(irreps_in1, irreps_in2, irreps_out)
out_e3nn = tp_e3nn(x1, x2, w)
out_oeq = tp_oeq(x1, x2, w)
print(torch.abs(out_e3nn - out_oeq).max().item())
# -> 5.960464477539063e-08

# multiplicity=2
irreps_in1 = o3.Irreps("8x0e+4x1e")
irreps_in2 = o3.Irreps("2x0e")
irreps_out = o3.Irreps("8x0e+4x1e")
tp_e3nn, tp_oeq, x1, x2, w = create_tps(irreps_in1, irreps_in2, irreps_out)
out_e3nn = tp_e3nn(x1, x2, w)
out_oeq = tp_oeq(x1, x2, w)
print(torch.abs(out_e3nn - out_oeq).max().item())
# -> 0.9126971960067749

Thanks in advance for your help!

[vbharadwaj-bk]

Ah, you need to call a weight reordering function that we provide to get this to match (no reordering needed for the first case). We'll modify your example later today to demonstrate the usage and resolve, it's not well-documented at the moment.

[vbharadwaj-bk]

Here's the reworked example that succeeds; note the reordering logic, which is only required if multiplicity is >1 for L2 or if the irrep dimensions exceed 32 for UVW tensor products. If attempting to access the reordering function yields "None", it's not required.

That said, the reordering logic could be made much simpler (with support to operate directly on torch tensors). We'll get this patched up and documented.

Edited because this is outdated; see the comment below.

[vbharadwaj-bk]

Update: we just merged a PR that exposes PyTorch versions of the weight reordering functions. Try cloning the latest version of the Github repo, executing an editable install, and running this modified verison of your example. We should be able to close this.

import e3nn.o3 as o3
import openequivariance as oeq
import torch

def create_tps(irreps_in1, irreps_in2, irreps_out):
    instructions = [
        (i_1, i_2, i_out, "uvw", True, 1.0)
        for i_1, (_, ir_1) in enumerate(irreps_in1)
        for i_2, (_, ir_2) in enumerate(irreps_in2)
        for i_out, (_, ir_out) in enumerate(irreps_out)
        if ir_out in ir_1 * ir_2
    ]
    kwargs = {
        "irreps_in1": irreps_in1,
        "irreps_in2": irreps_in2,
        "irreps_out": irreps_out,
        "instructions": instructions,
        "internal_weights": False,
        "shared_weights": False,
        "path_normalization": "element",
        "irrep_normalization": "component"
    }
    tp_e3nn = o3.TensorProduct(**kwargs).cuda()
    problem = oeq.TPProblem(**kwargs)
    tp_oeq = oeq.TensorProduct(problem)
    
    x1 = irreps_in1.randn(1, -1, device="cuda")
    x2 = irreps_in2.randn(1, -1, device="cuda")
    w = torch.randn(1, tp_e3nn.weight_numel, device="cuda")
    return tp_e3nn, tp_oeq, x1, x2, w


# multiplicity=1
torch.manual_seed(0)
irreps_in1 = o3.Irreps("8x0e+4x1e")
irreps_in2 = o3.Irreps("1x0e")
irreps_out = o3.Irreps("8x0e+4x1e")
tp_e3nn, tp_oeq, x1, x2, w = create_tps(irreps_in1, irreps_in2, irreps_out)
out_e3nn = tp_e3nn(x1, x2, w)
out_oeq = tp_oeq(x1, x2, w)
print(torch.abs(out_e3nn - out_oeq).max().item())
# -> 5.960464477539063e-08

# multiplicity=2
irreps_in1 = o3.Irreps("8x0e+4x1e")
irreps_in2 = o3.Irreps("2x0e")
irreps_out = o3.Irreps("8x0e+4x1e")
tp_e3nn, tp_oeq, x1, x2, w = create_tps(irreps_in1, irreps_in2, irreps_out)
out_oeq = tp_oeq(x1, x2, tp_oeq.reorder_weights_from_e3nn(w))
out_e3nn = tp_e3nn(x1, x2, w)
print(torch.abs(out_e3nn - out_oeq).max().item())
# -> 5.960464477539063e-08

[mjrs33]

Thank you for the quick fix!

The updated implementation works perfectly, and the issue with the multiplicities is now resolved. I really appreciate the fast turnaround and the clarification.
Thanks again for enabling broader use cases with this improvement — it makes further development much smoother!

[vbharadwaj-bk]

Note, by the way, that the reordering function (while differentiable) itself is somewhat slow. While it can be used to verify correctness, if you find that reordering impacts performance consider a) training a model from scratch with our kernels, so you can disregard the ordering difference or b) applying the weight reordering function once to last FCLayer of a trained model that produces the weights.