Update broadcast tests xfails

Signed-off-by: Adrian Lundell <[email protected]>
Change-Id: I6224ca5a890556ca9cf0929440f7aea8da1b5028

Author: AdrianLundell

backends/cortex_m/ops/operators.py

@@ -138,6 +138,9 @@ def quantized_add_meta(
     output_multiplier: int,
     output_shift: int,
 ) -> torch.Tensor:
+    assert (
+        self.shape == other.shape
+    ), "Broadcasting not yet supported in Cortex-M backend."
     broadcasted_shape = torch.broadcast_shapes(self.shape, other.shape)
     return torch.empty(broadcasted_shape, dtype=torch.int8, device=self.device)
 
@@ -156,6 +159,9 @@ def quantized_add_impl(
     output_multiplier: int,
     output_shift: int,
 ) -> torch.Tensor:
+    assert (
+        self.shape == other.shape
+    ), "Broadcasting not yet supported in Cortex-M backend."
     self_shifted = (self.to(torch.int32) - self_zero_point) << SHIFT_INT8
     self_fp = requantize_cmsis(self_shifted, self_multiplier, self_shift)
 
@@ -197,6 +203,9 @@ def quantized_mul_meta(
     output_shift: int,
 ) -> torch.Tensor:
     # Broadcast to output shape
+    assert (
+        self.shape == other.shape
+    ), "Broadcasting not yet supported in Cortex-M backend."
     broadcasted_shape = torch.broadcast_shapes(self.shape, other.shape)
     return torch.empty(broadcasted_shape, dtype=torch.int8, device=self.device)
 
@@ -214,6 +223,9 @@ def quantized_mul_impl(
     # CMSIS-NN kernel multiplies raw int8 tensors (after zero-point offset) and
     # only uses the output multiplier/shift for rescaling. Mirror that here to
     # keep the composite implementation numerically aligned with the backend.
+    assert (
+        self.shape == other.shape
+    ), "Broadcasting not yet supported in Cortex-M backend."
     self_int = self.to(torch.int32) - self_zero_point
     other_int = other.to(torch.int32) - other_zero_point
     result_fp = self_int * other_int

backends/cortex_m/test/ops/test_add.py

@@ -139,7 +139,7 @@ class CortexMAlphaAdd(ModelAlpha):
 }
 
 
-dialect_xfails = {
+xfails = {
     "self_scalar": (
         "'float' object has not attribute 'fake_mode' - scalar only ops not supported.",
         AttributeError,
@@ -152,46 +152,21 @@ class CortexMAlphaAdd(ModelAlpha):
         "Expecting kwargs for aten op IR to be empty - alpha arg not supported.",
         AssertionError,
     ),
+    "broadcast_1": "Broadcasting not yet supported in Cortex-M backend",
+    "broadcast_2": "Broadcasting not yet supported in Cortex-M backend",
+    "broadcast_3": "Broadcasting not yet supported in Cortex-M backend",
 }
 
 
-@parametrize("test_case", test_cases, xfails=dialect_xfails)
+@parametrize("test_case", test_cases, xfails=xfails)
 def test_dialect_add(test_case):
     tester = CortexMTester(test_case.model, test_case.example_inputs)
     tester.test_dialect(
         test_case.model.ops_before_transforms, test_case.model.ops_after_transforms
     )
 
 
-implementation_xfails = {
-    "self_scalar": (
-        "'float' object has not attribute 'fake_mode' - scalar only ops not supported.",
-        AttributeError,
-    ),
-    "scalar_scalar": (
-        "'float' object has not attribute 'fake_mode' - scalar only ops not supported.",
-        AttributeError,
-    ),
-    "broadcast_1": (
-        " assert failed (input1.sizes() == input2.sizes()): Input1 and Input2 must have the same sizes.",
-        RuntimeError,
-    ),
-    "broadcast_2": (
-        " assert failed (input1.sizes() == input2.sizes()): Input1 and Input2 must have the same sizes.",
-        RuntimeError,
-    ),
-    "broadcast_3": (
-        " assert failed (input1.sizes() == input2.sizes()): Input1 and Input2 must have the same sizes.",
-        RuntimeError,
-    ),
-    "alpha": (
-        "Expecting kwargs for aten op IR to be empty - alpha arg not supported.",
-        AssertionError,
-    ),
-}
-
-
-@parametrize("test_case", test_cases, xfails=implementation_xfails)
+@parametrize("test_case", test_cases, xfails=xfails)
 def test_implementation_add(test_case):
     tester = CortexMTester(test_case.model, test_case.example_inputs)
     tester.test_implementation()

backends/cortex_m/test/ops/test_mul.py

@@ -59,16 +59,6 @@ class CortexMTensorMul(Model):
     }
 
 
-class CortexMTensorMulBroadCast(Model):
-    ops_before_transforms = {
-        "executorch_exir_dialects_edge__ops_aten_mul_Tensor": 1,
-    }
-
-    ops_after_transforms = {
-        "executorch_exir_dialects_edge__ops_aten_mul_Tensor": 1,
-    }
-
-
 test_cases = {
     "self_scalar": McuTestCase(
         CortexMSelfMul(),
@@ -107,15 +97,15 @@ class CortexMTensorMulBroadCast(Model):
         (1000.0, torch.ones(1)),
     ),
     "broadcast_1": McuTestCase(
-        CortexMTensorMulBroadCast(),
+        CortexMTensorMul(),
         (torch.ones(1), torch.ones(2, 2, 2, 2)),
     ),
     "broadcast_2": McuTestCase(
-        CortexMTensorMulBroadCast(),
+        CortexMTensorMul(),
         (torch.ones((2, 1, 1, 1)), torch.ones(1)),
     ),
     "broadcast_3": McuTestCase(
-        CortexMTensorMulBroadCast(),
+        CortexMTensorMul(),
         (
             ramp_tensor(-2, 2, (2, 1, 2, 1)),
             ramp_tensor(-5, 5, (1, 2, 1, 2)),
@@ -125,8 +115,17 @@ class CortexMTensorMulBroadCast(Model):
 
 
 xfail_cases = {
-    "self_scalar": "lift_constant_tensor_pass assumes fake tensors for scalars",
-    "scalar_scalar": "lift_constant_tensor_pass assumes fake tensors for scalars",
+    "self_scalar": (
+        "'float' object has not attribute 'fake_mode' - scalar only ops not supported.",
+        AttributeError,
+    ),
+    "scalar_scalar": (
+        "'float' object has not attribute 'fake_mode' - scalar only ops not supported.",
+        AttributeError,
+    ),
+    "broadcast_1": "Broadcasting not yet supported in Cortex-M backend",
+    "broadcast_2": "Broadcasting not yet supported in Cortex-M backend",
+    "broadcast_3": "Broadcasting not yet supported in Cortex-M backend",
 }