Fixed more stuff.

Author: vbharadwaj-bk

openequivariance/openequivariance/core/utils.py

@@ -96,6 +96,11 @@ def filter_and_analyze_problem(problem):
         f"Connection mode must be 'uvu' or 'uvw', got {problem.instructions[0].connection_mode}"
     )
 
+    if problem.layout == "ir_mul":
+        assert problem.instructions[0].connection_mode == "uvu", (
+            "layout='ir_mul' is only supported for pure 'uvu' problems"
+        )
+
     assert problem.irrep_dtype == problem.weight_dtype, (
         f"irrep_dtype and weight_dtype must be the same, got {problem.irrep_dtype} and {problem.weight_dtype}"
     )

tests/batch_test.py

@@ -273,8 +273,31 @@ def tp_and_problem(self, problem, extra_tp_constructor_args, with_jax):
             return tp, tp.config
 
 
+def ir_mul_representative_uvu_problems():
+    return [
+        oeq.TPProblem(
+            "5x5e",
+            "1x3e",
+            "5x5e",
+            [(0, 0, 0, "uvu", True)],
+            shared_weights=False,
+            internal_weights=False,
+            label="ir_mul_repr_5x1x5_l535",
+        ),
+        oeq.TPProblem(
+            "13x5e",
+            "1x3e",
+            "13x5e",
+            [(0, 0, 0, "uvu", True)],
+            shared_weights=False,
+            internal_weights=False,
+            label="ir_mul_repr_13x1x13_l535",
+        ),
+    ]
+
+
 class TestIrMulLayoutMACE(TPCorrectness):
-    production_model_tpps = mace_problems()
+    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):
@@ -284,6 +307,23 @@ def problem(self, request, dtype):
         return problem
 
 
+def test_ir_mul_rejects_uvw_problem(dtype):
+    problem = oeq.TPProblem(
+        "5x5e",
+        "1x3e",
+        "5x5e",
+        [(0, 0, 0, "uvw", True)],
+        shared_weights=False,
+        internal_weights=False,
+        irrep_dtype=dtype,
+        weight_dtype=dtype,
+        layout="ir_mul",
+    )
+
+    with pytest.raises(AssertionError, match="layout='ir_mul'"):
+        oeq.TensorProduct(problem)
+
+
 class TestTorchToSubmodule:
     """Test that TensorProduct works correctly as a submodule when parent's .to() is called"""
 

tests/conv_test.py

@@ -284,8 +284,31 @@ 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 [
+        oeq.TPProblem(
+            "5x5e",
+            "1x3e",
+            "5x5e",
+            [(0, 0, 0, "uvu", True)],
+            shared_weights=False,
+            internal_weights=False,
+            label="ir_mul_repr_5x1x5_l535",
+        ),
+        oeq.TPProblem(
+            "13x5e",
+            "1x3e",
+            "13x5e",
+            [(0, 0, 0, "uvu", True)],
+            shared_weights=False,
+            internal_weights=False,
+            label="ir_mul_repr_13x1x13_l535",
+        ),
+    ]
+
+
 class TestIrMulLayout(ConvCorrectness):
-    production_model_tpps = mace_problems()
+    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):