updating tests to use absltest

Author: rdyro

tests/test_fn_wrapping.py

@@ -1,6 +1,7 @@
 import sys
 from pathlib import Path
 
+from absl.testing import parameterized, absltest
 import torch
 from torch import Size
 import jax
@@ -15,105 +16,104 @@
 
 from utils import jax_randn  # noqa: E402
 from torch2jax import torch2jax  # noqa: E402
-from torch2jax.compat import torch2jax as _torch2jax_flat  # noqa: E402
 
 ####################################################################################################
 
 
-def test_single_output_fn():
-    shape = (10, 2)
-
-    def torch_fn(x, y):
-        return (x + 1 - y.reshape(x.shape)) / torch.norm(y)
-
-    device_list = ["cpu", "cuda"] if torch.cuda.is_available() else ["cpu"]
-    dtype_list = [jnp.float32, jnp.float64]
-
-    x = jax_randn(shape, device="cpu", dtype=jnp.float64)
-    y = jax_randn(shape, device="cpu", dtype=jnp.float64).reshape(-1)
-    jax_fn = torch2jax(torch_fn, x, y, output_shapes=Size(shape))
-
-    for device in device_list:
-        for dtype in dtype_list:
-            x = jax_randn(shape, device=device, dtype=dtype)
-            y = jax_randn(shape, device=device, dtype=dtype).reshape(-1)
-
-            # non-jit version
-            out = jax_fn(x, y)
-            assert isinstance(out, Array)
-            expected = (x + 1 - y.reshape(x.shape)) / jnp.linalg.norm(y)
-            err = jnp.linalg.norm(out - expected) / jnp.linalg.norm(expected)
-            assert err < 1e-5
-
-            # jit version
-            @jax.jit
-            def complication_fn(x, y):
-                a = jax_fn(x, y)
-                y2 = y.reshape(x.shape)
-                b, c = x - y2 + 1, x + y2 + 1
-                d = jnp.linalg.norm(x) - jnp.linalg.norm(y)
-                return a, b, c, d
-
-            out = complication_fn(x, y)
-            assert isinstance(out, (list, tuple)) and len(out) == 4
-            out1 = out[0]
-            assert isinstance(out1, Array)
-            expected = (x + 1 - y.reshape(x.shape)) / jnp.linalg.norm(y)
-            err = jnp.linalg.norm(out1 - expected) / jnp.linalg.norm(expected)
-            assert err < 1e-5, f"Error is quite high: {err:.4e}"
-
-
-def test_multi_output_fn():
-    shape = (10, 2)
-
-    def torch_fn(x, y):
-        a = (x + 1 - y.reshape(x.shape)) / torch.norm(y)
-        b = (x - y.reshape(x.shape)).reshape(-1)[:5]
-        return a, b
-
-    device_list = ["cpu", "cuda"] if torch.cuda.is_available() else ["cpu"]
-    dtype_list = [jnp.float32, jnp.float64]
-    x = jax_randn(shape, device="cpu", dtype=jnp.float64)
-    y = jax_randn(shape, device="cpu", dtype=jnp.float64).reshape(-1)
-    jax_fn = torch2jax(torch_fn, x, y, output_shapes=(Size(shape), Size((5,))))
-
-    for device in device_list:
-        for dtype in dtype_list:
-            x = jax_randn(shape, device=device, dtype=dtype)
-            y = jax_randn(shape, device=device, dtype=dtype).reshape(-1)
-
-            # non-jit version
-            out = jax_fn(x, y)
-            assert isinstance(out, (list, tuple)) and len(out) == 2
-            assert all(isinstance(z, Array) for z in out)
-            expected1 = (x + 1 - y.reshape(x.shape)) / jnp.linalg.norm(y)
-            expected2 = (x - y.reshape(x.shape)).reshape(-1)[:5]
-            err1 = jnp.linalg.norm(out[0] - expected1) / jnp.linalg.norm(expected1)
-            err2 = jnp.linalg.norm(out[1] - expected2) / jnp.linalg.norm(expected2)
-            assert err1 < 1e-5 and err2 < 1e-5
-
-            # jit version
-            @jax.jit
-            def complication_fn(x, y):
-                a = jax_fn(x, y)
-                y2 = y.reshape(x.shape)
-                b, c = x - y2 + 1, x + y2 + 1
-                d = jnp.linalg.norm(x) - jnp.linalg.norm(y)
-                return a, b, c, d
-
-            out = complication_fn(x, y)
-            assert isinstance(out, (list, tuple)) and len(out) == 4
-            assert all(isinstance(z, Array) for z in out[0])
-            out1 = out[0]
-            expected1 = (x + 1 - y.reshape(x.shape)) / jnp.linalg.norm(y)
-            expected2 = (x - y.reshape(x.shape)).reshape(-1)[:5]
-            err1 = jnp.linalg.norm(out1[0] - expected1) / jnp.linalg.norm(expected2)
-            err2 = jnp.linalg.norm(out1[1] - expected2) / jnp.linalg.norm(expected2)
-            assert err1 < 1e-5 and err2 < 1e-5
+class TestFnWrapping(parameterized.TestCase):
+    @parameterized.product(device=["cpu", "cuda"], dtype=[jnp.float32, jnp.float64])
+    def test_single_output_fn(self, device, dtype):
+        if not torch.cuda.is_available() and device == "cuda":
+            self.skipTest("Skipping CUDA test when CUDA is not available.")
+        shape = (10, 2)
+
+        def torch_fn(x, y):
+            return (x + 1 - y.reshape(x.shape)) / torch.norm(y)
+
+        device_list = ["cpu", "cuda"] if torch.cuda.is_available() else ["cpu"]
+        dtype_list = [jnp.float32, jnp.float64]
+
+        x = jax_randn(shape, device="cpu", dtype=jnp.float64)
+        y = jax_randn(shape, device="cpu", dtype=jnp.float64).reshape(-1)
+        jax_fn = torch2jax(torch_fn, x, y, output_shapes=Size(shape))
+
+        x = jax_randn(shape, device=device, dtype=dtype)
+        y = jax_randn(shape, device=device, dtype=dtype).reshape(-1)
+
+        # non-jit version
+        out = jax_fn(x, y)
+        assert isinstance(out, Array)
+        expected = (x + 1 - y.reshape(x.shape)) / jnp.linalg.norm(y)
+        err = jnp.linalg.norm(out - expected) / jnp.linalg.norm(expected)
+        assert err < 1e-5
+
+        # jit version
+        @jax.jit
+        def complication_fn(x, y):
+            a = jax_fn(x, y)
+            y2 = y.reshape(x.shape)
+            b, c = x - y2 + 1, x + y2 + 1
+            d = jnp.linalg.norm(x) - jnp.linalg.norm(y)
+            return a, b, c, d
+
+        out = complication_fn(x, y)
+        assert isinstance(out, (list, tuple)) and len(out) == 4
+        out1 = out[0]
+        assert isinstance(out1, Array)
+        expected = (x + 1 - y.reshape(x.shape)) / jnp.linalg.norm(y)
+        err = jnp.linalg.norm(out1 - expected) / jnp.linalg.norm(expected)
+        assert err < 1e-5, f"Error is quite high: {err:.4e}"
+
+    @parameterized.product(device=["cpu", "cuda"], dtype=[jnp.float32, jnp.float64])
+    def test_multi_output_fn(self, device, dtype):
+        if not torch.cuda.is_available() and device == "cuda":
+            self.skipTest("Skipping CUDA test when CUDA is not available.")
+
+        shape = (10, 2)
+
+        def torch_fn(x, y):
+            a = (x + 1 - y.reshape(x.shape)) / torch.norm(y)
+            b = (x - y.reshape(x.shape)).reshape(-1)[:5]
+            return a, b
+
+        x = jax_randn(shape, device="cpu", dtype=jnp.float64)
+        y = jax_randn(shape, device="cpu", dtype=jnp.float64).reshape(-1)
+        jax_fn = torch2jax(torch_fn, x, y, output_shapes=(Size(shape), Size((5,))))
+
+        x = jax_randn(shape, device=device, dtype=dtype)
+        y = jax_randn(shape, device=device, dtype=dtype).reshape(-1)
+
+        # non-jit version
+        out = jax_fn(x, y)
+        assert isinstance(out, (list, tuple)) and len(out) == 2
+        assert all(isinstance(z, Array) for z in out)
+        expected1 = (x + 1 - y.reshape(x.shape)) / jnp.linalg.norm(y)
+        expected2 = (x - y.reshape(x.shape)).reshape(-1)[:5]
+        err1 = jnp.linalg.norm(out[0] - expected1) / jnp.linalg.norm(expected1)
+        err2 = jnp.linalg.norm(out[1] - expected2) / jnp.linalg.norm(expected2)
+        assert err1 < 1e-5 and err2 < 1e-5
+
+        # jit version
+        @jax.jit
+        def complication_fn(x, y):
+            a = jax_fn(x, y)
+            y2 = y.reshape(x.shape)
+            b, c = x - y2 + 1, x + y2 + 1
+            d = jnp.linalg.norm(x) - jnp.linalg.norm(y)
+            return a, b, c, d
+
+        out = complication_fn(x, y)
+        assert isinstance(out, (list, tuple)) and len(out) == 4
+        assert all(isinstance(z, Array) for z in out[0])
+        out1 = out[0]
+        expected1 = (x + 1 - y.reshape(x.shape)) / jnp.linalg.norm(y)
+        expected2 = (x - y.reshape(x.shape)).reshape(-1)[:5]
+        err1 = jnp.linalg.norm(out1[0] - expected1) / jnp.linalg.norm(expected2)
+        err2 = jnp.linalg.norm(out1[1] - expected2) / jnp.linalg.norm(expected2)
+        assert err1 < 1e-5 and err2 < 1e-5
 
 
 ####################################################################################################
 
 if __name__ == "__main__":
