Repository: google/jax
Issue: BUG: checkify of a vmap-ed while loop fails with an obscure error

### Description

Consider the following code:
```python
def test_checkify_with_batched_while():
    def func_to_check_with_loop(n, v):
        def while_condition(loop_state):
            counter, _ = loop_state
            return counter > 0

        def loop_body(loop_state):
            counter, value = loop_state
            checkify.check(value > 0, 'value must be positive')
            return counter - 1, value - 1

        _, result = jax.lax.while_loop(while_condition, loop_body, (n, v))
        return result

    batched_func = jax.vmap(func_to_check_with_loop)
    checked_f = checkify.checkify(batched_func)

    err, val = checked_f(jnp.asarray([1, 2, 3]), jnp.asarray([5, 2, 4]))
    err.throw()
    assert jnp.all(val == jnp.asarray([4, 0, 1]))
```

The code above fails when calling `checked_f` with the following error (truncated):
```
def _pred_bcast_select_mhlo(
        pred_aval: core.ShapedArray, pred: ir.Value, xs: Sequence[ir.Value],
        ys: Sequence[ir.Value], x_y_aval: core.AbstractValue) -> Sequence[ir.Value]:
      if x_y_aval is core.abstract_token:
        x, = xs
        y, = ys
        return [mhlo.AfterAllOp(mlir.aval_to_ir_type(x_y_aval), [x, y]).result]
      else:
        assert isinstance(x_y_aval, core.ShapedArray), x_y_aval
        x, = xs
        y, = ys
        assert x.type == y.type, (x.type, y.type)
>       assert (pred_aval.shape == x_y_aval.shape[:len(pred_aval.shape)]), (
                pred_aval.shape, x_y_aval)
E       jax._src.traceback_util.UnfilteredStackTrace: AssertionError: ((3,), ShapedArray(bool[]))
```

If one modifies the code by either removing `checkify`, or making a non-batched version by removing `vmap`, everything works as expected.

Can you please confirm that this is indeed a bug in JAX and, if so, provide a workaround?

### What jax/jaxlib version are you using?

0.3.16

### Which accelerator(s) are you using?

CPU

### Additional System Info

Mac

Changed files:
tests/checkify_test.py, jax/_src/checkify.py

diff --git a/jax/_src/checkify.py b/jax/_src/checkify.py
index 4243a9209..5a48a14c1 100644
--- a/jax/_src/checkify.py
+++ b/jax/_src/checkify.py
@@ -662,18 +662,54 @@ def ignore_error_output_jaxpr(jaxpr):
   new_jaxpr = jaxpr.replace(outvars=jaxpr.outvars[3:])
   return core.ClosedJaxpr(new_jaxpr, consts)
 
+def batch_error(err, code, payload, batch_shape):
+  err = jnp.broadcast_to(err, batch_shape)
+  code = jnp.broadcast_to(code, batch_shape)
+  payload = jnp.broadcast_to(payload, batch_shape+(3,))
+  return err, code, payload
+
+def unbatch_error(err, code, payload):
+  err = err.ravel()[0]
+  code = code.ravel()[0]
+  payload = payload.reshape(-1, 3)[0]
+  return err, code, payload
+
+def trivial_batched_jaxpr(jaxpr, batch_shape, batched_err):
+  fun = core.jaxpr_as_fun(jaxpr)
+
+  def g(err, code, payload, *a):
+    err_args = unbatch_error(err, code, payload)
+    err, code, payload, *out = fun(*err_args, *a)
+    err, code, payload = batch_error(err, code, payload, batch_shape)
+    return (err, code, payload, *out)
+
+  error_avals = map(lambda x: core.raise_to_shaped(core.get_aval(x)), batched_err)
+  new_jaxpr, _, literals_out = pe.trace_to_jaxpr_dynamic(
+      lu.wrap_init(g), [*error_avals, *jaxpr.in_avals[3:]])
+  return core.ClosedJaxpr(new_jaxpr, literals_out)
+
 def while_loop_error_check(error, enabled_errors, *in_flat, cond_nconsts,
                            cond_jaxpr, body_nconsts, body_jaxpr):
