Ruff.

Author: vbharadwaj-bk

openequivariance/openequivariance/benchmark/correctness.py

@@ -16,7 +16,7 @@
 )
 from openequivariance.core.e3nn_lite import TPProblem
 from openequivariance.core.TensorProductBase import TensorProductBase
-from openequivariance.core.utils import IrrepLayoutUtils
+from openequivariance.core.utils import transpose_irrep_layout
 
 logger = getLogger()
 
@@ -87,29 +87,31 @@ def correctness_forward(
     outputs = []
 
     for i, impl in enumerate([test_implementation, reference_implementation]):
-        is_test_impl = (i == 0)
+        is_test_impl = i == 0
         tp = instantiate_implementation(impl, problem)
         uses_cue = impl == CUETensorProduct or isinstance(tp, CUETensorProduct)
-        run_in1, run_in2, run_weights, run_out = [ buf.copy() for buf in (in1, in2, weights, out) ] 
+        run_in1, run_in2, run_weights, run_out = [
+            buf.copy() for buf in (in1, in2, weights, out)
+        ]
 
         if problem.shared_weights and uses_cue:
             run_weights = run_weights[np.newaxis, :]
 
-        # Transpose inputs, if necessary, for the test implementation 
+        # Transpose inputs, if necessary, for the test implementation
         if is_test_impl:
             run_in1, run_in2 = [
-                IrrepLayoutUtils.transpose_irrep_layout(
-                    arr, irreps, "mul_ir", tp.config.layout 
-                )                for arr, irreps in zip(
-                    (run_in1, run_in2), 
-                    (problem.irreps_in1, problem.irreps_in2)
+                transpose_irrep_layout(arr, irreps, "mul_ir", tp.config.layout)
+                for arr, irreps in zip(
+                    (run_in1, run_in2), (problem.irreps_in1, problem.irreps_in2)
                 )
             ]
 
-        tp.forward_cpu(L1_in=run_in1, L2_in=run_in2, L3_out=run_out, weights=run_weights)
+        tp.forward_cpu(
+            L1_in=run_in1, L2_in=run_in2, L3_out=run_out, weights=run_weights
+        )
 
         if is_test_impl:
-            run_out = IrrepLayoutUtils.transpose_irrep_layout(
+            run_out = transpose_irrep_layout(
                 run_out, problem.irreps_out, tp.config.layout, "mul_ir"
             )
 
@@ -147,7 +149,15 @@ def correctness_backward(
         is_test_impl = i == 0
         tp = instantiate_implementation(impl, problem)
 
-        run_in1, run_in2, run_L3_grad, run_weights, run_weights_grad, run_in1_grad, run_in2_grad = [
+        (
+            run_in1,
+            run_in2,
+            run_L3_grad,
+            run_weights,
+            run_weights_grad,
+            run_in1_grad,
+            run_in2_grad,
+        ) = [
             buf.copy()
             for buf in (in1, in2, out_grad, weights, weights_grad, in1_grad, in2_grad)
         ]
@@ -159,9 +169,7 @@ def correctness_backward(
 
         if is_test_impl:
             run_in1, run_in2, run_L3_grad = [
-                IrrepLayoutUtils.transpose_irrep_layout(
-                    arr, irreps, "mul_ir", tp.config.layout
-                )
+                transpose_irrep_layout(arr, irreps, "mul_ir", tp.config.layout)
                 for arr, irreps in zip(
                     (run_in1, run_in2, run_L3_grad),
                     (problem.irreps_in1, problem.irreps_in2, problem.irreps_out),
@@ -180,9 +188,7 @@ def correctness_backward(
 
         if is_test_impl:
             run_in1_grad, run_in2_grad = [
-                IrrepLayoutUtils.transpose_irrep_layout(
-                    arr, irreps, tp.config.layout, "mul_ir"
-                )
+                transpose_irrep_layout(arr, irreps, tp.config.layout, "mul_ir")
                 for arr, irreps in zip(
                     (run_in1_grad, run_in2_grad),
                     (problem.irreps_in1, problem.irreps_in2),
@@ -254,9 +260,7 @@ def correctness_double_backward(
 
         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
-                )
+                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),
                     (
@@ -281,9 +285,7 @@ def correctness_double_backward(
 
         if is_test_impl:
             out_dgrad, in1_grad, in2_grad = [
-                IrrepLayoutUtils.transpose_irrep_layout(
-                    arr, irreps, tp.config.layout, "mul_ir"
-                )
+                transpose_irrep_layout(arr, irreps, tp.config.layout, "mul_ir")
                 for arr, irreps in zip(
                     (out_dgrad, in1_grad, in2_grad),
                     (problem.irreps_out, problem.irreps_in1, problem.irreps_in2),
@@ -359,9 +361,7 @@ def correctness_forward_conv(
 
         if is_test_impl:
             run_in1, run_in2 = [
-                IrrepLayoutUtils.transpose_irrep_layout(
-                    arr, irreps, "mul_ir", conv.config.layout
-                )
+                transpose_irrep_layout(arr, irreps, "mul_ir", conv.config.layout)
                 for arr, irreps in zip(
                     (run_in1, run_in2),
                     (conv.config.irreps_in1, conv.config.irreps_in2),
@@ -375,7 +375,7 @@ def correctness_forward_conv(
                 graph=graph,
             )
 
-            run_out = IrrepLayoutUtils.transpose_irrep_layout(
+            run_out = transpose_irrep_layout(
                 run_out, conv.config.irreps_out, conv.config.layout, "mul_ir"
             )
         else:
@@ -410,13 +410,9 @@ def correctness_forward_conv(
 
         for _ 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, rep_weights = [buf.copy() for buf in (in1, in2, weights)]
             rep_in1, rep_in2 = [
-                IrrepLayoutUtils.transpose_irrep_layout(
-                    arr, irreps, "mul_ir", conv.config.layout
-                )
+                transpose_irrep_layout(arr, irreps, "mul_ir", conv.config.layout)
                 for arr, irreps in zip(
                     (rep_in1, rep_in2),
                     (conv.config.irreps_in1, conv.config.irreps_in2),
@@ -430,7 +426,7 @@ def correctness_forward_conv(
                 graph=graph,
             )
 
-            repeated_run = IrrepLayoutUtils.transpose_irrep_layout(
+            repeated_run = transpose_irrep_layout(
                 repeated_run, conv.config.irreps_out, conv.config.layout, "mul_ir"
             )
 
@@ -471,9 +467,15 @@ def correctness_backward_conv(
         is_test_impl = i == 0
         tp = impl if is_test_impl else impl(reference_problem)
 
-        run_in1, run_in2, run_out_grad, run_weights, run_weights_grad, run_in1_grad, run_in2_grad = [
-            buf.copy() for buf in buffers
-        ]
+        (
+            run_in1,
+            run_in2,
+            run_out_grad,
+            run_weights,
+            run_weights_grad,
+            run_in1_grad,
+            run_in2_grad,
+        ) = [buf.copy() for buf in buffers]
 
         if not is_test_impl and high_precision_ref:
             (
@@ -488,12 +490,14 @@ def correctness_backward_conv(
 
         if is_test_impl:
             run_in1, run_in2, run_out_grad = [
-                IrrepLayoutUtils.transpose_irrep_layout(
-                    arr, irreps, "mul_ir", conv.config.layout
-                )
+                transpose_irrep_layout(arr, irreps, "mul_ir", conv.config.layout)
                 for arr, irreps in zip(
                     (run_in1, run_in2, run_out_grad),
-                    (conv.config.irreps_in1, conv.config.irreps_in2, conv.config.irreps_out),
+                    (
+                        conv.config.irreps_in1,
+                        conv.config.irreps_in2,
+                        conv.config.irreps_out,
+                    ),
                 )
             ]
 
@@ -510,9 +514,7 @@ def correctness_backward_conv(
 
         if is_test_impl:
             run_in1_grad, run_in2_grad = [
-                IrrepLayoutUtils.transpose_irrep_layout(
-                    arr, irreps, conv.config.layout, "mul_ir"
-                )
+                transpose_irrep_layout(arr, irreps, conv.config.layout, "mul_ir")
                 for arr, irreps in zip(
                     (run_in1_grad, run_in2_grad),
                     (conv.config.irreps_in1, conv.config.irreps_in2),
@@ -581,9 +583,7 @@ def correctness_double_backward_conv(
         ]
         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
-                )
+                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),
                     (
@@ -609,9 +609,7 @@ def correctness_double_backward_conv(
 
         if is_test_impl:
             out_dgrad, in1_grad, in2_grad = [
-                IrrepLayoutUtils.transpose_irrep_layout(
-                    arr, irreps, tp.config.layout, "mul_ir"
-                )
+                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),

openequivariance/openequivariance/core/ConvolutionBase.py

@@ -5,11 +5,10 @@
 from openequivariance.benchmark.logging import bcolors, getLogger
 from openequivariance.benchmark.random_buffer_utils import (
     get_random_buffers_backward_conv,
-    get_random_buffers_double_backward_conv,
     get_random_buffers_forward_conv,
 )
 from openequivariance.core.e3nn_lite import wigner_3j
-from openequivariance.core.utils import IrrepLayoutUtils, benchmark
+from openequivariance.core.utils import benchmark
 
 logger = getLogger()
 

openequivariance/openequivariance/core/utils.py

@@ -209,50 +209,52 @@ def hash_str_64(s: str) -> int:
     return int.from_bytes(hashlib.sha256(s.encode()).digest()[:7], "big")
 
 
-class IrrepLayoutUtils:
-    @staticmethod
-    def transpose_irrep_layout(
-        array: np.ndarray,
-        irreps: Irreps,
-        src_layout: str,
-        dst_layout: str,
-    ) -> np.ndarray:
-        """
-        Transpose irrep-packed feature arrays between `mul_ir` and `ir_mul` layouts.
-
-        Expected input shape is `[..., irreps.dim]`. A new array is returned.
-        If `src_layout == dst_layout`, this returns a copy.
-        """
-        if src_layout not in ("mul_ir", "ir_mul"):
-            raise ValueError(f"Unsupported src_layout: {src_layout}")
-        if dst_layout not in ("mul_ir", "ir_mul"):
-            raise ValueError(f"Unsupported dst_layout: {dst_layout}")
-
-        x = np.asarray(array)
-        out = np.empty_like(x)
-
-        if src_layout == dst_layout:
-            out[...] = x
-            return out
-
-        slices = irreps.slices()
-        for ir_idx, mul_ir in enumerate(irreps):
-            mul = mul_ir.mul
-            dim = mul_ir.ir.dim
-            seg = slices[ir_idx]
-            block = x[..., seg.start : seg.stop]
-
-            if src_layout == "ir_mul" and dst_layout == "mul_ir":
-                out[..., seg.start : seg.stop] = block.reshape(
-                    *block.shape[:-1], dim, mul
-                ).swapaxes(-1, -2).reshape(*block.shape[:-1], mul * dim)
-            elif src_layout == "mul_ir" and dst_layout == "ir_mul":
-                out[..., seg.start : seg.stop] = block.reshape(
-                    *block.shape[:-1], mul, dim
-                ).swapaxes(-1, -2).reshape(*block.shape[:-1], dim * mul)
-            else:
-                raise ValueError(
-                    f"Unsupported layout transpose: {src_layout} -> {dst_layout}"
-                )
+def transpose_irrep_layout(
+    array: np.ndarray,
+    irreps: Irreps,
+    src_layout: str,
+    dst_layout: str,
+) -> np.ndarray:
+    """
+    Transpose irrep-packed feature arrays between `mul_ir` and `ir_mul` layouts.
+
+    Expected input shape is `[..., irreps.dim]`. A new array is returned.
+    If `src_layout == dst_layout`, this returns a copy.
+    """
+    if src_layout not in ("mul_ir", "ir_mul"):
+        raise ValueError(f"Unsupported src_layout: {src_layout}")
+    if dst_layout not in ("mul_ir", "ir_mul"):
+        raise ValueError(f"Unsupported dst_layout: {dst_layout}")
 
+    x = np.asarray(array)
+    out = np.empty_like(x)
+
+    if src_layout == dst_layout:
+        out[...] = x
         return out
+
+    slices = irreps.slices()
+    for ir_idx, mul_ir in enumerate(irreps):
+        mul = mul_ir.mul
+        dim = mul_ir.ir.dim
+        seg = slices[ir_idx]
+        block = x[..., seg.start : seg.stop]
+
+        if src_layout == "ir_mul" and dst_layout == "mul_ir":
+            out[..., seg.start : seg.stop] = (
+                block.reshape(*block.shape[:-1], dim, mul)
+                .swapaxes(-1, -2)
+                .reshape(*block.shape[:-1], mul * dim)
+            )
+        elif src_layout == "mul_ir" and dst_layout == "ir_mul":
+            out[..., seg.start : seg.stop] = (
+                block.reshape(*block.shape[:-1], mul, dim)
+                .swapaxes(-1, -2)
+                .reshape(*block.shape[:-1], dim * mul)
+            )
+        else:
+            raise ValueError(
+                f"Unsupported layout transpose: {src_layout} -> {dst_layout}"
+            )
+
+    return out

tests/batch_test.py

@@ -272,9 +272,9 @@ def tp_and_problem(self, problem, extra_tp_constructor_args, with_jax):
             tp.to(switch_map[problem.irrep_dtype])
             return tp, tp.config
 
+
 class TestIrMulLayoutMACE(TPCorrectness):
-    production_model_tpps = mace_problems() + \
-        [
+    production_model_tpps = mace_problems() + [
         oeq.TPProblem(
             "5x5e",
             "1x3e",
@@ -293,7 +293,7 @@ class TestIrMulLayoutMACE(TPCorrectness):
             internal_weights=False,
             label="ir_mul_repr_13x1x13_l535",
         ),
-        ] 
+    ]
 
     @pytest.fixture(params=production_model_tpps, ids=lambda x: x.label, scope="class")
     def problem(self, request, dtype):

tests/conv_test.py

@@ -293,8 +293,7 @@ def conv_object(self, request, problem, extra_conv_constructor_args):
 
 
 class TestIrMulLayout(ConvCorrectness):
-    production_model_tpps = mace_problems() + \
-        [
+    production_model_tpps = mace_problems() + [
         oeq.TPProblem(
             "5x5e",
             "1x3e",

tests/input_validation_test.py

@@ -154,4 +154,4 @@ def test_ir_mul_rejects_uvw_problem(dtype):
     )
 
     with pytest.raises(AssertionError, match="layout='ir_mul'"):
-        TensorProduct(problem)
+        TensorProduct(problem)