Improve usability of chexify.

PiperOrigin-RevId: 578145061

chex/_src/asserts_chexify.py

@@ -189,6 +189,16 @@ def _chexified_fn(*args, **kwargs):
           'Nested @chexify wrapping is disallowed. '
           'Make sure that you only wrap the function at the outermost level.')
 
+    if _ai.has_tracers((args, kwargs)):
+      raise RuntimeError(
+          '@chexify must be applied on top of all (p)jit/pmap transformations'
+          ' (otherwise it will result in `UnexpectedTracerError`). If you have'
+          ' functions that use value assertions, do not wrap them'
+          ' individually -- just wrap the outermost function after'
+          ' applying all your JAX transformations. See the example at '
+          'https://github.com/google-deepmind/chex#static-and-value-aka-runtime-assertions'
+      )
+
     if async_check:
       # Check completed calls.
       while async_check_futures and async_check_futures[0].done():

chex/_src/asserts_chexify_test.py

@@ -288,6 +288,36 @@ def fn(x, y):
       chexified_fn(jnp.array([2]))
       chexified_fn.wait_checks()  # Fail: not equal
 
+  def test_wrong_order_of_wrapping(self):
+
+    @chexify_async
+    def inner_fn(x, y):
+      asserts.assert_trees_all_equal(x, y)
+      return jax.tree_map(jnp.add, x, y)
+
+    def outer_fn(x, y):
+      z = inner_fn(x, y)
+      return jax.tree_map(jnp.square, z)
+
+    x = jnp.array([1])
+    y = jnp.array([1])
+    with self.assertRaisesRegex(
+        RuntimeError, '@chexify must be applied on top of all'
+    ):
+      jax.jit(inner_fn)(x, y)
+
+    with self.assertRaisesRegex(
+        RuntimeError, '@chexify must be applied on top of all'
+    ):
+      jax.jit(outer_fn)(x, y)
+
+    with self.assertRaisesRegex(
+        RuntimeError, '@chexify must be applied on top of all'
+    ):
+      jax.jit(chexify_async(outer_fn))(x, y)
+
+    outer_fn(x, y)
+
 
 class AssertsChexifyTestSuite(variants.TestCase):
   """Test suite for chexify assertions."""