{"record":{"id":"e6ae26feaf51b1c0","repo":"jax-ml/jax","slug":"io-effect-not-supported-in-vmap-of-cond","errorCode":null,"errorMessage":"IO effect not supported in vmap-of-cond.","messagePattern":"IO effect not supported in vmap-of-cond\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/control_flow/conditionals.py","lineNumber":472,"sourceCode":"\ndef _bcast_select_n(pred, *cases):\n  if np.ndim(pred) != np.ndim(cases[0]):\n    idx = list(range(np.ndim(pred)))\n    pred = lax.broadcast_in_dim(pred, np.shape(cases[0]), idx)\n  return lax.select_n(pred, *cases)\n\ndef _cond_batching_rule(axis_data, args, dims, *, branches, **params):\n  index, *ops = args\n  index_dim, *op_dims = dims\n  # TODO(sharadmv): clean this up by adding a specific blocklist\n  if any(isinstance(eff, RefEffect) for branch in branches for eff in\n      branch.effects):\n    raise NotImplementedError(\n        \"State effect not supported in vmap-of-cond.\")\n  from jax._src.callback import _IOEffect, _OrderedIOEffect\n  if any(eff in branch.effects for eff in [_IOEffect, _OrderedIOEffect]\n      for branch in branches):\n    raise NotImplementedError(\n        \"IO effect not supported in vmap-of-cond.\")\n\n  if \"branches_platforms\" in params and (index_dim is not None):\n    # If we end up with a mapped index for a platform_dependent cond, we can\n    # replace the index with a fresh call to platform_index. See #29329.\n    index = platform_index_p.bind(platforms=params[\"branches_platforms\"])\n    index_dim = None\n\n  if index_dim is not None:\n    # Convert to a lax.select. While we could get away with not broadcasting\n    # some operands yet, because all outputs must be broadcast together anyway\n    # for the select we broadcast the input operands for simplicity and leave\n    # optimizations to XLA.\n    # TODO(mattjj,frostig): assumes branches are side-effect-free, revise!\n    index, *ops = (\n        batching.bdim_at_front(x, d, axis_data.size,\n                               mesh_axis=axis_data.explicit_mesh_axis)\n        for x, d in zip(args, dims)","sourceCodeStart":454,"sourceCodeEnd":490,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/control_flow/conditionals.py#L454-L490","documentation":"Raised when a cond whose branches perform host callbacks / IO effects (host_callback.io_callback, debug callbacks, or ordered IO) is batched by vmap. IO effects cannot be soundly replicated across a batch axis inside cond, so the batching rule rejects them.","triggerScenarios":"jax.vmap over a lax.cond whose true/false branches call host_callback.io_callback, debug_callback, or any primitive carrying _IOEffect/_OrderedIOEffect.","commonSituations":"Mixing logging, printing, or host-side side effects with vmap inside conditionals; using io_callback for I/O in batched inference code.","solutions":["Move the io_callback out of the cond and guard it with jnp.where on a concrete flag, or call it unconditionally and mask outputs","Replace io_callback with pure JAX ops or compile-time (Python) branching when the predicate is concrete","Use lax.map instead of vmap if per-tracing IO is acceptable","Remove the callback from the batched path entirely"],"exampleFix":"// before\n@jax.vmap\ndef f(x):\n  return lax.cond(x > 0, x, lambda x: hcb.io_callback(log, x), x, lambda x: x)\n// after\n@jax.vmap\ndef f(x):\n  y = jnp.where(x > 0, x, x)  # pure compute\n  # do IO outside vmap on the result\n  return y","handlingStrategy":"validation","validationCode":"import jax\nj = jax.make_jaxpr(branch)(*args)\nfrom jax._src.effects import control_flow_allowed_effects\nassert not control_flow_allowed_effects.filter_not_in(j.effects), 'IO effects present'","typeGuard":"def is_pure_enough(fn, *args) -> bool:\n    j = jax.make_jaxpr(fn)(*args)\n    return not control_flow_allowed_effects.filter_not_in(j.effects)","tryCatchPattern":"try: jax.vmap(model)(xs)\\nexcept NotImplementedError as e:\\n    if 'IO effect' in str(e): move callbacks out of cond and retry\\n    else: raise","preventionTips":["Keep io_callback/debug_callback out of batched conditionals","Do IO at trace boundaries: before vmap or on final results","Prefer pure ops inside cond; branch selection via jnp.where"],"tags":["jax","vmap","cond","io-callback","host-callback"],"backgroundTag":"jax-vmap-unsupported-effect","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}