More bugfixes.

Author: vbharadwaj-bk

openequivariance/openequivariance/core/ComputationSchedule.py

@@ -301,6 +301,7 @@ def __init__(self, input, mult_threshold):
             path_normalization="none",
             internal_weights=False,
             shared_weights=input.shared_weights,
+            layout=input.layout,
         )
 
         assert self.output.weight_numel == input.weight_numel
@@ -595,6 +596,7 @@ def calculate_backward_smem(
                 path_normalization="none",
                 internal_weights=False,
                 shared_weights=config.shared_weights,
+                layout=config.layout,
             )
 
             weight_offset = 0

openequivariance/openequivariance/core/utils.py

@@ -240,11 +240,11 @@ def transpose_irrep_layout(
             if src_layout == "ir_mul" and dst_layout == "mul_ir":
                 out[..., seg.start : seg.stop] = block.reshape(
                     *block.shape[:-1], dim, mul
-                ).reshape(*block.shape[:-1], mul * dim)
+                ).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
-                ).reshape(*block.shape[:-1], dim * mul)
+                ).swapaxes(-1, -2).reshape(*block.shape[:-1], dim * mul)
             else:
                 raise ValueError(
                     f"Unsupported layout transpose: {src_layout} -> {dst_layout}"

tests/batch_test.py

@@ -273,14 +273,14 @@ def tp_and_problem(self, problem, extra_tp_constructor_args, with_jax):
             return tp, tp.config
 
 
-class TestMulIrLayoutMACE(TPCorrectness):
+class TestIrMulLayoutMACE(TPCorrectness):
     production_model_tpps = mace_problems()
 
     @pytest.fixture(params=production_model_tpps, ids=lambda x: x.label, scope="class")
     def problem(self, request, dtype):
         problem = request.param.clone()
         problem.irrep_dtype, problem.weight_dtype = dtype, dtype
-        problem.layout = "mul_ir"
+        problem.layout = "ir_mul"
         return problem