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


Discussion:
Given that many library users will likely use checkify on the top-most function (i.e. `train_step`), it would be nice to have some transformation that would allow the user to locally exclude functions from the checkify transform even if the calling code has been transformed. It should help to work around problems like the one in this issue until the new functionality is perfectly stable and will probably be useful in general.

For instance, in my case I would have liked to just label the function with the vmap-ed while loop as non-checkifiable, while continuing to check the rest of the code.

One other example where this transformation can be useful can be illustrated by the following code:
```python
def entropy(logits):
    chex.assert_rank(logits, 1)
    log_probs = jax.nn.log_softmax(logits)
    probs = jnp.exp(log_probs)
    return -jnp.sum(jnp.where(probs > 0, probs * log_probs, jnp.zeros_like(probs)))
```
`NaN` can arise in this code for zero probabilities, but it is a valid behavior for this particular function as it relies on `jnp.where` to filter them out. Therefore, I don't want `NaN` checks enabled for this particular function, but these checks will be very useful for the rest of the code, so I would like to be able to enable them globally.
Hello, can someone take a look at this please?
Sorry for the delayed response! I've been out of office for the past 2 weeks.

Thanks for the clear and minimal repro, I'll send a fix for this soon with your repro as a test-case.
Also, good point on adding some ability to locally exclude functions from being checkified. We've been thinking of a way to support "nullifying" any enabled checks from an outer checkify. I'll take this as a +1 to get that feature added! 


Resolving patch:
diff --git a/jax/_src/checkify.py b/jax/_src/checkify.py
--- 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
 


Test patch:
diff --git a/tests/checkify_test.py b/tests/checkify_test.py
--- a/tests/checkify_test.py
+++ b/tests/checkify_test.py
@@ -875,5 +875,35 @@ def f(x):
     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())