+  batch_shape = cond_jaxpr.out_avals[0].shape
+  if batch_shape:
+    err_args = batch_error(error.err, error.code, error.payload, batch_shape)
+  else:
+    err_args = [error.err, error.code, error.payload]
+
   c_consts, b_consts, carry = split_list(in_flat, [cond_nconsts, body_nconsts])
 
   # Check if the first cond application will error.
   checked_cond_jaxpr, msgs_cond = checkify_jaxpr(cond_jaxpr, error,
                                                  enabled_errors)
+  if batch_shape:
+    checked_cond_jaxpr = trivial_batched_jaxpr(checked_cond_jaxpr, batch_shape, err_args)
   cond_err, cond_code, cond_payload, _ = core.jaxpr_as_fun(checked_cond_jaxpr)(
-      error.err, error.code, error.payload, *c_consts, *carry)
+      *err_args, *c_consts, *carry)
 
   checked_body_jaxpr_, msgs_body = checkify_while_body_jaxpr(
     cond_jaxpr, body_jaxpr, error, enabled_errors, c_consts)
+  if batch_shape:
+    checked_body_jaxpr_ = trivial_batched_jaxpr(checked_body_jaxpr_, batch_shape, err_args)
   to_move = [False] * 3 + [True] * body_nconsts + [False] * len(carry)
   checked_body_jaxpr = pe.move_binders_to_front(checked_body_jaxpr_, to_move)
 
@@ -686,6 +722,9 @@ def while_loop_error_check(error, enabled_errors, *in_flat, cond_nconsts,
       *new_in_flat, cond_nconsts=cond_nconsts, cond_jaxpr=compat_cond_jaxpr,
       body_nconsts=body_nconsts, body_jaxpr=checked_body_jaxpr)
   new_msgs = {**error.msgs, **msgs_body, **msgs_cond}
+  if batch_shape:
+    err, code, payload = unbatch_error(err, code, payload)
+
   return out, Error(err, code, new_msgs, payload)
 error_checks[lax.while_p] = while_loop_error_check
 
diff --git a/tests/checkify_test.py b/tests/checkify_test.py
index 66ddd69de..68b839ae0 100644
--- a/tests/checkify_test.py
+++ b/tests/checkify_test.py
@@ -875,5 +875,35 @@ class AssertPrimitiveTests(jtu.JaxTestCase):
     self.assertIsNotNone(err.get())
     self.assertIn("should be positive", err.get())
 
+  def test_checkify_of_vmap_of_while(self):
+    @jax.vmap
+    def fun(n, v):
+      def while_cond(s):
+        counter, value = s
+        checkify.check(value < 6, "value needs to be less than 6!")
+        return counter > 0
+
+      def while_body(s):
+        counter, value = s
+        checkify.check(value >= 0, "value needs to be positive!")
+        return counter/value, value - 1.
+
+      _, result = jax.lax.while_loop(while_cond, while_body, (n, v))
+      return result
+
+    checked_f = checkify.checkify(fun, errors=checkify.all_checks)
+
+    err, _ = checked_f(jnp.asarray([1., 2., 3.]), jnp.asarray([5., 2., 4.]))
+    self.assertIsNotNone(err.get())
+    self.assertStartsWith(err.get(), "divided by zero")
+
+    err, _ = checked_f(jnp.asarray([1., 2., 3.]), jnp.asarray([5., 2., -4.]))
+    self.assertIsNotNone(err.get())
+    self.assertStartsWith(err.get(), "value needs to be positive")
+
+    err, _ = checked_f(jnp.asarray([1., 2., 3.]), jnp.asarray([6., 2., -4.]))
+    self.assertIsNotNone(err.get())
+    self.assertStartsWith(err.get(), "value needs to be less than 6")
+
 if __name__ == "__main__":
   absltest.main(testLoader=jtu.JaxTestLoader())