More compaction.

Author: vbharadwaj-bk

openequivariance/openequivariance/core/ConvolutionBase.py

@@ -144,13 +144,6 @@ def test_correctness_forward(
         check_reproducible=True,
         high_precision_ref=False,
     ):
-        def maybe_transpose_input_for_test_impl(x, irreps):
-            if self.config.layout == "ir_mul":
-                return IrrepLayoutUtils.transpose_irrep_layout(
-                    x, irreps, "mul_ir", "ir_mul"
-                )
-            return x
-
         if reference_implementation is None:
             from openequivariance._torch.E3NNConv import E3NNConv
 
@@ -192,23 +185,29 @@ def maybe_transpose_input_for_test_impl(x, irreps):
 
         ref_out[:] = ref_tp.forward(**args).cpu().numpy()
 
-        test_out = out.copy()
+        run_in1, run_in2, run_weights, test_out = [
+            buf.copy() for buf in (in1, in2, weights, out)
+        ]
+        run_in1, run_in2 = [
+            IrrepLayoutUtils.transpose_irrep_layout(
+                arr, irreps, "mul_ir", self.config.layout
+            )
+            for arr, irreps in zip(
+                (run_in1, run_in2),
+                (self.config.irreps_in1, self.config.irreps_in2),
+            )
+        ]
         self.forward_cpu(
-            L1_in=maybe_transpose_input_for_test_impl(
-                in1.copy(), self.config.irreps_in1
-            ),
-            L2_in=maybe_transpose_input_for_test_impl(
-                in2.copy(), self.config.irreps_in2
-            ),
-            weights=weights.copy(),
+            L1_in=run_in1,
+            L2_in=run_in2,
+            weights=run_weights,
             L3_out=test_out,
             graph=graph,
         )
 
-        if self.config.layout == "ir_mul":
-            test_out = IrrepLayoutUtils.transpose_irrep_layout(
-                test_out, self.config.irreps_out, "ir_mul", "mul_ir"
-            )
+        test_out = IrrepLayoutUtils.transpose_irrep_layout(
+            test_out, self.config.irreps_out, self.config.layout, "mul_ir"
+        )
 
         for name, to_check, ground_truth in [("output", ref_out, test_out)]:
             result[name] = check_similiarity(name, to_check, ground_truth, thresh)
@@ -221,22 +220,29 @@ def maybe_transpose_input_for_test_impl(x, irreps):
 
             for i in range(num_trials):
                 repeated_run = out.copy()
+                rep_in1, rep_in2, rep_weights = [
+                    buf.copy() for buf in (in1, in2, weights)
+                ]
+                rep_in1, rep_in2 = [
+                    IrrepLayoutUtils.transpose_irrep_layout(
+                        arr, irreps, "mul_ir", self.config.layout
+                    )
+                    for arr, irreps in zip(
+                        (rep_in1, rep_in2),
+                        (self.config.irreps_in1, self.config.irreps_in2),
+                    )
+                ]
                 self.forward_cpu(
-                    L1_in=maybe_transpose_input_for_test_impl(
-                        in1.copy(), self.config.irreps_in1
-                    ),
-                    L2_in=maybe_transpose_input_for_test_impl(
-                        in2.copy(), self.config.irreps_in2
-                    ),
-                    weights=weights.copy(),
+                    L1_in=rep_in1,
+                    L2_in=rep_in2,
+                    weights=rep_weights,
                     L3_out=repeated_run,
                     graph=graph,
                 )
 
