enable test_functional; failing many

tripy/tests/integration/test_functional.py

@@ -30,14 +30,17 @@
 
 class TestFunctional:
     @pytest.mark.parametrize("kind", ["cpu", "gpu"])
-    def test_add_two_tensors(self, kind):
+    def test_add_two_tensors(self, kind, compile_fixture):
         module = cp if kind == "gpu" else np
         arr = module.array([2, 3], dtype=np.float32)
         a = tp.Tensor(arr, device=tp.device(kind))
         b = tp.Tensor(module.ones(2, dtype=module.float32), device=tp.device(kind))
 
-        c = a + b
-        out = c + c
+        def func(a, b):
+            c = a + b
+            return c + c
+
+        out = compile_fixture(func, a, b)
         assert (cp.from_dlpack(out).get() == np.array([6.0, 8.0], dtype=np.float32)).all()
 
     @pytest.mark.parametrize(
@@ -49,7 +52,7 @@ def test_add_two_tensors(self, kind):
             ((1, 3, 7), (4, 3, 1)),  # broadcast at differnt dim of both operand
         ],
     )
-    def test_static_broadcast_add_two_tensors(self, dim_a, dim_b):
+    def test_static_broadcast_add_two_tensors(self, dim_a, dim_b, compile_fixture):
         cp_a = cp.arange(np.prod(dim_a)).reshape(dim_a).astype(np.float32)
         cp_b = cp.arange(np.prod(dim_b)).reshape(dim_b).astype(np.float32)
         a = tp.Tensor(cp_a, shape=dim_a, device=tp.device("gpu"))
@@ -59,7 +62,7 @@ def func(a, b):
             c = a + b
             return c
 
-        out = func(a, b)
+        out = compile_fixture(func, a, b)
         assert (cp.from_dlpack(out) == cp.array(cp_a + cp_b)).all()
 
     def test_multi_output_trace(self):
@@ -179,36 +182,36 @@ def test_weights_loading_from_torch(self, kind):
 class TestCopyFunctional:
     @pytest.mark.parametrize("src", ["cpu", "gpu"])
     @pytest.mark.parametrize("dst", ["cpu", "gpu"])
-    def test_single_copy(self, src, dst):
+    def test_single_copy(self, src, dst, compile_fixture):
         a = tp.Tensor([1, 2], device=tp.device(src))
-        out = tp.copy(a, tp.device(dst))
-        out = out.eval()
+        out = compile_fixture(tp.copy, a, tp.device(dst))
         assert out.device.kind == dst
-        assert out.data() == [1, 2]
+        assert out == [1, 2]
 
-    def test_multiple_copy_1(self):
+    def test_multiple_copy_1(self, compile_fixture):
         a = tp.Tensor([1, 2])
-        a = tp.copy(a, tp.device("gpu"))
-        a = tp.copy(a, tp.device("cpu"))
-        out = a.eval()
-        assert out.device.kind == "cpu"
-        assert out.data() == [1, 2]
+        a = compile_fixture(a, tp.device("gpu"))
+        a = compile_fixture(a, tp.device("cpu"))
+        assert a.device.kind == "cpu"
+        assert a.data() == [1, 2]
 
-    def test_multiple_copy_2(self):
+    def test_multiple_copy_2(self, compile_fixture):
         a = tp.Tensor([1, 2])
-        a = tp.copy(a, tp.device("cpu"))
-        a = tp.copy(a, tp.device("gpu"))
-        out = a.eval()
-        assert out.device.kind == "gpu"
-        assert out.data() == [1, 2]
+        a = compile_fixture(tp.copy, a, tp.device("cpu"))
+        a = compile_fixture(tp.copy, a, tp.device("gpu"))
+        assert a.device.kind == "gpu"
+        assert a.data() == [1, 2]
 
     @pytest.mark.skip("Remove copy op in Tripy: mlir-tensorrt #756")
-    def test_with_ops(self):
+    def test_with_ops(self, compile_fixture):
         a = tp.Tensor([1, 2])
         b = tp.Tensor([2, 3])
-        out = a + b
-        out = tp.copy(out, tp.device("cpu"))
-        out = out.eval()
+
+        def func(a, b):
+            out = a + b
+            out = tp.copy(out, tp.device("cpu"))
+
+        out = compile_fixture(func, a, b)
         assert out.device.kind == "cpu"
         assert out.data() == [3, 5]