warp-lang 1.9.1__py3-none-manylinux_2_34_aarch64.whl → 1.10.0rc2__py3-none-manylinux_2_34_aarch64.whl

warp/tests/interop/test_jax.py

@@ -21,9 +21,12 @@ from typing import Any
 import numpy as np
 
 import warp as wp
-from warp.jax import get_jax_device
+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
@@ -46,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(
@@ -63,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]:
@@ -146,7 +161,7 @@ def test_jax_kernel_basic(test, device, use_ffi=False):
 
         jax_triple = jax_kernel(triple_kernel, quiet=True)  # suppress deprecation warnings
 
-    n = 64
+    n = ARRAY_SIZE
 
     @jax.jit
     def f():
@@ -157,6 +172,8 @@ def test_jax_kernel_basic(test, device, use_ffi=False):
     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)
 
@@ -175,6 +192,7 @@ def test_jax_kernel_scalar(test, device, use_ffi=False):
 
         kwargs = {"quiet": True}
 
+    # use a smallish size to ensure arange * 3 doesn't overflow
     n = 64
 
     for T in scalar_types:
@@ -196,6 +214,8 @@ def test_jax_kernel_scalar(test, device, use_ffi=False):
             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)
 
@@ -218,6 +238,7 @@ def test_jax_kernel_vecmat(test, device, use_ffi=False):
         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_
@@ -237,6 +258,8 @@ def test_jax_kernel_vecmat(test, device, use_ffi=False):
             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)
 
@@ -255,7 +278,7 @@ def test_jax_kernel_multiarg(test, device, use_ffi=False):
 
         jax_multiarg = jax_kernel(multiarg_kernel, quiet=True)
 
-    n = 64
+    n = ARRAY_SIZE
 
     @jax.jit
     def f():
@@ -268,6 +291,8 @@ def test_jax_kernel_multiarg(test, device, use_ffi=False):
     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)
@@ -292,40 +317,32 @@ def test_jax_kernel_launch_dims(test, device, use_ffi=False):
     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,), **kwargs
+    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), **kwargs
+    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
 
@@ -342,11 +359,17 @@ def test_jax_kernel_launch_dims(test, device, use_ffi=False):
 
 
 @wp.kernel
-def add_kernel(a: wp.array(dtype=int), b: wp.array(dtype=int), output: wp.array(dtype=int)):
+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]
 
 
+@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]
+
+
 @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()