-                if self.config.layout == "ir_mul":
-                    repeated_run = IrrepLayoutUtils.transpose_irrep_layout(
-                        repeated_run, self.config.irreps_out, "ir_mul", "mul_ir"
-                    )
+                repeated_run = IrrepLayoutUtils.transpose_irrep_layout(
+                    repeated_run, self.config.irreps_out, self.config.layout, "mul_ir"
+                )
 
                 for name, to_check, ground_truth in [
                     ("output", repeated_run, test_out)
@@ -413,13 +419,6 @@ def test_correctness_backward(
         reference_implementation=None,
         high_precision_ref=False,
     ):
-        def maybe_transpose_input_for_test_impl(x, irreps):
-            if self.config.layout == "ir_mul":
-                return IrrepLayoutUtils.transpose_irrep_layout(
-                    x, irreps, "mul_ir", "ir_mul"
-                )
-            return x
-
         if reference_implementation is None:
             from openequivariance._torch.E3NNConv import E3NNConv
 
@@ -469,34 +468,39 @@ def maybe_transpose_input_for_test_impl(x, irreps):
         test_in1_grad = in1_grad.copy()
         test_in2_grad = in2_grad.copy()
 
-        test_L3_grad = out_grad.copy()
-        if self.config.layout == "ir_mul":
-            test_L3_grad = IrrepLayoutUtils.transpose_irrep_layout(
-                test_L3_grad, self.config.irreps_out, "mul_ir", "ir_mul"
+        test_in1, test_in2, test_L3_grad = [
+            buf.copy() for buf in (in1, in2, out_grad)
+        ]
+        test_in1, test_in2, test_L3_grad = [
+            IrrepLayoutUtils.transpose_irrep_layout(
+                arr, irreps, "mul_ir", self.config.layout
+            )
+            for arr, irreps in zip(
+                (test_in1, test_in2, test_L3_grad),
+                (self.config.irreps_in1, self.config.irreps_in2, self.config.irreps_out),
             )
+        ]
 
         self.backward_cpu(
-            L1_in=maybe_transpose_input_for_test_impl(
-                in1.copy(), self.config.irreps_in1
-            ),
+            L1_in=test_in1,
             L1_grad=test_in1_grad,
-            L2_in=maybe_transpose_input_for_test_impl(
-                in2.copy(), self.config.irreps_in2
-            ),
+            L2_in=test_in2,
             L2_grad=test_in2_grad,
             L3_grad=test_L3_grad,
             weights=weights.copy(),
             weights_grad=test_weights_grad,
             graph=graph,
         )
 
-        if self.config.layout == "ir_mul":
-            test_in1_grad = IrrepLayoutUtils.transpose_irrep_layout(
-                test_in1_grad, self.config.irreps_in1, "ir_mul", "mul_ir"
+        test_in1_grad, test_in2_grad = [
+            IrrepLayoutUtils.transpose_irrep_layout(
+                arr, irreps, self.config.layout, "mul_ir"
             )
-            test_in2_grad = IrrepLayoutUtils.transpose_irrep_layout(
-                test_in2_grad, self.config.irreps_in2, "ir_mul", "mul_ir"
+            for arr, irreps in zip(
+                (test_in1_grad, test_in2_grad),
+                (self.config.irreps_in1, self.config.irreps_in2),
             )
+        ]
 
         for name, to_check, ground_truth, threshold in [
             ("weight_grad", test_weights_grad, ref_weights_grad, thresh),
@@ -515,13 +519,6 @@ def test_correctness_double_backward(
         reference_implementation=None,
         high_precision_ref=False,
     ):
-        def maybe_transpose_input_for_test_impl(tp, x, irreps):
-            if tp is self and tp.config.layout == "ir_mul":
-                return IrrepLayoutUtils.transpose_irrep_layout(
-                    x, irreps, "mul_ir", "ir_mul"
-                )
-            return x
-
         buffers = get_random_buffers_double_backward_conv(
             self.config, graph.node_count, graph.nnz, prng_seed
         )
@@ -542,6 +539,7 @@ def maybe_transpose_input_for_test_impl(tp, x, irreps):
         result = {"thresh": thresh}
         tensors = []
         for i, tp in enumerate([self, reference_tp]):
+            is_test_impl = i == 0
             buffers_copy = [buf.copy() for buf in buffers]
 
             if i == 1 and high_precision_ref:
@@ -558,21 +556,25 @@ def maybe_transpose_input_for_test_impl(tp, x, irreps):
                 weights_dgrad, not tp.config.shared_weights
             )
 
-            db_in1 = maybe_transpose_input_for_test_impl(tp, in1, tp.config.irreps_in1)
-            db_in2 = maybe_transpose_input_for_test_impl(tp, in2, tp.config.irreps_in2)
-            db_out_grad = out_grad
-            db_in1_dgrad = in1_dgrad
-            db_in2_dgrad = in2_dgrad
-            if tp is self and tp.config.layout == "ir_mul":
-                db_out_grad = IrrepLayoutUtils.transpose_irrep_layout(
-                    out_grad, tp.config.irreps_out, "mul_ir", "ir_mul"
-                )
-                db_in1_dgrad = IrrepLayoutUtils.transpose_irrep_layout(
-                    in1_dgrad, tp.config.irreps_in1, "mul_ir", "ir_mul"
-                )
-                db_in2_dgrad = IrrepLayoutUtils.transpose_irrep_layout(
-                    in2_dgrad, tp.config.irreps_in2, "mul_ir", "ir_mul"
-                )
+            db_in1, db_in2, db_out_grad, db_in1_dgrad, db_in2_dgrad = [
+                buf.copy() for buf in (in1, in2, out_grad, in1_dgrad, in2_dgrad)
+            ]
+            if is_test_impl:
+                db_in1, db_in2, db_out_grad, db_in1_dgrad, db_in2_dgrad = [
+                    IrrepLayoutUtils.transpose_irrep_layout(
+                        arr, irreps, "mul_ir", tp.config.layout
+                    )
+                    for arr, irreps in zip(
+                        (db_in1, db_in2, db_out_grad, db_in1_dgrad, db_in2_dgrad),
+                        (
+                            tp.config.irreps_in1,
+                            tp.config.irreps_in2,
+                            tp.config.irreps_out,
+                            tp.config.irreps_in1,
+                            tp.config.irreps_in2,
+                        ),
+                    )
+                ]
 
             in1_grad, in2_grad, weights_grad, out_dgrad = tp.double_backward_cpu(
                 db_in1,
@@ -585,16 +587,16 @@ def maybe_transpose_input_for_test_impl(tp, x, irreps):
                 graph,
             )
 
-            if tp is self and tp.config.layout == "ir_mul":
-                out_dgrad = IrrepLayoutUtils.transpose_irrep_layout(
-                    out_dgrad, tp.config.irreps_out, "ir_mul", "mul_ir"
-                )
-                in1_grad = IrrepLayoutUtils.transpose_irrep_layout(
-                    in1_grad, tp.config.irreps_in1, "ir_mul", "mul_ir"
-                )
-                in2_grad = IrrepLayoutUtils.transpose_irrep_layout(
-                    in2_grad, tp.config.irreps_in2, "ir_mul", "mul_ir"
-                )
+            if is_test_impl:
+                out_dgrad, in1_grad, in2_grad = [
+                    IrrepLayoutUtils.transpose_irrep_layout(
+                        arr, irreps, tp.config.layout, "mul_ir"
+                    )
+                    for arr, irreps in zip(
+                        (out_dgrad, in1_grad, in2_grad),
+                        (tp.config.irreps_out, tp.config.irreps_in1, tp.config.irreps_in2),
+                    )
+                ]
 
             tensors.append(
                 (

tests/conv_test.py

@@ -284,8 +284,9 @@ def conv_object(self, request, problem, extra_conv_constructor_args):
         return module.to(switch_map[problem.irrep_dtype])
 
 
-def ir_mul_representative_uvu_problems():
-    return [
+class TestIrMulLayout(ConvCorrectness):
+    production_model_tpps = mace_problems() + \
+        [
         oeq.TPProblem(
             "5x5e",
             "1x3e",
@@ -306,10 +307,6 @@ def ir_mul_representative_uvu_problems():
         ),
     ]
 
-
-class TestIrMulLayout(ConvCorrectness):
-    production_model_tpps = mace_problems() + ir_mul_representative_uvu_problems()
-
     @pytest.fixture(params=production_model_tpps, ids=lambda x: x.label, scope="class")
     def problem(self, request, dtype):
         problem = request.param.clone()