rewrite test not to include float0 broadcast

jax-ml · Nov 30, 2023 · 43ed74f · 43ed74f
1 parent 5b3fc1b
commit 43ed74f
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/api_test.py b/tests/api_test.py
@@ -5519,7 +5519,7 @@ def test_vjp_caching(self):
   def test_vjp_caching_static_argnums(self):
     identity = jax.remat(lambda x, y: jax.jit(lambda x: 2 * x if y else x)(x),
                          static_argnums=(1,))
-    _, f_vjp = jax.vjp(identity, 1., True)
+    _, f_vjp = jax.vjp(lambda x: identity(x, True), 1.)
     with jtu.count_jit_and_pmap_compiles() as count:  # noqa: F841
       for _ in range(20):
         f_vjp(1.)[0].block_until_ready()