@@ -408,6 +431,39 @@ def in_out_kernel(
     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(
@@ -432,6 +488,15 @@ def in_out_func(
     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
@@ -443,16 +508,18 @@ def test_ffi_jax_kernel_add(test, device):
 
     @jax.jit
     def f():
-        n = 10
-        a = jp.arange(n, dtype=jp.int32)
-        b = jp.ones(n, dtype=jp.int32)
+        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, 11, dtype=np.int32)
+    expected = np.arange(1, ARRAY_SIZE + 1, dtype=np.float32)
 
     assert_np_equal(result, expected)
 
@@ -465,7 +532,8 @@ def test_ffi_jax_kernel_sincos(test, device):
     from warp.jax_experimental.ffi import jax_kernel
 
     jax_sincos = jax_kernel(sincos_kernel, num_outputs=2)
-    n = 32
+
+    n = ARRAY_SIZE
 
     @jax.jit
     def f():
@@ -475,6 +543,8 @@ def test_ffi_jax_kernel_sincos(test, device):
     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)
 
@@ -498,6 +568,8 @@ def test_ffi_jax_kernel_diagonal(test, device):
         # 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()
 
@@ -527,12 +599,14 @@ def test_ffi_jax_kernel_in_out(test, device):
     f = jax.jit(jax_func)
 
     with jax.default_device(wp.device_to_jax(device)):
-        a = jp.ones(10, dtype=jp.float32)
-        b = jp.arange(10, dtype=jp.float32)
+        a = jp.ones(ARRAY_SIZE, dtype=jp.float32)
+        b = jp.arange(ARRAY_SIZE, dtype=jp.float32)
         b, c = f(a, b)
 
-    assert_np_equal(b, np.arange(1, 11, dtype=np.float32))
-    assert_np_equal(c, np.full(10, 2, dtype=np.float32))
+    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")
@@ -546,14 +620,16 @@ def test_ffi_jax_kernel_scale_vec_constant(test, device):
 
     @jax.jit
     def f():
-        a = jp.arange(10, dtype=jp.float32).reshape((5, 2))  # array of vec2
+        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()
 
-    expected = 2 * np.arange(10, dtype=np.float32).reshape((5, 2))
+    wp.synchronize_device(device)
+
+    expected = 2 * np.arange(ARRAY_SIZE, dtype=np.float32).reshape((ARRAY_SIZE // 2, 2))
 
     assert_np_equal(b, expected)
 
@@ -572,13 +648,15 @@ def test_ffi_jax_kernel_scale_vec_static(test, device):
     def f(a, s):
         return jax_scale_vec(a, s)
 
-    a = jp.arange(10, dtype=jp.float32).reshape((5, 2))  # array of vec2
+    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)
 
-    expected = 3 * np.arange(10, dtype=np.float32).reshape((5, 2))
+    wp.synchronize_device(device)
+
+    expected = 3 * np.arange(ARRAY_SIZE, dtype=np.float32).reshape((ARRAY_SIZE // 2, 2))
 
     assert_np_equal(b, expected)
 
@@ -605,6 +683,8 @@ def test_ffi_jax_kernel_launch_dims_default(test, device):
     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)
@@ -641,6 +721,8 @@ def test_ffi_jax_kernel_launch_dims_custom(test, device):
     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)
 
@@ -662,8 +744,8 @@ def test_ffi_jax_callable_scale_constant(test, device):
     @jax.jit
     def f():
         # inputs
-        a = jp.arange(10, dtype=jp.float32)
-        b = jp.arange(10, dtype=jp.float32).reshape((5, 2))  # wp.vec2
+        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
@@ -676,8 +758,10 @@ def test_ffi_jax_callable_scale_constant(test, device):
     with jax.default_device(wp.device_to_jax(device)):
         result1, result2 = f()
 
-    expected1 = 2 * np.arange(10, dtype=np.float32)
-    expected2 = 2 * np.arange(10, dtype=np.float32).reshape((5, 2))
+    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)
@@ -704,13 +788,15 @@ def test_ffi_jax_callable_scale_static(test, device):
 
     with jax.default_device(wp.device_to_jax(device)):
         # inputs
-        a = jp.arange(10, dtype=jp.float32)
-        b = jp.arange(10, dtype=jp.float32).reshape((5, 2))  # wp.vec2
+        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)
 
-    expected1 = 3 * np.arange(10, dtype=np.float32)
-    expected2 = 3 * np.arange(10, dtype=np.float32).reshape((5, 2))
+    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)
@@ -728,12 +814,224 @@ def test_ffi_jax_callable_in_out(test, device):
     f = jax.jit(jax_func)
 
     with jax.default_device(wp.device_to_jax(device)):
-        a = jp.ones(10, dtype=jp.float32)
-        b = jp.arange(10, dtype=jp.float32)
+        a = jp.ones(ARRAY_SIZE, dtype=jp.float32)
+        b = jp.arange(ARRAY_SIZE, dtype=jp.float32)
         b, c = f(a, b)
 
-    assert_np_equal(b, np.arange(1, 11, dtype=np.float32))
-    assert_np_equal(c, np.full(10, 2, dtype=np.float32))
+    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")
@@ -770,7 +1068,7 @@ def test_ffi_callback(test, device):
     # register callback
     register_ffi_callback("warp_func", warp_func)
 
-    n = 10
+    n = ARRAY_SIZE
 
     with jax.default_device(wp.device_to_jax(device)):
         # inputs
@@ -788,8 +1086,344 @@ def test_ffi_callback(test, device):
         # call it
         c, d = call(a, b, scale=s)
 
-    assert_np_equal(c, 2 * np.arange(10, dtype=np.float32))
-    assert_np_equal(d, 2 * np.arange(10, dtype=np.float32).reshape((5, 2)))
+    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):
@@ -936,10 +1570,99 @@ try:
         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_ffi_jax_kernel_autodiff_pmap_multi_output",
+            test_ffi_jax_kernel_autodiff_pmap_multi_output,
+            devices=None,
+        )
 
 except Exception as e:
     print(f"Skipping Jax tests due to exception: {e}")