-    test_single_output_fn()
-    test_multi_output_fn()
+    absltest.main()

tests/test_grad_fallback.py

@@ -3,6 +3,7 @@
 import sys
 from pathlib import Path
 
+from absl.testing import parameterized, absltest
 import torch
 from torch.autograd import Function
 import jax
@@ -19,76 +20,78 @@
 ####################################################################################################
 
 
-def test_torch2jax_with_vjp_vjp_fallback():
-    shape = (5, 7)
+class TestGradFallback(parameterized.TestCase):
+    @parameterized.product(device=["cpu", "cuda"], dtype=[jnp.float32, jnp.float64])
+    def test_torch2jax_with_vjp_vjp_fallback(self, device, dtype):
+        if device == "cuda" and not torch.cuda.is_available():
+            self.skipTest("Skipping CUDA tests when CUDA is not available")
+        shape = (5, 7)
 
-    class OldInterfaceFunction(Function):
-        @staticmethod
-        def forward(ctx, x, y):
-            ctx.save_for_backward(x, y)
-            return x**2 + y**2
+        class OldInterfaceFunction(Function):
+            @staticmethod
+            def forward(ctx, x, y):
+                ctx.save_for_backward(x, y)
+                return x**2 + y**2
 
-        @staticmethod
-        def backward(ctx, grad_output):
-            x, y = ctx.saved_tensors
-            return 2 * x * grad_output, 2 * y * grad_output
+            @staticmethod
+            def backward(ctx, grad_output):
+                x, y = ctx.saved_tensors
+                return 2 * x * grad_output, 2 * y * grad_output
 
