Update tripy/tests/integration/test_equal.py

Co-authored-by: pranavm-nvidia <[email protected]>
Signed-off-by: Jhalak Patel <[email protected]>

Author: jhalakpatel

tripy/tests/integration/test_equal.py

@@ -25,44 +25,23 @@
 
 class TestEqual:
     @pytest.mark.parametrize(
-        "tensor_a, tensor_b, dtype",
+        "tensor_a, tensor_b, dtype, expected",
         [
-            ([1, 2, 3], [1, 2, 3], tp.int64),
-            ([1, 2, 3], [1, 2, 4], tp.int32),
-            ([1.0, 2.0, 3.0], [1.0, 2.0, 3.0], tp.float16),
-            ([1.0, 2.0, 3.0], [1.0, 2.0, 1.0], tp.bfloat16),
-            ([1.0, 2.0, 3.0], [1.0, 2.0, 3.00001], tp.float32),
-            ([True, False, True], [True, False, True], tp.bool),
+            ([1, 2, 3], [1, 2, 3], tp.int64, True),
+            ([1, 2, 3], [1, 2, 4], tp.int32, False),
+            ([1.0, 2.0, 3.0], [1.0, 2.0, 3.0], tp.float16, True),
+            ([1.0, 2.0, 3.0], [1.0, 2.0, 1.0], tp.bfloat16, False),
+            ([1.0, 2.0, 3.0], [1.0, 2.0, 3.00001], tp.float32, False),
+            ([True, False, True], [True, False, True], tp.bool, True),
         ],
     )
-    def test_equal(self, tensor_a, tensor_b, dtype):
-        # Convert to torch tensors for comparison
-        torch_a = torch.tensor(tensor_a, dtype=self.torch_dtype(dtype))
-        torch_b = torch.tensor(tensor_b, dtype=self.torch_dtype(dtype))
+    def test_equal(self, tensor_a, tensor_b, dtype, expected):
+        a = tp.Tensor(tensor_a, dtype=dtype)
+        b = tp.Tensor(tensor_b, dtype=dtype)
 
-        # Convert to tripy tensors
-        tp_a = tp.Tensor(tensor_a, dtype=dtype)
-        tp_b = tp.Tensor(tensor_b, dtype=dtype)
+        out = tp.equal(a, b)
 
-        # Compare results
-        torch_result = torch.equal(torch_a, torch_b)
-        tp_result = tp.equal(tp_a, tp_b)
-
-        assert torch_result == tp_result
-
-    @staticmethod
-    def torch_dtype(tp_dtype):
-        # Map tripy dtypes to torch dtypes
-        dtype_map = {
-            tp.float32: torch.float32,
-            tp.float16: torch.float16,
-            tp.bfloat16: torch.bfloat16,
-            tp.int32: torch.int32,
-            tp.int64: torch.int64,
-            tp.int8: torch.int8,
-            tp.bool: torch.bool,
-        }
-        return dtype_map[tp_dtype]
+        assert out == expected
 
     def test_equal_different_shapes(self):
         a = tp.Tensor([1, 2, 3])