jax-ml/jax · error · ValueError

Checkify does not support batched while-loops (checkify-of-v

Error message

Checkify does not support batched while-loops (checkify-of-vmap-of-while). 
Hint: if possible, move the vmap to the outer level to get vmap-of-checkify-of-while.

What it means

Checkify cannot transform a while_loop whose cond_jaxpr produces a batched (non-scalar) boolean — i.e. checkify(vmap(while_loop(...))). It raises with a hint to reorder the transformations.

Source

Thrown at jax/_src/checkify.py:864

      debug_info=body_jaxpr.debug_info.with_unknown_names())
  err_vals, err_tree = jtu.tree_flatten(error)
  err_vals = map(core.typeof, err_vals)
  flat_err_and_in_vals = [*err_vals, *c_consts_avals, *body_jaxpr.in_avals]
  jaxpr, out_tree, error_effects = jaxpr_to_checkify_jaxpr(
      jaxpr, enabled_errors, err_tree, *flat_err_and_in_vals)
  return jaxpr, out_tree, error_effects


@weakref_lru_cache
def ignore_error_output_jaxpr(jaxpr, num_error_vals: int):
  """Constructs a checked jaxpr which does not output its error value."""
  return jaxpr.replace(outvars=jaxpr.outvars[num_error_vals:])

def while_loop_error_check(error, enabled_errors, *in_flat, cond_nconsts,
                           cond_jaxpr, body_nconsts, body_jaxpr):
  if cond_jaxpr.out_avals[0].shape:
    # TODO(lenamartens, sharadmv): support batched while.
    raise ValueError('Checkify does not support batched while-loops '
                     '(checkify-of-vmap-of-while). \nHint: if possible, move '
                     'the vmap to the outer level to get '
                     'vmap-of-checkify-of-while.')

  c_consts, b_consts, carry = split_list(in_flat, [cond_nconsts, body_nconsts])
  # Check if the first cond application will error.
  error, _ = checkify_jaxpr(cond_jaxpr, enabled_errors, error, *c_consts, *carry)

  _, _, error_effects = checkify_while_body_jaxpr(cond_jaxpr, body_jaxpr,
                                                  enabled_errors, error,
                                                  cond_nconsts)
  # merged error!
  error = error._add_placeholder_effects(error_effects)
  err_vals, err_tree = jtu.tree_flatten(error)
  checked_body_jaxpr_, body_out_tree, _ = checkify_while_body_jaxpr(
      cond_jaxpr, body_jaxpr, enabled_errors, error, cond_nconsts)
  num_error_vals = len(err_vals)
  to_move = ([False] * num_error_vals + [True] * cond_nconsts

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Restructure to vmap-of-checkify-of-while: apply checkify inside, vmap outside
2. Replace lax.while_loop with lax.fori_loop or Python-level loop when batching is needed
3. Batch manually by stacking state and using a scalar any() cond

Example fix

# before
checked = checkify.checkify(jax.vmap(train_step_while))
# after
checked = jax.vmap(checkify.checkify(train_step_while))

def cond_scalar(cond_fun, *args):
    # ensure cond returns a scalar before using with vmap+checkify
    out = cond_fun(*args)
    assert np.ndim(out) == 0, 'batched cond: restructure to vmap-of-checkify-of-while'

When it happens

Trigger: Applying checkify.checkify on a function where vmap-of-lax.while_loop (or scan containing batched while) appears, making cond_jaxpr.out_avals[0] have a non-empty shape.

Common situations: Ensemble/batched optimization loops written with while_loop then vmap'd, then wrapped in checkify for error monitoring.

Related errors

• Ordered IO effects not supported in vmap.
• Unordered IO effects not supported in while_loop with batche
• Mapped away dimension of inputs passed to vmap should be sha
• Unmapped values passed to vmap cannot be sharded along the m
• {name} wrapped function must be passed at least one argument

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/d626a34382c873c6. Report an issue: GitHub.