Made more progress.

vbharadwaj-bk · vbharadwaj-bk · commit 85e988f634d9 · 2026-03-21T18:47:28.000-07:00
diff --git a/openequivariance/openequivariance/_torch/NPDoubleBackwardMixin.py b/openequivariance/openequivariance/_torch/NPDoubleBackwardMixin.py
@@ -1,7 +1,5 @@
 import torch
 
-from openequivariance.core.utils import IrrepLayoutUtils
-
 
 class NumpyDoubleBackwardMixin:
     """
@@ -15,30 +13,12 @@ def double_backward_cpu(
     ):
         assert self.torch_op
 
-        layout = self.config.layout
-
-        in1_kernel = IrrepLayoutUtils.transpose_irrep_layout(
-            in1, self.config.irreps_in1, layout, "mul_ir"
-        )
-        in2_kernel = IrrepLayoutUtils.transpose_irrep_layout(
-            in2, self.config.irreps_in2, layout, "mul_ir"
-        )
-        out_grad_kernel = IrrepLayoutUtils.transpose_irrep_layout(
-            out_grad, self.config.irreps_out, layout, "mul_ir"
-        )
-        in1_dgrad_kernel = IrrepLayoutUtils.transpose_irrep_layout(
-            in1_dgrad, self.config.irreps_in1, layout, "mul_ir"
-        )
-        in2_dgrad_kernel = IrrepLayoutUtils.transpose_irrep_layout(
-            in2_dgrad, self.config.irreps_in2, layout, "mul_ir"
-        )
-
-        in1_torch = torch.tensor(in1_kernel).to("cuda").requires_grad_(True)
-        in2_torch = torch.tensor(in2_kernel).to("cuda").requires_grad_(True)
+        in1_torch = torch.tensor(in1).to("cuda").requires_grad_(True)
+        in2_torch = torch.tensor(in2).to("cuda").requires_grad_(True)
         weights_torch = torch.tensor(weights).to("cuda").requires_grad_(True)
-        out_grad_torch = torch.tensor(out_grad_kernel).to("cuda").requires_grad_(True)
-        in1_dgrad_torch = torch.tensor(in1_dgrad_kernel).to("cuda")
-        in2_dgrad_torch = torch.tensor(in2_dgrad_kernel).to("cuda")
+        out_grad_torch = torch.tensor(out_grad).to("cuda").requires_grad_(True)
+        in1_dgrad_torch = torch.tensor(in1_dgrad).to("cuda")
+        in2_dgrad_torch = torch.tensor(in2_dgrad).to("cuda")
         weights_dgrad_torch = torch.tensor(weights_dgrad).to("cuda")
         out_torch = self.forward(in1_torch, in2_torch, weights_torch)
 
@@ -61,16 +41,6 @@ def double_backward_cpu(
         c_np = c.detach().cpu().numpy()
         d_np = d.detach().cpu().numpy()
 
-        a_np = IrrepLayoutUtils.transpose_irrep_layout(
-            a_np, self.config.irreps_in1, "mul_ir", layout
-        )
-        b_np = IrrepLayoutUtils.transpose_irrep_layout(
-            b_np, self.config.irreps_in2, "mul_ir", layout
-        )
-        d_np = IrrepLayoutUtils.transpose_irrep_layout(
-            d_np, self.config.irreps_out, "mul_ir", layout
-        )
-
         return (a_np, b_np, c_np, d_np)
 
 
@@ -84,30 +54,12 @@ def double_backward_cpu(
     ):
         assert self.torch_op
 
-        layout = self.config.layout
-
-        in1_kernel = IrrepLayoutUtils.transpose_irrep_layout(
-            in1, self.config.irreps_in1, layout, "mul_ir"
-        )
-        in2_kernel = IrrepLayoutUtils.transpose_irrep_layout(
-            in2, self.config.irreps_in2, layout, "mul_ir"
-        )
-        out_grad_kernel = IrrepLayoutUtils.transpose_irrep_layout(
-            out_grad, self.config.irreps_out, layout, "mul_ir"
-        )
-        in1_dgrad_kernel = IrrepLayoutUtils.transpose_irrep_layout(
-            in1_dgrad, self.config.irreps_in1, layout, "mul_ir"
-        )
-        in2_dgrad_kernel = IrrepLayoutUtils.transpose_irrep_layout(
-            in2_dgrad, self.config.irreps_in2, layout, "mul_ir"
-        )
-
-        in1_torch = torch.tensor(in1_kernel).to("cuda").requires_grad_(True)
-        in2_torch = torch.tensor(in2_kernel).to("cuda").requires_grad_(True)
+        in1_torch = torch.tensor(in1).to("cuda").requires_grad_(True)
+        in2_torch = torch.tensor(in2).to("cuda").requires_grad_(True)
         weights_torch = torch.tensor(weights).to("cuda").requires_grad_(True)
-        out_grad_torch = torch.tensor(out_grad_kernel).to("cuda").requires_grad_(True)
-        in1_dgrad_torch = torch.tensor(in1_dgrad_kernel).to("cuda")
-        in2_dgrad_torch = torch.tensor(in2_dgrad_kernel).to("cuda")
+        out_grad_torch = torch.tensor(out_grad).to("cuda").requires_grad_(True)
+        in1_dgrad_torch = torch.tensor(in1_dgrad).to("cuda")
+        in2_dgrad_torch = torch.tensor(in2_dgrad).to("cuda")
         weights_dgrad_torch = torch.tensor(weights_dgrad).to("cuda")
 
         torch_rows = torch.tensor(graph.rows, device="cuda")
@@ -142,14 +94,4 @@ def double_backward_cpu(
         c_np = c.detach().cpu().numpy()
         d_np = d.detach().cpu().numpy()
 
-        a_np = IrrepLayoutUtils.transpose_irrep_layout(
-            a_np, self.config.irreps_in1, "mul_ir", layout
-        )
-        b_np = IrrepLayoutUtils.transpose_irrep_layout(
-            b_np, self.config.irreps_in2, "mul_ir", layout
-        )
-        d_np = IrrepLayoutUtils.transpose_irrep_layout(
-            d_np, self.config.irreps_out, "mul_ir", layout
-        )
-
         return (a_np, b_np, c_np, d_np)
diff --git a/openequivariance/openequivariance/_torch/TensorProduct.py b/openequivariance/openequivariance/_torch/TensorProduct.py
@@ -11,11 +11,7 @@
 )
 from openequivariance.benchmark.logging_utils import getLogger
 from openequivariance.core.LoopUnrollTP import LoopUnrollTP
-from openequivariance.core.utils import (
-    IrrepLayoutUtils,
-    dtype_to_enum,
-    torch_to_oeq_dtype,
-)
+from openequivariance.core.utils import dtype_to_enum, torch_to_oeq_dtype
 
 logger = getLogger()
 
@@ -150,24 +146,12 @@ def forward_cpu(
             weights, not self.config.shared_weights
         )
 
-        layout = self.config.layout
-
-        L1_in_kernel = IrrepLayoutUtils.transpose_irrep_layout(
-            L1_in, self.config.irreps_in1, layout, "mul_ir"
-        )
-        L2_in_kernel = IrrepLayoutUtils.transpose_irrep_layout(
-            L2_in, self.config.irreps_in2, layout, "mul_ir"
-        )
-
-        torch_L1_in = torch.tensor(L1_in_kernel, device="cuda")
-        torch_L2_in = torch.tensor(L2_in_kernel, device="cuda")
+        torch_L1_in = torch.tensor(L1_in, device="cuda")
+        torch_L2_in = torch.tensor(L2_in, device="cuda")
         torch_weights = torch.tensor(weights_chunked, device="cuda")
         torch_L3_out = self.forward(torch_L1_in, torch_L2_in, torch_weights)
 
-        L3_kernel = torch_L3_out.numpy(force=True)
-        L3_out[:] = IrrepLayoutUtils.transpose_irrep_layout(
-            L3_kernel, self.config.irreps_out, "mul_ir", layout
-        )
+        L3_out[:] = torch_L3_out.numpy(force=True)
 
     def backward_cpu(
         self, L1_in, L1_grad, L2_in, L2_grad, L3_grad, weights, weights_grad
@@ -176,37 +160,18 @@ def backward_cpu(
             weights, not self.config.shared_weights
         )
 
-        layout = self.config.layout
-
-        L1_in_kernel = IrrepLayoutUtils.transpose_irrep_layout(
-            L1_in, self.config.irreps_in1, layout, "mul_ir"
-        )
-        L2_in_kernel = IrrepLayoutUtils.transpose_irrep_layout(
-            L2_in, self.config.irreps_in2, layout, "mul_ir"
-        )
-        L3_grad_kernel = IrrepLayoutUtils.transpose_irrep_layout(
-            L3_grad, self.config.irreps_out, layout, "mul_ir"
-        )
-
-        torch_L1_in = torch.tensor(L1_in_kernel, requires_grad=True, device="cuda")
-        torch_L2_in = torch.tensor(L2_in_kernel, requires_grad=True, device="cuda")
+        torch_L1_in = torch.tensor(L1_in, requires_grad=True, device="cuda")
+        torch_L2_in = torch.tensor(L2_in, requires_grad=True, device="cuda")
         torch_weights = torch.tensor(weights_chunked, requires_grad=True, device="cuda")
 
         torch_out = self.forward(torch_L1_in, torch_L2_in, torch_weights)
 
-        torch_L3_grad_in = torch.tensor(L3_grad_kernel, device="cuda")
+        torch_L3_grad_in = torch.tensor(L3_grad, device="cuda")
 
         torch_out.backward(gradient=torch_L3_grad_in)
 
-        L1_grad_kernel = torch_L1_in.grad.numpy(force=True)
-        L2_grad_kernel = torch_L2_in.grad.numpy(force=True)
-
-        L1_grad[:] = IrrepLayoutUtils.transpose_irrep_layout(
-            L1_grad_kernel, self.config.irreps_in1, "mul_ir", layout
-        )
-        L2_grad[:] = IrrepLayoutUtils.transpose_irrep_layout(
-            L2_grad_kernel, self.config.irreps_in2, "mul_ir", layout
-        )
+        L1_grad[:] = torch_L1_in.grad.numpy(force=True)
+        L2_grad[:] = torch_L2_in.grad.numpy(force=True)
         weights_grad[:] = torch_weights.grad.numpy(force=True)
 
         weights_grad[:] = self.reorder_weights_to_e3nn(
diff --git a/openequivariance/openequivariance/benchmark/correctness_utils.py b/openequivariance/openequivariance/benchmark/correctness_utils.py
@@ -1,17 +1,18 @@
 from typing import Optional, Union
 
-from openequivariance.core.TensorProductBase import TensorProductBase
-from openequivariance.core.e3nn_lite import TPProblem
+import numpy as np
+import numpy.linalg as la
+
 from openequivariance._torch.CUETensorProduct import CUETensorProduct
+from openequivariance.benchmark.logging_utils import bcolors, getLogger
 from openequivariance.benchmark.random_buffer_utils import (
-    get_random_buffers_forward,
     get_random_buffers_backward,
     get_random_buffers_double_backward,
+    get_random_buffers_forward,
 )
-
-from openequivariance.benchmark.logging_utils import getLogger, bcolors
-import numpy as np
-import numpy.linalg as la
+from openequivariance.core.e3nn_lite import TPProblem
+from openequivariance.core.TensorProductBase import TensorProductBase
+from openequivariance.core.utils import IrrepLayoutUtils
 
 logger = getLogger()
 
@@ -81,7 +82,7 @@ def correctness_forward(
 
     in1, in2, weights, out = get_random_buffers_forward(problem, batch_size, prng_seed)
 
-    # run reference
+    # run reference (always in mul_ir)
     ref_tp = reference_implementation(problem)
 
     ref_out = out.copy()
@@ -93,13 +94,31 @@ def correctness_forward(
     if problem.shared_weights and test_implementation == CUETensorProduct:
         weights_copy = weights[np.newaxis, :]
 
-    # run test
+    # run test (may require ir_mul conversion)
     test_tp = instantiate_implementation(test_implementation, problem)
+    test_layout = getattr(test_tp.config, "layout", "mul_ir")
+
+    test_in1 = in1.copy()
+    test_in2 = in2.copy()
     test_out = out.copy()
+
+    if test_layout == "ir_mul":
+        test_in1 = IrrepLayoutUtils.transpose_irrep_layout(
+            test_in1, problem.irreps_in1, "mul_ir", "ir_mul"
+        )
+        test_in2 = IrrepLayoutUtils.transpose_irrep_layout(
+            test_in2, problem.irreps_in2, "mul_ir", "ir_mul"
+        )
+
     test_tp.forward_cpu(
-        L1_in=in1.copy(), L2_in=in2.copy(), L3_out=test_out, weights=weights_copy
+        L1_in=test_in1, L2_in=test_in2, L3_out=test_out, weights=weights_copy
     )
 
+    if test_layout == "ir_mul":
+        test_out = IrrepLayoutUtils.transpose_irrep_layout(
+            test_out, problem.irreps_out, "ir_mul", "mul_ir"
+        )
+
     for name, to_check, ground_truth in [("output", ref_out, test_out)]:
         result[name] = check_similiarity(
             name, to_check, ground_truth, correctness_threshold
@@ -144,7 +163,7 @@ def correctness_backward(
         weights_grad=ref_weights_grad,
     )
 
-    # run test version
+    # run test version (may require ir_mul conversion)
     test_weights_grad = weights_grad.copy()
     test_in1_grad = in1_grad.copy()
     test_in2_grad = in2_grad.copy()
@@ -156,16 +175,41 @@ def correctness_backward(
         test_weights_grad = test_weights_grad[np.newaxis, :]
 
     test_tp = instantiate_implementation(test_implementation, problem)
+    test_layout = getattr(test_tp.config, "layout", "mul_ir")
+
+    test_in1 = in1.copy()
+    test_in2 = in2.copy()
+    test_L3_grad = out_grad.copy()
+
+    if test_layout == "ir_mul":
+        test_in1 = IrrepLayoutUtils.transpose_irrep_layout(
+            test_in1, problem.irreps_in1, "mul_ir", "ir_mul"
+        )
+        test_in2 = IrrepLayoutUtils.transpose_irrep_layout(
+            test_in2, problem.irreps_in2, "mul_ir", "ir_mul"
+        )
+        test_L3_grad = IrrepLayoutUtils.transpose_irrep_layout(
+            test_L3_grad, problem.irreps_out, "mul_ir", "ir_mul"
+        )
+
     test_tp.backward_cpu(
-        L1_in=in1.copy(),
+        L1_in=test_in1,
         L1_grad=test_in1_grad,
-        L2_in=in2.copy(),
+        L2_in=test_in2,
         L2_grad=test_in2_grad,
-        L3_grad=out_grad.copy(),
+        L3_grad=test_L3_grad,
         weights=weights_copy,
         weights_grad=test_weights_grad,
     )
 
+    if test_layout == "ir_mul":
+        test_in1_grad = IrrepLayoutUtils.transpose_irrep_layout(
+            test_in1_grad, problem.irreps_in1, "ir_mul", "mul_ir"
+        )
+        test_in2_grad = IrrepLayoutUtils.transpose_irrep_layout(
+            test_in2_grad, problem.irreps_in2, "ir_mul", "mul_ir"
+        )
+
     weight_threshold = (
         correctness_threshold * batch_size
         if problem.shared_weights
@@ -210,7 +254,9 @@ def correctness_double_backward(
     result = {"thresh": correctness_threshold, "batch_size": batch_size}
 
     tensors = []
-    for _, impl in enumerate([test_implementation, reference_implementation]):
+    for is_test_impl, impl in enumerate(
+        [test_implementation, reference_implementation]
+    ):
         tp = instantiate_implementation(impl, problem)
         weights_reordered = tp.reorder_weights_from_e3nn(
             weights, has_batch_dim=not problem.shared_weights
@@ -222,15 +268,53 @@ def correctness_double_backward(
         if impl == CUETensorProduct and problem.shared_weights:
             weights_reordered = weights_reordered[np.newaxis, :]
 
+        tp_layout = getattr(tp.config, "layout", "mul_ir")
+        apply_test_layout = is_test_impl == 0 and tp_layout == "ir_mul"
+
+        db_in1 = in1
+        db_in2 = in2
+        db_out_grad = out_grad
+        db_in1_dgrad = in1_dgrad
+        db_in2_dgrad = in2_dgrad
+
+        if apply_test_layout:
+            db_in1 = IrrepLayoutUtils.transpose_irrep_layout(
+                in1, problem.irreps_in1, "mul_ir", "ir_mul"
+            )
+            db_in2 = IrrepLayoutUtils.transpose_irrep_layout(
+                in2, problem.irreps_in2, "mul_ir", "ir_mul"
+            )
+            db_out_grad = IrrepLayoutUtils.transpose_irrep_layout(
+                out_grad, problem.irreps_out, "mul_ir", "ir_mul"
+            )
+            db_in1_dgrad = IrrepLayoutUtils.transpose_irrep_layout(
+                in1_dgrad, problem.irreps_in1, "mul_ir", "ir_mul"
+            )
+            db_in2_dgrad = IrrepLayoutUtils.transpose_irrep_layout(
+                in2_dgrad, problem.irreps_in2, "mul_ir", "ir_mul"
+            )
+
         in1_grad, in2_grad, weights_grad, out_dgrad = tp.double_backward_cpu(
-            in1,
-            in2,
-            out_grad,
+            db_in1,
+            db_in2,
+            db_out_grad,
             weights_reordered,
             weights_dgrad_reordered,
-            in1_dgrad,
-            in2_dgrad,
+            db_in1_dgrad,
+            db_in2_dgrad,
         )
+
+        if apply_test_layout:
+            out_dgrad = IrrepLayoutUtils.transpose_irrep_layout(
+                out_dgrad, problem.irreps_out, "ir_mul", "mul_ir"
+            )
+            in1_grad = IrrepLayoutUtils.transpose_irrep_layout(
+                in1_grad, problem.irreps_in1, "ir_mul", "mul_ir"
+            )
+            in2_grad = IrrepLayoutUtils.transpose_irrep_layout(
+                in2_grad, problem.irreps_in2, "ir_mul", "mul_ir"
+            )
+
         tensors.append(
             (
                 out_dgrad,
