warp-lang 1.9.0__py3-none-win_amd64.whl → 1.10.0rc2__py3-none-win_amd64.whl

warp/tests/interop/test_jax.py

@@ -15,13 +15,18 @@
 
 import os
 import unittest
+from functools import partial
 from typing import Any
 
 import numpy as np
 
 import warp as wp
+from warp._src.jax import get_jax_device
 from warp.tests.unittest_utils import *
 
+# default array size for tests
+ARRAY_SIZE = 1024 * 1024
+
 
 # basic kernel with one input and output
 @wp.kernel
@@ -44,6 +49,18 @@ def triple_kernel_vecmat(input: wp.array(dtype=Any), output: wp.array(dtype=Any)
     output[tid] = input.dtype.dtype(3) * input[tid]
 
 
+@wp.kernel
+def inc_1d_kernel(x: wp.array(dtype=float), y: wp.array(dtype=float)):
+    tid = wp.tid()
+    y[tid] = x[tid] + 1.0
+
+
+@wp.kernel
+def inc_2d_kernel(x: wp.array2d(dtype=float), y: wp.array2d(dtype=float)):
+    i, j = wp.tid()
+    y[i, j] = x[i, j] + 1.0
+
+
 # kernel with multiple inputs and outputs
 @wp.kernel
 def multiarg_kernel(
@@ -61,7 +78,7 @@ def multiarg_kernel(
 
 
 # various types for testing
-scalar_types = wp.types.scalar_types
+scalar_types = wp._src.types.scalar_types
 vector_types = []
 matrix_types = []
 for dim in [2, 3, 4]:
@@ -132,15 +149,19 @@ def test_device_conversion(test, device):
     test.assertEqual(warp_device, device)
 
 
-@unittest.skipUnless(_jax_version() >= (0, 4, 25), "Jax version too old")
-def test_jax_kernel_basic(test, device):
+def test_jax_kernel_basic(test, device, use_ffi=False):
     import jax.numpy as jp
 
-    from warp.jax_experimental import jax_kernel
+    if use_ffi:
+        from warp.jax_experimental.ffi import jax_kernel
 
-    n = 64
+        jax_triple = jax_kernel(triple_kernel)
+    else:
+        from warp.jax_experimental.custom_call import jax_kernel
+
+        jax_triple = jax_kernel(triple_kernel, quiet=True)  # suppress deprecation warnings
 
-    jax_triple = jax_kernel(triple_kernel)
+    n = ARRAY_SIZE
 
     @jax.jit
     def f():
@@ -151,18 +172,27 @@ def test_jax_kernel_basic(test, device):
     with jax.default_device(wp.device_to_jax(device)):
         y = f()
 
+    wp.synchronize_device(device)
+
     result = np.asarray(y).reshape((n,))
     expected = 3 * np.arange(n, dtype=np.float32)
 
     assert_np_equal(result, expected)
 
 
-@unittest.skipUnless(_jax_version() >= (0, 4, 25), "Jax version too old")
-def test_jax_kernel_scalar(test, device):
+def test_jax_kernel_scalar(test, device, use_ffi=False):
     import jax.numpy as jp
 
-    from warp.jax_experimental import jax_kernel
+    if use_ffi:
+        from warp.jax_experimental.ffi import jax_kernel
+
+        kwargs = {}
+    else:
+        from warp.jax_experimental.custom_call import jax_kernel
 
+        kwargs = {"quiet": True}
+
+    # use a smallish size to ensure arange * 3 doesn't overflow
     n = 64
 
     for T in scalar_types:
@@ -173,7 +203,7 @@ def test_jax_kernel_scalar(test, device):
             # get the concrete overload
             kernel_instance = triple_kernel_scalar.add_overload([wp.array(dtype=T), wp.array(dtype=T)])
 
-            jax_triple = jax_kernel(kernel_instance)
+            jax_triple = jax_kernel(kernel_instance, **kwargs)
 
             @jax.jit
             def f(jax_triple=jax_triple, jp_dtype=jp_dtype):
@@ -184,22 +214,31 @@ def test_jax_kernel_scalar(test, device):
             with jax.default_device(wp.device_to_jax(device)):
                 y = f()
 
+            wp.synchronize_device(device)
+
             result = np.asarray(y).reshape((n,))
             expected = 3 * np.arange(n, dtype=np_dtype)
 
             assert_np_equal(result, expected)
 
 
-@unittest.skipUnless(_jax_version() >= (0, 4, 25), "Jax version too old")
-def test_jax_kernel_vecmat(test, device):
+def test_jax_kernel_vecmat(test, device, use_ffi=False):
     import jax.numpy as jp
 
-    from warp.jax_experimental import jax_kernel
+    if use_ffi:
+        from warp.jax_experimental.ffi import jax_kernel
+
+        kwargs = {}
+    else:
+        from warp.jax_experimental.custom_call import jax_kernel
+
+        kwargs = {"quiet": True}
 
     for T in [*vector_types, *matrix_types]:
         jp_dtype = wp.dtype_to_jax(T._wp_scalar_type_)
         np_dtype = wp.dtype_to_numpy(T._wp_scalar_type_)
 
+        # use a smallish size to ensure arange * 3 doesn't overflow
         n = 64 // T._length_
         scalar_shape = (n, *T._shape_)
         scalar_len = n * T._length_
@@ -208,7 +247,7 @@ def test_jax_kernel_vecmat(test, device):
             # get the concrete overload
             kernel_instance = triple_kernel_vecmat.add_overload([wp.array(dtype=T), wp.array(dtype=T)])
 
-            jax_triple = jax_kernel(kernel_instance)
+            jax_triple = jax_kernel(kernel_instance, **kwargs)
 
             @jax.jit
             def f(jax_triple=jax_triple, jp_dtype=jp_dtype, scalar_len=scalar_len, scalar_shape=scalar_shape):
@@ -219,21 +258,27 @@ def test_jax_kernel_vecmat(test, device):
             with jax.default_device(wp.device_to_jax(device)):
                 y = f()
 
+            wp.synchronize_device(device)
+
             result = np.asarray(y).reshape(scalar_shape)
             expected = 3 * np.arange(scalar_len, dtype=np_dtype).reshape(scalar_shape)
 
             assert_np_equal(result, expected)
 
 
-@unittest.skipUnless(_jax_version() >= (0, 4, 25), "Jax version too old")
-def test_jax_kernel_multiarg(test, device):
+def test_jax_kernel_multiarg(test, device, use_ffi=False):
     import jax.numpy as jp
 
-    from warp.jax_experimental import jax_kernel
+    if use_ffi:
+        from warp.jax_experimental.ffi import jax_kernel
 
-    n = 64
+        jax_multiarg = jax_kernel(multiarg_kernel, num_outputs=2)
+    else:
+        from warp.jax_experimental.custom_call import jax_kernel
 
-    jax_multiarg = jax_kernel(multiarg_kernel)
+        jax_multiarg = jax_kernel(multiarg_kernel, quiet=True)
+
+    n = ARRAY_SIZE
 
     @jax.jit
     def f():
@@ -246,6 +291,8 @@ def test_jax_kernel_multiarg(test, device):
     with jax.default_device(wp.device_to_jax(device)):
         x, y = f()
 
+    wp.synchronize_device(device)
+
     result_x, result_y = np.asarray(x), np.asarray(y)
     expected_x = np.full(n, 3, dtype=np.float32)
     expected_y = np.full(n, 5, dtype=np.float32)
@@ -254,50 +301,48 @@ def test_jax_kernel_multiarg(test, device):
     assert_np_equal(result_y, expected_y)
 
 
-@unittest.skipUnless(_jax_version() >= (0, 4, 25), "Jax version too old")
-def test_jax_kernel_launch_dims(test, device):
+def test_jax_kernel_launch_dims(test, device, use_ffi=False):
     import jax.numpy as jp
 
-    from warp.jax_experimental import jax_kernel
+    if use_ffi:
+        from warp.jax_experimental.ffi import jax_kernel
+
+        kwargs = {}
+    else:
+        from warp.jax_experimental.custom_call import jax_kernel
+
+        kwargs = {"quiet": True}
 
     n = 64
     m = 32
 
     # Test with 1D launch dims
-    @wp.kernel
-    def add_one_kernel(x: wp.array(dtype=float), y: wp.array(dtype=float)):
-        tid = wp.tid()
-        y[tid] = x[tid] + 1.0
-
-    jax_add_one = jax_kernel(
-        add_one_kernel, launch_dims=(n - 2,)
+    jax_inc_1d = jax_kernel(
+        inc_1d_kernel, launch_dims=(n - 2,), **kwargs
     )  # Intentionally not the same as the first dimension of the input
 
     @jax.jit
     def f_1d():
         x = jp.arange(n, dtype=jp.float32)
-        return jax_add_one(x)
+        return jax_inc_1d(x)
 
     # Test with 2D launch dims
-    @wp.kernel
-    def add_one_2d_kernel(x: wp.array2d(dtype=float), y: wp.array2d(dtype=float)):
-        i, j = wp.tid()
-        y[i, j] = x[i, j] + 1.0
-
-    jax_add_one_2d = jax_kernel(
-        add_one_2d_kernel, launch_dims=(n - 2, m - 2)
+    jax_inc_2d = jax_kernel(
+        inc_2d_kernel, launch_dims=(n - 2, m - 2), **kwargs
     )  # Intentionally not the same as the first dimension of the input
 
     @jax.jit
     def f_2d():
         x = jp.zeros((n, m), dtype=jp.float32) + 3.0
-        return jax_add_one_2d(x)
+        return jax_inc_2d(x)
 
     # run on the given device
     with jax.default_device(wp.device_to_jax(device)):
         y_1d = f_1d()
         y_2d = f_2d()
 
+    wp.synchronize_device(device)
+
     result_1d = np.asarray(y_1d).reshape((n - 2,))
     expected_1d = np.arange(n - 2, dtype=np.float32) + 1.0
 
@@ -308,57 +353,1315 @@ def test_jax_kernel_launch_dims(test, device):
     assert_np_equal(result_2d, expected_2d)
 
 
-class TestJax(unittest.TestCase):
-    pass
+# =========================================================================================================
+# JAX FFI
+# =========================================================================================================
 
 
-# try adding Jax tests if Jax is installed correctly
-try:
-    # prevent Jax from gobbling up GPU memory
-    os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"
-    os.environ["XLA_PYTHON_CLIENT_ALLOCATOR"] = "platform"
+@wp.kernel
+def add_kernel(a: wp.array(dtype=float), b: wp.array(dtype=float), output: wp.array(dtype=float)):
+    tid = wp.tid()
+    output[tid] = a[tid] + b[tid]
 
-    import jax
 
-    # NOTE: we must enable 64-bit types in Jax to test the full gamut of types
-    jax.config.update("jax_enable_x64", True)
+@wp.kernel
+def axpy_kernel(x: wp.array(dtype=float), y: wp.array(dtype=float), alpha: float, out: wp.array(dtype=float)):
+    tid = wp.tid()
+    out[tid] = alpha * x[tid] + y[tid]
 
-    # check which Warp devices work with Jax
-    # CUDA devices may fail if Jax cannot find a CUDA Toolkit
-    test_devices = get_test_devices()
-    jax_compatible_devices = []
-    jax_compatible_cuda_devices = []
-    for d in test_devices:
-        try:
-            with jax.default_device(wp.device_to_jax(d)):
-                j = jax.numpy.arange(10, dtype=jax.numpy.float32)
-                j += 1
-            jax_compatible_devices.append(d)
-            if d.is_cuda:
-                jax_compatible_cuda_devices.append(d)
-        except Exception as e:
-            print(f"Skipping Jax DLPack tests on device '{d}' due to exception: {e}")
 
-    add_function_test(TestJax, "test_dtype_from_jax", test_dtype_from_jax, devices=None)
-    add_function_test(TestJax, "test_dtype_to_jax", test_dtype_to_jax, devices=None)
+@wp.kernel
+def sincos_kernel(angle: wp.array(dtype=float), sin_out: wp.array(dtype=float), cos_out: wp.array(dtype=float)):
+    tid = wp.tid()
+    sin_out[tid] = wp.sin(angle[tid])
+    cos_out[tid] = wp.cos(angle[tid])
 
-    if jax_compatible_devices:
-        add_function_test(TestJax, "test_device_conversion", test_device_conversion, devices=jax_compatible_devices)
 
-    if jax_compatible_cuda_devices:
-        add_function_test(TestJax, "test_jax_kernel_basic", test_jax_kernel_basic, devices=jax_compatible_cuda_devices)
-        add_function_test(
-            TestJax, "test_jax_kernel_scalar", test_jax_kernel_scalar, devices=jax_compatible_cuda_devices
-        )
-        add_function_test(
-            TestJax, "test_jax_kernel_vecmat", test_jax_kernel_vecmat, devices=jax_compatible_cuda_devices
-        )
-        add_function_test(
-            TestJax, "test_jax_kernel_multiarg", test_jax_kernel_multiarg, devices=jax_compatible_cuda_devices
-        )
+@wp.kernel
+def diagonal_kernel(output: wp.array(dtype=wp.mat33)):
+    tid = wp.tid()
+    d = float(tid + 1)
+    output[tid] = wp.mat33(d, 0.0, 0.0, 0.0, d * 2.0, 0.0, 0.0, 0.0, d * 3.0)
+
+
+@wp.kernel
+def scale_kernel(a: wp.array(dtype=float), s: float, output: wp.array(dtype=float)):
+    tid = wp.tid()
+    output[tid] = a[tid] * s
+
+
+@wp.kernel
+def scale_vec_kernel(a: wp.array(dtype=wp.vec2), s: float, output: wp.array(dtype=wp.vec2)):
+    tid = wp.tid()
+    output[tid] = a[tid] * s
+
+
+@wp.kernel
+def accum_kernel(a: wp.array(dtype=float), b: wp.array(dtype=float)):
+    tid = wp.tid()
+    b[tid] += a[tid]
+
+
+@wp.kernel
+def matmul_kernel(
+    a: wp.array2d(dtype=float),  # NxK
+    b: wp.array2d(dtype=float),  # KxM
+    c: wp.array2d(dtype=float),  # NxM
+):
+    # launch dims should be (N, M)
+    i, j = wp.tid()
+    N = a.shape[0]
+    K = a.shape[1]
+    M = b.shape[1]
+    if i < N and j < M:
+        s = wp.float32(0)
+        for k in range(K):
+            s += a[i, k] * b[k, j]
+        c[i, j] = s
+
+
+@wp.kernel
+def in_out_kernel(
+    a: wp.array(dtype=float),  # input only
+    b: wp.array(dtype=float),  # input and output
+    c: wp.array(dtype=float),  # output only
+):
+    tid = wp.tid()
+    b[tid] += a[tid]
+    c[tid] = 2.0 * a[tid]
+
+
+@wp.kernel
+def multi_out_kernel(
+    a: wp.array(dtype=float), b: wp.array(dtype=float), s: float, c: wp.array(dtype=float), d: wp.array(dtype=float)
+):
+    tid = wp.tid()
+    c[tid] = a[tid] + b[tid]
+    d[tid] = s * a[tid]
+
+
+@wp.kernel
+def multi_out_kernel_v2(
+    a: wp.array(dtype=float), b: wp.array(dtype=float), s: float, c: wp.array(dtype=float), d: wp.array(dtype=float)
+):
+    tid = wp.tid()
+    c[tid] = a[tid] * a[tid]
+    d[tid] = a[tid] * b[tid] * s
+
+
+@wp.kernel
+def multi_out_kernel_v3(
+    a: wp.array(dtype=float), b: wp.array(dtype=float), s: float, c: wp.array(dtype=float), d: wp.array(dtype=float)
+):
+    tid = wp.tid()
+    c[tid] = a[tid] ** 2.0
+    d[tid] = a[tid] * b[tid] * s
+
+
+@wp.kernel
+def scale_sum_square_kernel(a: wp.array(dtype=float), b: wp.array(dtype=float), s: float, c: wp.array(dtype=float)):
+    tid = wp.tid()
+    c[tid] = (a[tid] * s + b[tid]) ** 2.0
+
+
+# The Python function to call.
+# Note the argument annotations, just like Warp kernels.
+def scale_func(
+    # inputs
+    a: wp.array(dtype=float),
+    b: wp.array(dtype=wp.vec2),
+    s: float,
+    # outputs
+    c: wp.array(dtype=float),
+    d: wp.array(dtype=wp.vec2),
+):
+    wp.launch(scale_kernel, dim=a.shape, inputs=[a, s], outputs=[c])
+    wp.launch(scale_vec_kernel, dim=b.shape, inputs=[b, s], outputs=[d])
+
+
+def in_out_func(
+    a: wp.array(dtype=float),  # input only
+    b: wp.array(dtype=float),  # input and output
+    c: wp.array(dtype=float),  # output only
+):
+    wp.launch(scale_kernel, dim=a.size, inputs=[a, 2.0], outputs=[c])
+    wp.launch(accum_kernel, dim=a.size, inputs=[a, b])  # modifies `b`
+
+
+def double_func(
+    # inputs
+    a: wp.array(dtype=float),
+    # outputs
+    b: wp.array(dtype=float),
+):
+    wp.launch(scale_kernel, dim=a.shape, inputs=[a, 2.0], outputs=[b])
+
+
+@unittest.skipUnless(_jax_version() >= (0, 5, 0), "Jax version too old")
+def test_ffi_jax_kernel_add(test, device):
+    # two inputs and one output
+    import jax.numpy as jp
+
+    from warp.jax_experimental.ffi import jax_kernel
+
+    jax_add = jax_kernel(add_kernel)
+
+    @jax.jit
+    def f():
+        n = ARRAY_SIZE
+        a = jp.arange(n, dtype=jp.float32)
+        b = jp.ones(n, dtype=jp.float32)
+        return jax_add(a, b)
+
+    with jax.default_device(wp.device_to_jax(device)):
+        (y,) = f()
+
+    wp.synchronize_device(device)
+
+    result = np.asarray(y)
+    expected = np.arange(1, ARRAY_SIZE + 1, dtype=np.float32)
+
+    assert_np_equal(result, expected)
+
+
+@unittest.skipUnless(_jax_version() >= (0, 5, 0), "Jax version too old")
+def test_ffi_jax_kernel_sincos(test, device):
+    # one input and two outputs
+    import jax.numpy as jp
+
+    from warp.jax_experimental.ffi import jax_kernel
+
+    jax_sincos = jax_kernel(sincos_kernel, num_outputs=2)
+
+    n = ARRAY_SIZE
+
+    @jax.jit
+    def f():
+        a = jp.linspace(0, 2 * jp.pi, n, dtype=jp.float32)
+        return jax_sincos(a)
+
+    with jax.default_device(wp.device_to_jax(device)):
+        s, c = f()
+
+    wp.synchronize_device(device)
+
+    result_s = np.asarray(s)
+    result_c = np.asarray(c)
+
+    a = np.linspace(0, 2 * np.pi, n, dtype=np.float32)
+    expected_s = np.sin(a)
+    expected_c = np.cos(a)
+
+    assert_np_equal(result_s, expected_s, tol=1e-4)
+    assert_np_equal(result_c, expected_c, tol=1e-4)
+
+
+@unittest.skipUnless(_jax_version() >= (0, 5, 0), "Jax version too old")
+def test_ffi_jax_kernel_diagonal(test, device):
+    # no inputs and one output
+    from warp.jax_experimental.ffi import jax_kernel
+
+    jax_diagonal = jax_kernel(diagonal_kernel)
+
+    @jax.jit
+    def f():
+        # launch dimensions determine output size
+        return jax_diagonal(launch_dims=4)
+
+    wp.synchronize_device(device)
+
+    with jax.default_device(wp.device_to_jax(device)):
+        (d,) = f()
+
+    result = np.asarray(d)
+    expected = np.array(
+        [
+            [[1.0, 0.0, 0.0], [0.0, 2.0, 0.0], [0.0, 0.0, 3.0]],
+            [[2.0, 0.0, 0.0], [0.0, 4.0, 0.0], [0.0, 0.0, 6.0]],
+            [[3.0, 0.0, 0.0], [0.0, 6.0, 0.0], [0.0, 0.0, 9.0]],
+            [[4.0, 0.0, 0.0], [0.0, 8.0, 0.0], [0.0, 0.0, 12.0]],
+        ],
+        dtype=np.float32,
+    )
+
+    assert_np_equal(result, expected)
+
+
+@unittest.skipUnless(_jax_version() >= (0, 5, 0), "Jax version too old")
+def test_ffi_jax_kernel_in_out(test, device):
+    # in-out args
+    import jax.numpy as jp
+
+    from warp.jax_experimental.ffi import jax_kernel
+
+    jax_func = jax_kernel(in_out_kernel, num_outputs=2, in_out_argnames=["b"])
+
+    f = jax.jit(jax_func)
+
+    with jax.default_device(wp.device_to_jax(device)):
+        a = jp.ones(ARRAY_SIZE, dtype=jp.float32)
+        b = jp.arange(ARRAY_SIZE, dtype=jp.float32)
+        b, c = f(a, b)
+
+    wp.synchronize_device(device)
+
+    assert_np_equal(b, np.arange(1, ARRAY_SIZE + 1, dtype=np.float32))
+    assert_np_equal(c, np.full(ARRAY_SIZE, 2, dtype=np.float32))
+
+
+@unittest.skipUnless(_jax_version() >= (0, 5, 0), "Jax version too old")
+def test_ffi_jax_kernel_scale_vec_constant(test, device):
+    # multiply vectors by scalar (constant)
+    import jax.numpy as jp
+
+    from warp.jax_experimental.ffi import jax_kernel
+
+    jax_scale_vec = jax_kernel(scale_vec_kernel)
+
+    @jax.jit
+    def f():
+        a = jp.arange(ARRAY_SIZE, dtype=jp.float32).reshape((ARRAY_SIZE // 2, 2))  # array of vec2
+        s = 2.0
+        return jax_scale_vec(a, s)
+
+    with jax.default_device(wp.device_to_jax(device)):
+        (b,) = f()
+
+    wp.synchronize_device(device)
+
+    expected = 2 * np.arange(ARRAY_SIZE, dtype=np.float32).reshape((ARRAY_SIZE // 2, 2))
+
+    assert_np_equal(b, expected)
+
+
+@unittest.skipUnless(_jax_version() >= (0, 5, 0), "Jax version too old")
+def test_ffi_jax_kernel_scale_vec_static(test, device):
+    # multiply vectors by scalar (static arg)
+    import jax.numpy as jp
+
+    from warp.jax_experimental.ffi import jax_kernel
+
+    jax_scale_vec = jax_kernel(scale_vec_kernel)
+
+    # NOTE: scalar arguments must be static compile-time constants
+    @partial(jax.jit, static_argnames=["s"])
+    def f(a, s):
+        return jax_scale_vec(a, s)
+
+    a = jp.arange(ARRAY_SIZE, dtype=jp.float32).reshape((ARRAY_SIZE // 2, 2))  # array of vec2
+    s = 3.0
+
+    with jax.default_device(wp.device_to_jax(device)):
+        (b,) = f(a, s)
+
+    wp.synchronize_device(device)
+
+    expected = 3 * np.arange(ARRAY_SIZE, dtype=np.float32).reshape((ARRAY_SIZE // 2, 2))
+
+    assert_np_equal(b, expected)
+
+
+@unittest.skipUnless(_jax_version() >= (0, 5, 0), "Jax version too old")
+def test_ffi_jax_kernel_launch_dims_default(test, device):
+    # specify default launch dims
+    import jax.numpy as jp
+
+    from warp.jax_experimental.ffi import jax_kernel
+
+    N, M, K = 3, 4, 2
+
+    jax_matmul = jax_kernel(matmul_kernel, launch_dims=(N, M))
+
+    @jax.jit
+    def f():
+        a = jp.full((N, K), 2, dtype=jp.float32)
+        b = jp.full((K, M), 3, dtype=jp.float32)
+
+        # use default launch dims
+        return jax_matmul(a, b)
+
+    with jax.default_device(wp.device_to_jax(device)):
+        (result,) = f()
+
+    wp.synchronize_device(device)
+
+    expected = np.full((3, 4), 12, dtype=np.float32)
+
+    test.assertEqual(result.shape, expected.shape)
+    assert_np_equal(result, expected)
+
+
+@unittest.skipUnless(_jax_version() >= (0, 5, 0), "Jax version too old")
+def test_ffi_jax_kernel_launch_dims_custom(test, device):
+    # specify custom launch dims per call
+    import jax.numpy as jp
+
+    from warp.jax_experimental.ffi import jax_kernel
+
+    jax_matmul = jax_kernel(matmul_kernel)
+
+    @jax.jit
+    def f():
+        N1, M1, K1 = 3, 4, 2
+        a1 = jp.full((N1, K1), 2, dtype=jp.float32)
+        b1 = jp.full((K1, M1), 3, dtype=jp.float32)
+
+        # use custom launch dims
+        result1 = jax_matmul(a1, b1, launch_dims=(N1, M1))
+
+        N2, M2, K2 = 4, 3, 2
+        a2 = jp.full((N2, K2), 2, dtype=jp.float32)
+        b2 = jp.full((K2, M2), 3, dtype=jp.float32)
+
+        # use different custom launch dims
+        result2 = jax_matmul(a2, b2, launch_dims=(N2, M2))
+
+        return result1[0], result2[0]
+
+    with jax.default_device(wp.device_to_jax(device)):
+        result1, result2 = f()
+
+    wp.synchronize_device(device)
+
+    expected1 = np.full((3, 4), 12, dtype=np.float32)
+    expected2 = np.full((4, 3), 12, dtype=np.float32)
+
+    test.assertEqual(result1.shape, expected1.shape)
+    test.assertEqual(result2.shape, expected2.shape)
+    assert_np_equal(result1, expected1)
+    assert_np_equal(result2, expected2)
+
+
+@unittest.skipUnless(_jax_version() >= (0, 5, 0), "Jax version too old")
+def test_ffi_jax_callable_scale_constant(test, device):
+    # scale two arrays using a constant
+    import jax.numpy as jp
+
+    from warp.jax_experimental.ffi import jax_callable
+
+    jax_func = jax_callable(scale_func, num_outputs=2)
+
+    @jax.jit
+    def f():
+        # inputs
+        a = jp.arange(ARRAY_SIZE, dtype=jp.float32)
+        b = jp.arange(ARRAY_SIZE, dtype=jp.float32).reshape((ARRAY_SIZE // 2, 2))  # wp.vec2
+        s = 2.0
+
+        # output shapes
+        output_dims = {"c": a.shape, "d": b.shape}
+
+        c, d = jax_func(a, b, s, output_dims=output_dims)
+
+        return c, d
+
+    with jax.default_device(wp.device_to_jax(device)):
+        result1, result2 = f()
+
+    wp.synchronize_device(device)
+
+    expected1 = 2 * np.arange(ARRAY_SIZE, dtype=np.float32)
+    expected2 = 2 * np.arange(ARRAY_SIZE, dtype=np.float32).reshape((ARRAY_SIZE // 2, 2))
+
+    assert_np_equal(result1, expected1)
+    assert_np_equal(result2, expected2)
+
+
+@unittest.skipUnless(_jax_version() >= (0, 5, 0), "Jax version too old")
+def test_ffi_jax_callable_scale_static(test, device):
+    # scale two arrays using a static arg
+    import jax.numpy as jp
 
+    from warp.jax_experimental.ffi import jax_callable
+
+    jax_func = jax_callable(scale_func, num_outputs=2)
+
+    # NOTE: scalar arguments must be static compile-time constants
+    @partial(jax.jit, static_argnames=["s"])
+    def f(a, b, s):
+        # output shapes
+        output_dims = {"c": a.shape, "d": b.shape}
+
+        c, d = jax_func(a, b, s, output_dims=output_dims)
+
+        return c, d
+
+    with jax.default_device(wp.device_to_jax(device)):
+        # inputs
+        a = jp.arange(ARRAY_SIZE, dtype=jp.float32)
+        b = jp.arange(ARRAY_SIZE, dtype=jp.float32).reshape((ARRAY_SIZE // 2, 2))  # wp.vec2
+        s = 3.0
+        result1, result2 = f(a, b, s)
+
+    wp.synchronize_device(device)
+
+    expected1 = 3 * np.arange(ARRAY_SIZE, dtype=np.float32)
+    expected2 = 3 * np.arange(ARRAY_SIZE, dtype=np.float32).reshape((ARRAY_SIZE // 2, 2))
+
+    assert_np_equal(result1, expected1)
+    assert_np_equal(result2, expected2)
+
+
+@unittest.skipUnless(_jax_version() >= (0, 5, 0), "Jax version too old")
+def test_ffi_jax_callable_in_out(test, device):
+    # in-out arguments
+    import jax.numpy as jp
+
+    from warp.jax_experimental.ffi import jax_callable
+
+    jax_func = jax_callable(in_out_func, num_outputs=2, in_out_argnames=["b"])
+
+    f = jax.jit(jax_func)
+
+    with jax.default_device(wp.device_to_jax(device)):
+        a = jp.ones(ARRAY_SIZE, dtype=jp.float32)
+        b = jp.arange(ARRAY_SIZE, dtype=jp.float32)
+        b, c = f(a, b)
+
+    wp.synchronize_device(device)
+
+    assert_np_equal(b, np.arange(1, ARRAY_SIZE + 1, dtype=np.float32))
+    assert_np_equal(c, np.full(ARRAY_SIZE, 2, dtype=np.float32))
+
+
+@unittest.skipUnless(_jax_version() >= (0, 5, 0), "Jax version too old")
+def test_ffi_jax_callable_graph_cache(test, device):
+    # test graph caching limits
+    import jax
+    import jax.numpy as jp
+
+    from warp.jax_experimental.ffi import (
+        GraphMode,
+        clear_jax_callable_graph_cache,
+        get_jax_callable_default_graph_cache_max,
+        jax_callable,
+        set_jax_callable_default_graph_cache_max,
+    )
+
+    # --- test with default cache settings ---
+
+    jax_double = jax_callable(double_func, graph_mode=GraphMode.WARP)
+    f = jax.jit(jax_double)
+    arrays = []
+
+    test.assertEqual(jax_double.graph_cache_max, get_jax_callable_default_graph_cache_max())
+
+    with jax.default_device(wp.device_to_jax(device)):
+        for i in range(10):
+            n = 10 + i
+            a = jp.arange(n, dtype=jp.float32)
+            (b,) = f(a)
+
+            assert_np_equal(b, 2 * np.arange(n, dtype=np.float32))
+
+            # ensure graph cache is always growing
+            test.assertEqual(jax_double.graph_cache_size, i + 1)
+
+            # keep JAX array alive to prevent the memory from being reused, thus forcing a new graph capture each time
+            arrays.append(a)
+
+    # --- test clearing one callable's cache ---
+
+    clear_jax_callable_graph_cache(jax_double)
+
+    test.assertEqual(jax_double.graph_cache_size, 0)
+
+    # --- test with a custom cache limit ---
+
+    graph_cache_max = 5
+    jax_double = jax_callable(double_func, graph_mode=GraphMode.WARP, graph_cache_max=graph_cache_max)
+    f = jax.jit(jax_double)
+    arrays = []
+
+    test.assertEqual(jax_double.graph_cache_max, graph_cache_max)
+
+    with jax.default_device(wp.device_to_jax(device)):
+        for i in range(10):
+            n = 10 + i
+            a = jp.arange(n, dtype=jp.float32)
+            (b,) = f(a)
+
+            assert_np_equal(b, 2 * np.arange(n, dtype=np.float32))
+
+            # ensure graph cache size is capped
+            test.assertEqual(jax_double.graph_cache_size, min(i + 1, graph_cache_max))
+
+            # keep JAX array alive to prevent the memory from being reused, thus forcing a new graph capture
+            arrays.append(a)
+
+    # --- test clearing all callables' caches ---
+
+    clear_jax_callable_graph_cache()
+
+    with wp.jax_experimental.ffi._FFI_REGISTRY_LOCK:
+        for c in wp.jax_experimental.ffi._FFI_CALLABLE_REGISTRY.values():
+            test.assertEqual(c.graph_cache_size, 0)
+
+    # --- test with a custom default cache limit ---
+
+    saved_max = get_jax_callable_default_graph_cache_max()
+    try:
+        set_jax_callable_default_graph_cache_max(5)
+        jax_double = jax_callable(double_func, graph_mode=GraphMode.WARP)
+        f = jax.jit(jax_double)
+        arrays = []
+
+        test.assertEqual(jax_double.graph_cache_max, get_jax_callable_default_graph_cache_max())
+
+        with jax.default_device(wp.device_to_jax(device)):
+            for i in range(10):
+                n = 10 + i
+                a = jp.arange(n, dtype=jp.float32)
+                (b,) = f(a)
+
+                assert_np_equal(b, 2 * np.arange(n, dtype=np.float32))
+
+                # ensure graph cache size is capped
+                test.assertEqual(
+                    jax_double.graph_cache_size,
+                    min(i + 1, get_jax_callable_default_graph_cache_max()),
+                )
+
+                # keep JAX array alive to prevent the memory from being reused, thus forcing a new graph capture
+                arrays.append(a)
+
+        clear_jax_callable_graph_cache()
+
+    finally:
+        set_jax_callable_default_graph_cache_max(saved_max)
+
+
+@unittest.skipUnless(_jax_version() >= (0, 5, 0), "Jax version too old")
+def test_ffi_jax_callable_pmap_mul(test, device):
+    import jax
+    import jax.numpy as jp
+
+    from warp.jax_experimental.ffi import jax_callable
+
+    j = jax_callable(double_func, num_outputs=1)
+
+    ndev = jax.local_device_count()
+    per_device = max(ARRAY_SIZE // ndev, 64)
+    x = jp.arange(ndev * per_device, dtype=jp.float32).reshape((ndev, per_device))
+
+    def per_device_func(v):
+        (y,) = j(v)
+        return y
+
+    y = jax.pmap(per_device_func)(x)
+
+    wp.synchronize()
+
+    assert_np_equal(np.asarray(y), 2 * np.asarray(x))
+
+
+@unittest.skipUnless(_jax_version() >= (0, 5, 0), "Jax version too old")
+def test_ffi_jax_callable_pmap_multi_output(test, device):
+    import jax
+    import jax.numpy as jp
+
+    from warp.jax_experimental.ffi import jax_callable
+
+    def multi_out_py(
+        a: wp.array(dtype=float),
+        b: wp.array(dtype=float),
+        s: float,
+        c: wp.array(dtype=float),
+        d: wp.array(dtype=float),
+    ):
+        wp.launch(multi_out_kernel, dim=a.shape, inputs=[a, b, s], outputs=[c, d])
+
+    j = jax_callable(multi_out_py, num_outputs=2)
+
+    ndev = jax.local_device_count()
+    per_device = max(ARRAY_SIZE // ndev, 64)
+    a = jp.arange(ndev * per_device, dtype=jp.float32).reshape((ndev, per_device))
+    b = jp.ones((ndev, per_device), dtype=jp.float32)
+    s = 3.0
+
+    def per_device_func(aa, bb):
+        c, d = j(aa, bb, s)
+        return c + d  # simple combine to exercise both outputs
+
+    out = jax.pmap(per_device_func)(a, b)
+
+    wp.synchronize()
+
+    a_np = np.arange(ndev * per_device, dtype=np.float32).reshape((ndev, per_device))
+    b_np = np.ones((ndev, per_device), dtype=np.float32)
+    ref = (a_np + b_np) + s * a_np
+    assert_np_equal(np.asarray(out), ref)
+
+
+@unittest.skipUnless(_jax_version() >= (0, 5, 0), "Jax version too old")
+def test_ffi_jax_callable_pmap_multi_stage(test, device):
+    import jax
+    import jax.numpy as jp
+
+    from warp.jax_experimental.ffi import jax_callable
+
+    def multi_stage_py(
+        a: wp.array(dtype=float),
+        b: wp.array(dtype=float),
+        alpha: float,
+        tmp: wp.array(dtype=float),
+        out: wp.array(dtype=float),
+    ):
+        wp.launch(add_kernel, dim=a.shape, inputs=[a, b], outputs=[tmp])
+        wp.launch(axpy_kernel, dim=a.shape, inputs=[tmp, b, alpha], outputs=[out])
+
+    j = jax_callable(multi_stage_py, num_outputs=2)
+
+    ndev = jax.local_device_count()
+    per_device = max(ARRAY_SIZE // ndev, 64)
+    a = jp.arange(ndev * per_device, dtype=jp.float32).reshape((ndev, per_device))
+    b = jp.ones((ndev, per_device), dtype=jp.float32)
+    alpha = 2.5
+
+    def per_device_func(aa, bb):
+        tmp, out = j(aa, bb, alpha)
+        return tmp + out
+
+    combined = jax.pmap(per_device_func)(a, b)
+
+    wp.synchronize()
+
+    a_np = np.arange(ndev * per_device, dtype=np.float32).reshape((ndev, per_device))
+    b_np = np.ones((ndev, per_device), dtype=np.float32)
+    tmp_ref = a_np + b_np
+    out_ref = alpha * (a_np + b_np) + b_np
+    ref = tmp_ref + out_ref
+    assert_np_equal(np.asarray(combined), ref)
+
+
+@unittest.skipUnless(_jax_version() >= (0, 5, 0), "Jax version too old")
+def test_ffi_callback(test, device):
+    # in-out arguments
+    import jax.numpy as jp
+
+    from warp.jax_experimental.ffi import register_ffi_callback
+
+    # the Python function to call
+    def warp_func(inputs, outputs, attrs, ctx):
+        # input arrays
+        a = inputs[0]
+        b = inputs[1]
+
+        # scalar attributes
+        s = attrs["scale"]
+
+        # output arrays
+        c = outputs[0]
+        d = outputs[1]
+
+        device = wp.device_from_jax(get_jax_device())
+        stream = wp.Stream(device, cuda_stream=ctx.stream)
+
+        with wp.ScopedStream(stream):
+            # launch with arrays of scalars
+            wp.launch(scale_kernel, dim=a.shape, inputs=[a, s], outputs=[c])
+
+            # launch with arrays of vec2
+            # NOTE: the input shapes are from JAX arrays, we need to strip the inner dimension for vec2 arrays
+            wp.launch(scale_vec_kernel, dim=b.shape[0], inputs=[b, s], outputs=[d])
+
+    # register callback
+    register_ffi_callback("warp_func", warp_func)
+
+    n = ARRAY_SIZE
+
+    with jax.default_device(wp.device_to_jax(device)):
+        # inputs
+        a = jp.arange(n, dtype=jp.float32)
+        b = jp.arange(n, dtype=jp.float32).reshape((n // 2, 2))  # array of wp.vec2
+        s = 2.0
+
+        # set up call
+        out_types = [
+            jax.ShapeDtypeStruct(a.shape, jp.float32),
+            jax.ShapeDtypeStruct(b.shape, jp.float32),  # array of wp.vec2
+        ]
+        call = jax.ffi.ffi_call("warp_func", out_types)
+
+        # call it
+        c, d = call(a, b, scale=s)
+
+    wp.synchronize_device(device)
+
+    assert_np_equal(c, 2 * np.arange(ARRAY_SIZE, dtype=np.float32))
+    assert_np_equal(d, 2 * np.arange(ARRAY_SIZE, dtype=np.float32).reshape((ARRAY_SIZE // 2, 2)))
+
+
+@unittest.skipUnless(_jax_version() >= (0, 5, 0), "Jax version too old")
+def test_ffi_jax_kernel_autodiff_simple(test, device):
+    import jax
+    import jax.numpy as jp
+
+    from warp.jax_experimental.ffi import jax_kernel
+
+    jax_func = jax_kernel(
+        scale_sum_square_kernel,
+        num_outputs=1,
+        enable_backward=True,
+    )
+
+    from functools import partial
+
+    @partial(jax.jit, static_argnames=["s"])
+    def loss(a, b, s):
+        out = jax_func(a, b, s)[0]
+        return jp.sum(out)
+
+    n = ARRAY_SIZE
+    a = jp.arange(n, dtype=jp.float32)
+    b = jp.ones(n, dtype=jp.float32)
+    s = 2.0
+
+    with jax.default_device(wp.device_to_jax(device)):
+        da, db = jax.grad(loss, argnums=(0, 1))(a, b, s)
+
+    wp.synchronize_device(device)
+
+    # reference gradients
+    # d/da sum((a*s + b)^2) = sum(2*(a*s + b) * s)
+    # d/db sum((a*s + b)^2) = sum(2*(a*s + b))
+    a_np = np.arange(n, dtype=np.float32)
+    b_np = np.ones(n, dtype=np.float32)
+    ref_da = 2.0 * (a_np * s + b_np) * s
+    ref_db = 2.0 * (a_np * s + b_np)
+
+    assert_np_equal(np.asarray(da), ref_da)
+    assert_np_equal(np.asarray(db), ref_db)
+
+
+@unittest.skipUnless(_jax_version() >= (0, 5, 0), "Jax version too old")
+def test_ffi_jax_kernel_autodiff_jit_of_grad_simple(test, device):
+    import jax
+    import jax.numpy as jp
+
+    from warp.jax_experimental.ffi import jax_kernel
+
+    jax_func = jax_kernel(scale_sum_square_kernel, num_outputs=1, enable_backward=True)
+
+    def loss(a, b, s):
+        out = jax_func(a, b, s)[0]
+        return jp.sum(out)
+
+    grad_fn = jax.grad(loss, argnums=(0, 1))
+
+    # more typical: jit(grad(...)) with static scalar
+    jitted_grad = jax.jit(lambda a, b, s: grad_fn(a, b, s), static_argnames=("s",))
+
+    n = ARRAY_SIZE
+    a = jp.arange(n, dtype=jp.float32)
+    b = jp.ones(n, dtype=jp.float32)
+    s = 2.0
+
+    with jax.default_device(wp.device_to_jax(device)):
+        da, db = jitted_grad(a, b, s)
+
+    wp.synchronize_device(device)
+
+    a_np = np.arange(n, dtype=np.float32)
+    b_np = np.ones(n, dtype=np.float32)
+    ref_da = 2.0 * (a_np * s + b_np) * s
+    ref_db = 2.0 * (a_np * s + b_np)
+
+    assert_np_equal(np.asarray(da), ref_da)
+    assert_np_equal(np.asarray(db), ref_db)
+
+
+@unittest.skipUnless(_jax_version() >= (0, 5, 0), "Jax version too old")
+def test_ffi_jax_kernel_autodiff_multi_output(test, device):
+    import jax
+    import jax.numpy as jp
+
+    from warp.jax_experimental.ffi import jax_kernel
+
+    jax_func = jax_kernel(multi_out_kernel_v3, num_outputs=2, enable_backward=True)
+
+    def caller(fn, a, b, s):
+        c, d = fn(a, b, s)
+        return jp.sum(c + d)
+
+    @jax.jit
+    def grads(a, b, s):
+        # mark s as static in the inner call via partial to avoid hashing
+        def _inner(a, b, s):
+            return caller(jax_func, a, b, s)
+
+        return jax.grad(lambda a, b: _inner(a, b, 2.0), argnums=(0, 1))(a, b)
+
+    n = ARRAY_SIZE
+    a = jp.arange(n, dtype=jp.float32)
+    b = jp.ones(n, dtype=jp.float32)
+    s = 2.0
+
+    with jax.default_device(wp.device_to_jax(device)):
+        da, db = grads(a, b, s)
+
+    wp.synchronize_device(device)
+
+    a_np = np.arange(n, dtype=np.float32)
+    b_np = np.ones(n, dtype=np.float32)
+    # d/da sum(c+d) = 2*a + b*s
+    ref_da = 2.0 * a_np + b_np * s
+    # d/db sum(c+d) = a*s
+    ref_db = a_np * s
+
+    assert_np_equal(np.asarray(da), ref_da)
+    assert_np_equal(np.asarray(db), ref_db)
+
+
+@unittest.skipUnless(_jax_version() >= (0, 5, 0), "Jax version too old")
+def test_ffi_jax_kernel_autodiff_jit_of_grad_multi_output(test, device):
+    import jax
+    import jax.numpy as jp
+
+    from warp.jax_experimental.ffi import jax_kernel
+
+    jax_func = jax_kernel(multi_out_kernel_v3, num_outputs=2, enable_backward=True)
+
+    def loss(a, b, s):
+        c, d = jax_func(a, b, s)
+        return jp.sum(c + d)
+
+    grad_fn = jax.grad(loss, argnums=(0, 1))
+    jitted_grad = jax.jit(lambda a, b, s: grad_fn(a, b, s), static_argnames=("s",))
+
+    n = ARRAY_SIZE
+    a = jp.arange(n, dtype=jp.float32)
+    b = jp.ones(n, dtype=jp.float32)
+    s = 2.0
+
+    with jax.default_device(wp.device_to_jax(device)):
+        da, db = jitted_grad(a, b, s)
+
+    wp.synchronize_device(device)
+
+    a_np = np.arange(n, dtype=np.float32)
+    b_np = np.ones(n, dtype=np.float32)
+    ref_da = 2.0 * a_np + b_np * s
+    ref_db = a_np * s
+
+    assert_np_equal(np.asarray(da), ref_da)
+    assert_np_equal(np.asarray(db), ref_db)
+
+
+@unittest.skipUnless(_jax_version() >= (0, 5, 0), "Jax version too old")
+def test_ffi_jax_kernel_autodiff_2d(test, device):
+    import jax
+    import jax.numpy as jp
+
+    from warp.jax_experimental.ffi import jax_kernel
+
+    jax_func = jax_kernel(inc_2d_kernel, num_outputs=1, enable_backward=True)
+
+    @jax.jit
+    def loss(a):
+        out = jax_func(a)[0]
+        return jp.sum(out)
+
+    n, m = 8, 6
+    a = jp.arange(n * m, dtype=jp.float32).reshape((n, m))
+
+    with jax.default_device(wp.device_to_jax(device)):
+        (da,) = jax.grad(loss, argnums=(0,))(a)
+
+    wp.synchronize_device(device)
+
+    ref = np.ones((n, m), dtype=np.float32)
+    assert_np_equal(np.asarray(da), ref)
+
+
+@unittest.skipUnless(_jax_version() >= (0, 5, 0), "Jax version too old")
+def test_ffi_jax_kernel_autodiff_vec2(test, device):
+    import jax
+    import jax.numpy as jp
+
+    from warp.jax_experimental.ffi import jax_kernel
+
+    jax_func = jax_kernel(scale_vec_kernel, num_outputs=1, enable_backward=True)
+
+    from functools import partial
+
+    @partial(jax.jit, static_argnames=("s",))
+    def loss(a, s):
+        out = jax_func(a, s)[0]
+        return jp.sum(out)
+
+    n = ARRAY_SIZE
+    a = jp.arange(n, dtype=jp.float32).reshape((n // 2, 2))
+    s = 3.0
+
+    with jax.default_device(wp.device_to_jax(device)):
+        (da,) = jax.grad(loss, argnums=(0,))(a, s)
+
+    wp.synchronize_device(device)
+
+    # d/da sum(a*s) = s
+    ref = np.full_like(np.asarray(a), s)
+    assert_np_equal(np.asarray(da), ref)
+
+
+@unittest.skipUnless(_jax_version() >= (0, 5, 0), "Jax version too old")
+def test_ffi_jax_kernel_autodiff_mat22(test, device):
+    import jax
+    import jax.numpy as jp
+
+    from warp.jax_experimental.ffi import jax_kernel
+
+    @wp.kernel
+    def scale_mat_kernel(a: wp.array(dtype=wp.mat22), s: float, out: wp.array(dtype=wp.mat22)):
+        tid = wp.tid()
+        out[tid] = a[tid] * s
+
+    jax_func = jax_kernel(scale_mat_kernel, num_outputs=1, enable_backward=True)
+
+    from functools import partial
+
+    @partial(jax.jit, static_argnames=("s",))
+    def loss(a, s):
+        out = jax_func(a, s)[0]
+        return jp.sum(out)
+
+    n = 12  # must be divisible by 4 for 2x2 matrices
+    a = jp.arange(n, dtype=jp.float32).reshape((n // 4, 2, 2))
+    s = 2.5
+
+    with jax.default_device(wp.device_to_jax(device)):
+        (da,) = jax.grad(loss, argnums=(0,))(a, s)
+
+    wp.synchronize_device(device)
+
+    ref = np.full((n // 4, 2, 2), s, dtype=np.float32)
+    assert_np_equal(np.asarray(da), ref)
+
+
+@unittest.skipUnless(_jax_version() >= (0, 5, 0), "Jax version too old")
+def test_ffi_jax_kernel_autodiff_static_required(test, device):
+    import jax
+    import jax.numpy as jp
+
+    from warp.jax_experimental.ffi import jax_kernel
+
+    # Require explicit static_argnames for scalar s
+    jax_func = jax_kernel(scale_sum_square_kernel, num_outputs=1, enable_backward=True)
+
+    def loss(a, b, s):
+        out = jax_func(a, b, s)[0]
+        return jp.sum(out)
+
+    n = ARRAY_SIZE
+    a = jp.arange(n, dtype=jp.float32)
+    b = jp.ones(n, dtype=jp.float32)
+    s = 1.5
+
+    with jax.default_device(wp.device_to_jax(device)):
+        da, db = jax.grad(loss, argnums=(0, 1))(a, b, s)
+
+    wp.synchronize_device(device)
+
+    a_np = np.arange(n, dtype=np.float32)
+    b_np = np.ones(n, dtype=np.float32)
+    ref_da = 2.0 * (a_np * s + b_np) * s
+    ref_db = 2.0 * (a_np * s + b_np)
+
+    assert_np_equal(np.asarray(da), ref_da)
+    assert_np_equal(np.asarray(db), ref_db)
+
+
+@unittest.skipUnless(_jax_version() >= (0, 5, 0), "Jax version too old")
+def test_ffi_jax_kernel_autodiff_pmap_triple(test, device):
+    import jax
+    import jax.numpy as jp
+
+    from warp.jax_experimental.ffi import jax_kernel
+
+    jax_mul = jax_kernel(triple_kernel, num_outputs=1, enable_backward=True)
+
+    ndev = jax.local_device_count()
+    per_device = ARRAY_SIZE // ndev
+    x = jp.arange(ndev * per_device, dtype=jp.float32).reshape((ndev, per_device))
+
+    def per_device_loss(x):
+        y = jax_mul(x)[0]
+        return jp.sum(y)
+
+    grads = jax.pmap(jax.grad(per_device_loss))(x)
+
+    wp.synchronize()
+
+    assert_np_equal(np.asarray(grads), np.full((ndev, per_device), 3.0, dtype=np.float32))
+
+
+@unittest.skipUnless(_jax_version() >= (0, 5, 0), "Jax version too old")
+def test_ffi_jax_kernel_autodiff_pmap_multi_output(test, device):
+    import jax
+    import jax.numpy as jp
+
+    from warp.jax_experimental.ffi import jax_kernel
+
+    jax_mo = jax_kernel(multi_out_kernel_v2, num_outputs=2, enable_backward=True)
+
+    ndev = jax.local_device_count()
+    per_device = ARRAY_SIZE // ndev
+    a = jp.arange(ndev * per_device, dtype=jp.float32).reshape((ndev, per_device))
+    b = jp.arange(ndev * per_device, dtype=jp.float32).reshape((ndev, per_device))
+    s = 2.0
+
+    def per_dev_loss(aa, bb):
+        c, d = jax_mo(aa, bb, s)
+        return jp.sum(c + d)
+
+    da, db = jax.pmap(jax.grad(per_dev_loss, argnums=(0, 1)))(a, b)
+
+    wp.synchronize()
+
+    a_np = np.arange(ndev * per_device, dtype=np.float32).reshape((ndev, per_device))
+    b_np = np.arange(ndev * per_device, dtype=np.float32).reshape((ndev, per_device))
+    ref_da = 2.0 * a_np + b_np * s
+    ref_db = a_np * s
+    assert_np_equal(np.asarray(da), ref_da)
+    assert_np_equal(np.asarray(db), ref_db)
+
+
+class TestJax(unittest.TestCase):
+    pass
+
+
+# try adding Jax tests if Jax is installed correctly
+try:
+    # prevent Jax from gobbling up GPU memory
+    os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"
+    os.environ["XLA_PYTHON_CLIENT_ALLOCATOR"] = "platform"
+
+    import jax
+
+    # NOTE: we must enable 64-bit types in Jax to test the full gamut of types
+    jax.config.update("jax_enable_x64", True)
+
+    # check which Warp devices work with Jax
+    # CUDA devices may fail if Jax cannot find a CUDA Toolkit
+    test_devices = get_test_devices()
+    jax_compatible_devices = []
+    jax_compatible_cuda_devices = []
+    for d in test_devices:
+        try:
+            with jax.default_device(wp.device_to_jax(d)):
+                j = jax.numpy.arange(10, dtype=jax.numpy.float32)
+                j += 1
+            jax_compatible_devices.append(d)
+            if d.is_cuda:
+                jax_compatible_cuda_devices.append(d)
+        except Exception as e:
+            print(f"Skipping Jax DLPack tests on device '{d}' due to exception: {e}")
+
+    add_function_test(TestJax, "test_dtype_from_jax", test_dtype_from_jax, devices=None)
+    add_function_test(TestJax, "test_dtype_to_jax", test_dtype_to_jax, devices=None)
+
+    if jax_compatible_devices:
+        add_function_test(TestJax, "test_device_conversion", test_device_conversion, devices=jax_compatible_devices)
+
+    if jax_compatible_cuda_devices:
+        # tests for both custom_call and ffi variants of jax_kernel(), selected by installed JAX version
+        if jax.__version_info__ < (0, 4, 25):
+            # no interop supported
+            ffi_opts = []
+        elif jax.__version_info__ < (0, 5, 0):
+            # only custom_call supported
+            ffi_opts = [False]
+        elif jax.__version_info__ < (0, 8, 0):
+            # both custom_call and ffi supported
+            ffi_opts = [False, True]
+        else:
+            # only ffi supported
+            ffi_opts = [True]
+
+        for use_ffi in ffi_opts:
+            suffix = "ffi" if use_ffi else "cc"
+            add_function_test(
+                TestJax,
+                f"test_jax_kernel_basic_{suffix}",
+                test_jax_kernel_basic,
+                devices=jax_compatible_cuda_devices,
+                use_ffi=use_ffi,
+            )
+            add_function_test(
+                TestJax,
+                f"test_jax_kernel_scalar_{suffix}",
+                test_jax_kernel_scalar,
+                devices=jax_compatible_cuda_devices,
+                use_ffi=use_ffi,
+            )
+            add_function_test(
+                TestJax,
+                f"test_jax_kernel_vecmat_{suffix}",
+                test_jax_kernel_vecmat,
+                devices=jax_compatible_cuda_devices,
+                use_ffi=use_ffi,
+            )
+            add_function_test(
+                TestJax,
+                f"test_jax_kernel_multiarg_{suffix}",
+                test_jax_kernel_multiarg,
+                devices=jax_compatible_cuda_devices,
+                use_ffi=use_ffi,
+            )
+            add_function_test(
+                TestJax,
+                f"test_jax_kernel_launch_dims_{suffix}",
+                test_jax_kernel_launch_dims,
+                devices=jax_compatible_cuda_devices,
+                use_ffi=use_ffi,
+            )
+
+        # ffi.jax_kernel() tests
+        add_function_test(
+            TestJax, "test_ffi_jax_kernel_add", test_ffi_jax_kernel_add, devices=jax_compatible_cuda_devices
+        )
+        add_function_test(
+            TestJax, "test_ffi_jax_kernel_sincos", test_ffi_jax_kernel_sincos, devices=jax_compatible_cuda_devices
+        )
+        add_function_test(
+            TestJax, "test_ffi_jax_kernel_diagonal", test_ffi_jax_kernel_diagonal, devices=jax_compatible_cuda_devices
+        )
+        add_function_test(
+            TestJax, "test_ffi_jax_kernel_in_out", test_ffi_jax_kernel_in_out, devices=jax_compatible_cuda_devices
+        )
+        add_function_test(
+            TestJax,
+            "test_ffi_jax_kernel_scale_vec_constant",
+            test_ffi_jax_kernel_scale_vec_constant,
+            devices=jax_compatible_cuda_devices,
+        )
+        add_function_test(
+            TestJax,
+            "test_ffi_jax_kernel_scale_vec_static",
+            test_ffi_jax_kernel_scale_vec_static,
+            devices=jax_compatible_cuda_devices,
+        )
+        add_function_test(
+            TestJax,
+            "test_ffi_jax_kernel_launch_dims_default",
+            test_ffi_jax_kernel_launch_dims_default,
+            devices=jax_compatible_cuda_devices,
+        )
+        add_function_test(
+            TestJax,
+            "test_ffi_jax_kernel_launch_dims_custom",
+            test_ffi_jax_kernel_launch_dims_custom,
+            devices=jax_compatible_cuda_devices,
+        )
+
+        # ffi.jax_callable() tests
+        add_function_test(
+            TestJax,
+            "test_ffi_jax_callable_scale_constant",
+            test_ffi_jax_callable_scale_constant,
+            devices=jax_compatible_cuda_devices,
+        )
+        add_function_test(
+            TestJax,
+            "test_ffi_jax_callable_scale_static",
+            test_ffi_jax_callable_scale_static,
+            devices=jax_compatible_cuda_devices,
+        )
+        add_function_test(
+            TestJax, "test_ffi_jax_callable_in_out", test_ffi_jax_callable_in_out, devices=jax_compatible_cuda_devices
+        )
+        add_function_test(
+            TestJax,
+            "test_ffi_jax_callable_graph_cache",
+            test_ffi_jax_callable_graph_cache,
+            devices=jax_compatible_cuda_devices,
+        )
+
+        # pmap tests
+        add_function_test(
+            TestJax,
+            "test_ffi_jax_callable_pmap_multi_output",
+            test_ffi_jax_callable_pmap_multi_output,
+            devices=None,
+        )
+        add_function_test(
+            TestJax,
+            "test_ffi_jax_callable_pmap_mul",
+            test_ffi_jax_callable_pmap_mul,
+            devices=None,
+        )
+        add_function_test(
+            TestJax,
+            "test_ffi_jax_callable_pmap_multi_stage",
+            test_ffi_jax_callable_pmap_multi_stage,
+            devices=None,
+        )
+
+        # ffi callback tests
+        add_function_test(TestJax, "test_ffi_callback", test_ffi_callback, devices=jax_compatible_cuda_devices)
+
+        # autodiff tests
+        add_function_test(
+            TestJax,
+            "test_ffi_jax_kernel_autodiff_simple",
+            test_ffi_jax_kernel_autodiff_simple,
+            devices=jax_compatible_cuda_devices,
+        )
+        add_function_test(
+            TestJax,
+            "test_ffi_jax_kernel_autodiff_jit_of_grad_simple",
+            test_ffi_jax_kernel_autodiff_jit_of_grad_simple,
+            devices=jax_compatible_cuda_devices,
+        )
+        add_function_test(
+            TestJax,
+            "test_ffi_jax_kernel_autodiff_multi_output",
+            test_ffi_jax_kernel_autodiff_multi_output,
+            devices=jax_compatible_cuda_devices,
+        )
+        add_function_test(
+            TestJax,
+            "test_ffi_jax_kernel_autodiff_jit_of_grad_multi_output",
+            test_ffi_jax_kernel_autodiff_jit_of_grad_multi_output,
+            devices=jax_compatible_cuda_devices,
+        )
+        add_function_test(
+            TestJax,
+            "test_ffi_jax_kernel_autodiff_2d",
+            test_ffi_jax_kernel_autodiff_2d,
+            devices=jax_compatible_cuda_devices,
+        )
+        add_function_test(
+            TestJax,
+            "test_ffi_jax_kernel_autodiff_vec2",
+            test_ffi_jax_kernel_autodiff_vec2,
+            devices=jax_compatible_cuda_devices,
+        )
+        add_function_test(
+            TestJax,
+            "test_ffi_jax_kernel_autodiff_mat22",
+            test_ffi_jax_kernel_autodiff_mat22,
+            devices=jax_compatible_cuda_devices,
+        )
+        add_function_test(
+            TestJax,
+            "test_ffi_jax_kernel_autodiff_static_required",
+            test_ffi_jax_kernel_autodiff_static_required,
+            devices=jax_compatible_cuda_devices,
+        )
+
+        # autodiff with pmap tests
+        add_function_test(
+            TestJax,
+            "test_ffi_jax_kernel_autodiff_pmap_triple",
+            test_ffi_jax_kernel_autodiff_pmap_triple,
+            devices=None,
+        )
         add_function_test(
-            TestJax, "test_jax_kernel_launch_dims", test_jax_kernel_launch_dims, devices=jax_compatible_cuda_devices
+            TestJax,
+            "test_ffi_jax_kernel_autodiff_pmap_multi_output",
+            test_ffi_jax_kernel_autodiff_pmap_multi_output,
+            devices=None,
         )
 
 except Exception as e: