{"record":{"id":"3a92be8d1f59bd8e","repo":"jax-ml/jax","slug":"ordered-io-effects-not-supported-in-vmap","errorCode":null,"errorMessage":"Ordered IO effects not supported in vmap.","messagePattern":"Ordered IO effects not supported in vmap\\.","errorType":"error_code","errorClass":"Exception","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/control_flow/loops.py","lineNumber":1808,"sourceCode":"  #   a1, a2 = next((a1, a2) for a1, a2 in zip(body_avals, body_jaxpr.in_avals)\n  #                 if not core.typecompat(a1, a2))\n  #   raise core.JaxprTypeError(f\"while_loop body function input type error: {a1} != {a2}\")\n\n\n  joined_effects = _join_while_effects(body_jaxpr, cond_jaxpr, body_nconsts,\n                                       cond_nconsts)\n  disallowed_effects = effects.control_flow_allowed_effects.filter_not_in(joined_effects)\n  if disallowed_effects:\n    raise NotImplementedError(\n        f'Effects not supported in `while`: {disallowed_effects}')\n  return body_jaxpr.out_avals, joined_effects\n\n\ndef _while_loop_batching_rule(axis_data, args, dims, cond_nconsts, cond_jaxpr,\n                              body_nconsts, body_jaxpr):\n  from jax._src.callback import _IOEffect, _OrderedIOEffect\n  if any(_OrderedIOEffect in fn.effects for fn in [body_jaxpr, cond_jaxpr]):\n    raise Exception(\"Ordered IO effects not supported in vmap.\")\n\n  orig_batched = [d is not None for d in dims]\n  cconst_bat, bconst_bat, init_bat = split_list(orig_batched, [cond_nconsts, body_nconsts])\n  cconsts, bconsts, init = split_list(args, [cond_nconsts, body_nconsts])\n  cconst_dims, bconst_dims, init_dims = split_list(dims, [cond_nconsts, body_nconsts])\n\n  carry_bat = init_bat\n  # Fixpoint computation of which carry are batched: either\n  # batched from init, or the carry out is batched. Each iteration promotes\n  # at least one carry to batched. We need at most len(carry) iterations to\n  # reach a fixpoint.\n  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:","sourceCodeStart":1790,"sourceCodeEnd":1826,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/control_flow/loops.py#L1790-L1826","documentation":"When vmap batches a while_loop, the loop must be executed consistently across the batch dimension, but _OrderedIOEffect (produced by ordered host callbacks like jax.debug.print / debug.print with ordered=True, or older io callbacks) implies a sequential ordering that cannot be preserved under batching. The batching rule therefore raises a plain Exception refusing to vmap such loops.","triggerScenarios":"Applying jax.vmap over a function that internally calls lax.while_loop (or fori_loop) whose cond/body contains jax.debug.print or another ordered-IO callback; nested vmap + while_loop in simulation/RL step loops with printing.","commonSituations":"Debug prints left in batched rollout/training code; converting a per-sample while-loop routine to vmap without removing debug callbacks; host-callback-based logging inside iterative solvers.","solutions":["Remove or gate the debug print/callback inside the while_loop body before vmap (e.g. only print in unbatched debug runs)","Use jax.debug.print(..., ordered=True->False) where ordering is not required — but note plain callback effects are also restricted under vmap, so hoisting out is safest","Debug by running the unvmapped version, or with config.disable_jit, then vmap the clean version"],"exampleFix":"// before\ndef body(c, x):\n    jax.debug.print('c {}', c, ordered=True)\n    return c + x\njax.vmap(lambda xs: jax.lax.scan(body, 0, xs))(batch)\n\n// after\ndef body(c, x):\n    return c + x  # no ordered callback\nout = jax.vmap(lambda xs: jax.lax.scan(body, 0, xs))(batch)\njax.debug.print('out {}', out)  # print outside","handlingStrategy":"fallback","validationCode":"import jax\njaxpr = jax.make_jaxpr(lambda: batched_fn(sample_batch))()  # hard to check pre-hoc; simplest guard:\n# grep-like check: ensure no jax.debug.print/ordered callbacks in the vmapped function source\nimport inspect\nsrc = inspect.getsource(fn_to_vmap)\nassert 'debug.print' not in src and 'debug.callback' not in src, 'ordered IO blocks vmap'","typeGuard":null,"tryCatchPattern":"try:\n    jax.vmap(fn)(batch)\nexcept Exception as e:\n    if 'Ordered IO effects not supported in vmap' in str(e):\n        result = jax.lax.map(fn, batch)  # or strip debug callbacks and re-vmap\n    else:\n        raise","preventionTips":["Keep debug prints out of any function that will be vmapped; toggle them behind a DEBUG flag","Validate vmapped pipelines with debug logging disabled","Prefer jax.debug.callback only in non-batched, non-jitted debug paths"],"tags":["jax","vmap","ordered-io","debug-print","while-loop","batching"],"backgroundTag":"ordered-io-callback-under-vmap","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}