{"record":{"id":"d626a34382c873c6","repo":"jax-ml/jax","slug":"checkify-does-not-support-batched-while-loops-che","errorCode":null,"errorMessage":"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.","messagePattern":"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\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/checkify.py","lineNumber":864,"sourceCode":"      debug_info=body_jaxpr.debug_info.with_unknown_names())\n  err_vals, err_tree = jtu.tree_flatten(error)\n  err_vals = map(core.typeof, err_vals)\n  flat_err_and_in_vals = [*err_vals, *c_consts_avals, *body_jaxpr.in_avals]\n  jaxpr, out_tree, error_effects = jaxpr_to_checkify_jaxpr(\n      jaxpr, enabled_errors, err_tree, *flat_err_and_in_vals)\n  return jaxpr, out_tree, error_effects\n\n\n@weakref_lru_cache\ndef ignore_error_output_jaxpr(jaxpr, num_error_vals: int):\n  \"\"\"Constructs a checked jaxpr which does not output its error value.\"\"\"\n  return jaxpr.replace(outvars=jaxpr.outvars[num_error_vals:])\n\ndef while_loop_error_check(error, enabled_errors, *in_flat, cond_nconsts,\n                           cond_jaxpr, body_nconsts, body_jaxpr):\n  if cond_jaxpr.out_avals[0].shape:\n    # TODO(lenamartens, sharadmv): support batched while.\n    raise ValueError('Checkify does not support batched while-loops '\n                     '(checkify-of-vmap-of-while). \\nHint: if possible, move '\n                     'the vmap to the outer level to get '\n                     'vmap-of-checkify-of-while.')\n\n  c_consts, b_consts, carry = split_list(in_flat, [cond_nconsts, body_nconsts])\n  # Check if the first cond application will error.\n  error, _ = checkify_jaxpr(cond_jaxpr, enabled_errors, error, *c_consts, *carry)\n\n  _, _, error_effects = checkify_while_body_jaxpr(cond_jaxpr, body_jaxpr,\n                                                  enabled_errors, error,\n                                                  cond_nconsts)\n  # merged error!\n  error = error._add_placeholder_effects(error_effects)\n  err_vals, err_tree = jtu.tree_flatten(error)\n  checked_body_jaxpr_, body_out_tree, _ = checkify_while_body_jaxpr(\n      cond_jaxpr, body_jaxpr, enabled_errors, error, cond_nconsts)\n  num_error_vals = len(err_vals)\n  to_move = ([False] * num_error_vals + [True] * cond_nconsts","sourceCodeStart":846,"sourceCodeEnd":882,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/checkify.py#L846-L882","documentation":"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.","triggerScenarios":"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.","commonSituations":"Ensemble/batched optimization loops written with while_loop then vmap'd, then wrapped in checkify for error monitoring.","solutions":["Restructure to vmap-of-checkify-of-while: apply checkify inside, vmap outside","Replace lax.while_loop with lax.fori_loop or Python-level loop when batching is needed","Batch manually by stacking state and using a scalar any() cond"],"exampleFix":"# before\nchecked = checkify.checkify(jax.vmap(train_step_while))\n# after\nchecked = jax.vmap(checkify.checkify(train_step_while))","handlingStrategy":"fallback","validationCode":"def cond_scalar(cond_fun, *args):\n    # ensure cond returns a scalar before using with vmap+checkify\n    out = cond_fun(*args)\n    assert np.ndim(out) == 0, 'batched cond: restructure to vmap-of-checkify-of-while'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Apply checkify innermost, vmap outermost by convention","Avoid lax.while_loop under vmap when planning to checkify"],"tags":["jax","checkify","vmap","while-loop","transformation-order"],"backgroundTag":"unsupported-transform-composition","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}