Add test case for condition where broadcast is not required

Author: sharavak

projects/pt1/e2e_testing/xfail_sets.py

@@ -2972,6 +2972,7 @@
     "HardtanhBackward_basic",
     "ElementwiseHeavisideModule_basic",
     "ElementwiseHeavisideIntModule_basic",
+    "ElementwiseHeavisideNoBroadcastModule_basic",
     "HstackBasicComplexModule_basic",
     "HstackBasicFloatModule_basic",
     "HstackBasicIntFloatModule_basic",
@@ -3989,6 +3990,7 @@
     "ElementwiseRreluWithNoiseTrainStaticModule_basic",
     "ElementwiseHeavisideModule_basic",
     "ElementwiseHeavisideIntModule_basic",
+    "ElementwiseHeavisideNoBroadcastModule_basic",
     "RreluWithNoiseBackwardEvalModule_basic",
     "RreluWithNoiseBackwardEvalStaticModule_basic",
     "RreluWithNoiseBackwardTrainModule_basic",

projects/pt1/python/torch_mlir_e2e_test/test_suite/elementwise_comparison.py

@@ -335,6 +335,26 @@ def ElementwiseHeavisideIntModule_basic(module, tu: TestUtils):
     )
 
 
+class ElementwiseHeavisideNoBroadcastModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args(
+        [None, ([-1, -1], torch.float32, True), ([-1, -1], torch.float32, True)]
+    )
+    def forward(self, x, values):
+        return torch.heaviside(x, values)
+
+
+@register_test_case(module_factory=lambda: ElementwiseHeavisideNoBroadcastModule())
+def ElementwiseHeavisideNoBroadcastModule_basic(module, tu: TestUtils):
+    module.forward(
+        tu.rand(5, 8),
+        tu.rand(5, 8),
+    )
+
+
 # ==============================================================================