More test cleaning.

Author: vbharadwaj-bk

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
 
-
-def ir_mul_representative_uvu_problems():
-    return [
+class TestIrMulLayoutMACE(TPCorrectness):
+    production_model_tpps = mace_problems() + \
+        [
         oeq.TPProblem(
             "5x5e",
             "1x3e",
@@ -293,11 +293,7 @@ def ir_mul_representative_uvu_problems():
             internal_weights=False,
             label="ir_mul_repr_13x1x13_l535",
         ),
-    ]
-
-
-class TestIrMulLayoutMACE(TPCorrectness):
-    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):
@@ -307,23 +303,6 @@ 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/input_validation_test.py

@@ -138,3 +138,20 @@ def test_cpp_checks_forward_dtype(executable_and_buffers, subtests):
             with pytest.raises(RuntimeError, match=r"Dtype mismatch"):
                 buffers[i] = buffers[i].to(dtype=torch.bfloat16)
                 executable(*buffers)
+
+
+def test_ir_mul_rejects_uvw_problem(dtype):
+    problem = 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'"):
+        TensorProduct(problem)