-    def torch_fn(x, y):
-        return OldInterfaceFunction.apply(x, y)
+        def torch_fn(x, y):
+            return OldInterfaceFunction.apply(x, y)
 
-    def expected_f_fn(x, y):
-        return x**2 + y**2
+        def expected_f_fn(x, y):
+            return x**2 + y**2
 
-    expected_g_fn = jax.grad(lambda *args: jnp.sum(expected_f_fn(*args)), argnums=(0, 1))
-    expected_h_fn = jax.grad(lambda *args: jnp.sum(expected_g_fn(*args)[0] + expected_g_fn(*args)[1]), argnums=(0, 1))
+        expected_g_fn = jax.grad(lambda *args: jnp.sum(expected_f_fn(*args)), argnums=(0, 1))
+        expected_h_fn = jax.grad(
+            lambda *args: jnp.sum(expected_g_fn(*args)[0] + expected_g_fn(*args)[1]), argnums=(0, 1)
+        )
 
-    xt, yt = torch.randn(shape), torch.randn(shape)
-    device_list = ["cuda", "cpu"] if torch.cuda.is_available() else ["cpu"]
-    dtype_list = [jnp.float32, jnp.float64]
+        xt, yt = torch.randn(shape), torch.randn(shape)
 
-    wrap_jax_f_fn = torch2jax_with_vjp(torch_fn, xt, yt, depth=2)
-    wrap_jax_g_fn = jax.grad(lambda x, y: jnp.sum(wrap_jax_f_fn(x, y)), argnums=(0, 1))
-    wrap_jax_h_fn = jax.grad(lambda x, y: jnp.sum(wrap_jax_g_fn(x, y)[0] + wrap_jax_g_fn(x, y)[1]), argnums=(0, 1))
+        wrap_jax_f_fn = torch2jax_with_vjp(torch_fn, xt, yt, depth=2)
+        wrap_jax_g_fn = jax.grad(lambda x, y: jnp.sum(wrap_jax_f_fn(x, y)), argnums=(0, 1))
+        wrap_jax_h_fn = jax.grad(lambda x, y: jnp.sum(wrap_jax_g_fn(x, y)[0] + wrap_jax_g_fn(x, y)[1]), argnums=(0, 1))
 
