{"record":{"id":"966c3217e82b464d","repo":"jax-ml/jax","slug":"unordered-io-effects-not-supported-in-while-loop-w","errorCode":null,"errorMessage":"Unordered IO effects not supported in while_loop with batched predicate","messagePattern":"Unordered IO effects not supported in while_loop with batched predicate","errorType":"error_code","errorClass":"Exception","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/control_flow/loops.py","lineNumber":1838,"sourceCode":"  for _ in range(1 + len(carry_bat)):\n    _, carry_bat_out = batching.batch_jaxpr(\n        body_jaxpr, axis_data, bconst_bat + carry_bat, instantiate=carry_bat)\n    if carry_bat == carry_bat_out:\n      break\n    carry_bat = safe_map(operator.or_, carry_bat, carry_bat_out)\n  else:\n    assert False, \"Fixpoint not reached\"\n\n  # Knowing how the carry is batched now, we can determine if the predicate is\n  # batched.\n  _, (pred_bat,) = batching.batch_jaxpr(\n      cond_jaxpr, axis_data, cconst_bat + carry_bat, instantiate=False)\n\n  if pred_bat:\n    # If the predicate is batched, we have to batch *all* of the carry\n    # regardless of if the body needs it.\n    if any(_IOEffect in fn.effects for fn in [body_jaxpr, cond_jaxpr]):\n      raise Exception(\"Unordered IO effects not supported in while_loop \"\n                      \"with batched predicate\")\n    carry_bat = [True] * len(carry_bat)\n    carry_dims = [0] * len(carry_bat)\n    body_jaxpr_batched, _ = batching.batch_jaxpr_axes(\n        body_jaxpr, axis_data, bconst_dims + carry_dims, carry_dims)\n    cond_jaxpr_batched, _ = batching.batch_jaxpr_axes(\n        cond_jaxpr, axis_data, cconst_dims + carry_dims, [0])\n  else:\n    # If the predicate is not batched, we can look at the `cond_jaxpr`'s out\n    # shape to determine the rank of the predicate. From this rank we pick the\n    # dims of the carry to be batched to ensure that the predicate shape is a\n    # prefix of the carry in and out shapes. We can then batch the `body_jaxpr`\n    # according to these new batch dims.\n    cond_rank = len(cond_jaxpr.out_avals[0].shape)\n    carry_dims = [cond_rank if b else None for b in carry_bat]\n    body_jaxpr_batched, _ = batching.batch_jaxpr_axes(\n        body_jaxpr, axis_data, bconst_dims + carry_dims, carry_dims)\n    # Now we need to rebatch the `cond_jaxpr` according to the new dims of the","sourceCodeStart":1820,"sourceCodeEnd":1856,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/control_flow/loops.py#L1820-L1856","documentation":"JAX's while_loop batching rule (e.g. under vmap) must execute the loop body for the whole batch whenever any predicate is batched, which requires batching all carry values. Unordered IO effects (like random state or prints in jax.debug) cannot be batched this way, so JAX refuses to vmap a while_loop whose cond or body performs IO effects when the loop condition itself is batched.","triggerScenarios":"Calling jax.vmap over a function containing lax.while_loop (or fori_loop/scan lowering to while) where the loop predicate depends on the batched axis AND the cond/body jaxprs contain _IOEffect (e.g. jax.debug.print, host callbacks, or unordered effectful primitives).","commonSituations":"Using jax.debug.print or host_callback inside a while_loop for logging, then wrapping the whole function in vmap or scan-based batching; also seen with custom primitives carrying IO effects under vmap.","solutions":["Remove IO effects (debug prints, callbacks) from the while_loop cond and body, or gate them with jax.debug.print(..., ordered=False)-free constructs outside the loop","Restructure so the predicate is not batched (move the batched condition out, or mask per-sample loops with a fixed trip count using lax.scan)","Collect values in the loop carry and print/inspect them after the vmap call outside of the loop","If a fixed number of iterations is possible, replace while_loop with lax.scan, which supports batching with effects"],"exampleFix":"// before\njax.vmap(lambda x: lax.while_loop(lambda c: c[0] < 10,\n                                   lambda c: (jax.debug.print('{}', c[0]), (c[0]+1, c[1]*x))[1],\n                                   (0, 1.0)))\n// after (no IO effect inside loop)\njit_fn = jax.vmap(lambda x: lax.fori_loop(0, 10, lambda i, acc: acc * x, 1.0))","handlingStrategy":"validation","validationCode":"from jax._src import effects\nimport jax\n# before vmap, check for IO effects by tracing\ndef has_io_effects(fun, *args):\n    jaxpr = jax.make_jaxpr(fun)(*args)\n    return any(getattr(e, 'name', '') == 'IO' or 'IO' in type(e).__name__ for e in jaxpr.effects)","typeGuard":"def loop_is_vmappable_with_batched_pred(cond_fn, body_fn, example_carry):\n    jaxpr_c = jax.make_jaxpr(cond_fn)(example_carry)\n    jaxpr_b = jax.make_jaxpr(body_fn)(example_carry)\n    return not any('IO' in type(e).__name__ for e in (*jaxpr_c.effects, *jaxpr_b.effects))","tryCatchPattern":null,"preventionTips":["Never place jax.debug.print or host callbacks inside while_loop cond/body that will be vmap'd","Return debug values through the carry and log after the vmap","Prefer fixed-trip-count loops (fori_loop/scan) when mixing with vmap and effects"],"tags":["jax","vmap","while-loop","io-effects","control-flow"],"backgroundTag":"jax-vmap-effect-not-supported","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}