-    for device in device_list:
-        for dtype in dtype_list:
-            x = jax_randn(shape, dtype=dtype, device=device)
-            y = jax_randn(shape, dtype=dtype, device=device)
-            f = wrap_jax_f_fn(x, y)
-            g = wrap_jax_g_fn(x, y)
-            h = wrap_jax_h_fn(x, y)
+        x = jax_randn(shape, dtype=dtype, device=device)
+        y = jax_randn(shape, dtype=dtype, device=device)
+        f = wrap_jax_f_fn(x, y)
+        g = wrap_jax_g_fn(x, y)
+        h = wrap_jax_h_fn(x, y)
 
-            f_expected = expected_f_fn(x, y)
-            g_expected = expected_g_fn(x, y)
-            h_expected = expected_h_fn(x, y)
+        f_expected = expected_f_fn(x, y)
+        g_expected = expected_g_fn(x, y)
+        h_expected = expected_h_fn(x, y)
 
-            # test output structure #############################
-            assert isinstance(f, Array)
-            assert isinstance(g, (tuple, list)) and len(g) == 2 and isinstance(g[0], Array) and isinstance(g[1], Array)
-            assert isinstance(h, (tuple, list)) and len(h) == 2 and isinstance(h[0], Array) and isinstance(h[1], Array)
+        # test output structure #############################
+        self.assertIsInstance(f, Array)
+        assert isinstance(g, (tuple, list)) and len(g) == 2 and isinstance(g[0], Array) and isinstance(g[1], Array)
+        assert isinstance(h, (tuple, list)) and len(h) == 2 and isinstance(h[0], Array) and isinstance(h[1], Array)
 
-            # test values not under JIT #########################
-            err_f = jnp.linalg.norm(f - f_expected)
-            err_g = jnp.linalg.norm(g[0] - g_expected[0]) + jnp.linalg.norm(g[1] - g_expected[1])
-            err_h = jnp.linalg.norm(h[0] - h_expected[0]) + jnp.linalg.norm(h[1] - h_expected[1])
+        # test values not under JIT #########################
+        err_f = jnp.linalg.norm(f - f_expected)
+        err_g = jnp.linalg.norm(g[0] - g_expected[0]) + jnp.linalg.norm(g[1] - g_expected[1])
+        err_h = jnp.linalg.norm(h[0] - h_expected[0]) + jnp.linalg.norm(h[1] - h_expected[1])
 
-            assert err_f < 1e-5, f"Error in f value is {err_f:.4e}"
-            assert err_g < 1e-5, f"Error in g value is {err_g:.4e}"
-            assert err_h < 1e-5, f"Error in h value is {err_h:.4e}"
+        self.assertLess(err_f, 1e-5)
+        self.assertLess(err_g, 1e-5)
+        self.assertLess(err_h, 1e-5)
 
-            # test values when under JIT ########################
-            f = jax.jit(wrap_jax_f_fn)(x, y)
-            g = jax.jit(wrap_jax_g_fn)(x, y)
-            h = jax.jit(wrap_jax_h_fn)(x, y)
+        # test values when under JIT ########################
+        f = jax.jit(wrap_jax_f_fn)(x, y)
+        g = jax.jit(wrap_jax_g_fn)(x, y)
+        h = jax.jit(wrap_jax_h_fn)(x, y)
 
-            err_f = jnp.linalg.norm(f - f_expected)
-            err_g = jnp.linalg.norm(g[0] - g_expected[0]) + jnp.linalg.norm(g[1] - g_expected[1])
-            err_h = jnp.linalg.norm(h[0] - h_expected[0]) + jnp.linalg.norm(h[1] - h_expected[1])
+        err_f = jnp.linalg.norm(f - f_expected)
+        err_g = jnp.linalg.norm(g[0] - g_expected[0]) + jnp.linalg.norm(g[1] - g_expected[1])
+        err_h = jnp.linalg.norm(h[0] - h_expected[0]) + jnp.linalg.norm(h[1] - h_expected[1])
 
-            assert err_f < 1e-5, f"Error in f value is {err_f:.4e}"
-            assert err_g < 1e-5, f"Error in g value is {err_g:.4e}"
-            assert err_h < 1e-5, f"Error in h value is {err_h:.4e}"
+        self.assertLess(err_f, 1e-5)
+        self.assertLess(err_g, 1e-5)
+        self.assertLess(err_h, 1e-5)
 
 
 if __name__ == "__main__":
-    test_torch2jax_with_vjp_vjp_fallback()
+    absltest.main()