jax-ml/jax · error · NotImplementedError

IO effect not supported in vmap-of-cond.

Error message

IO effect not supported in vmap-of-cond.

What it means

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.

Source

Thrown at jax/_src/lax/control_flow/conditionals.py:472

def _bcast_select_n(pred, *cases):
  if np.ndim(pred) != np.ndim(cases[0]):
    idx = list(range(np.ndim(pred)))
    pred = lax.broadcast_in_dim(pred, np.shape(cases[0]), idx)
  return lax.select_n(pred, *cases)

def _cond_batching_rule(axis_data, args, dims, *, branches, **params):
  index, *ops = args
  index_dim, *op_dims = dims
  # TODO(sharadmv): clean this up by adding a specific blocklist
  if any(isinstance(eff, RefEffect) for branch in branches for eff in
      branch.effects):
    raise NotImplementedError(
        "State effect not supported in vmap-of-cond.")
  from jax._src.callback import _IOEffect, _OrderedIOEffect
  if any(eff in branch.effects for eff in [_IOEffect, _OrderedIOEffect]
      for branch in branches):
    raise NotImplementedError(
        "IO effect not supported in vmap-of-cond.")

  if "branches_platforms" in params and (index_dim is not None):
    # If we end up with a mapped index for a platform_dependent cond, we can
    # replace the index with a fresh call to platform_index. See #29329.
    index = platform_index_p.bind(platforms=params["branches_platforms"])
    index_dim = None

  if index_dim is not None:
    # Convert to a lax.select. While we could get away with not broadcasting
    # some operands yet, because all outputs must be broadcast together anyway
    # for the select we broadcast the input operands for simplicity and leave
    # optimizations to XLA.
    # TODO(mattjj,frostig): assumes branches are side-effect-free, revise!
    index, *ops = (
        batching.bdim_at_front(x, d, axis_data.size,
                               mesh_axis=axis_data.explicit_mesh_axis)
        for x, d in zip(args, dims)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. 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
2. Replace io_callback with pure JAX ops or compile-time (Python) branching when the predicate is concrete
3. Use lax.map instead of vmap if per-tracing IO is acceptable
4. Remove the callback from the batched path entirely

Example fix

// before
@jax.vmap
def f(x):
  return lax.cond(x > 0, x, lambda x: hcb.io_callback(log, x), x, lambda x: x)
// after
@jax.vmap
def f(x):
  y = jnp.where(x > 0, x, x)  # pure compute
  # do IO outside vmap on the result
  return y

import jax
j = jax.make_jaxpr(branch)(*args)
from jax._src.effects import control_flow_allowed_effects
assert not control_flow_allowed_effects.filter_not_in(j.effects), 'IO effects present'

def is_pure_enough(fn, *args) -> bool:
    j = jax.make_jaxpr(fn)(*args)
    return not control_flow_allowed_effects.filter_not_in(j.effects)

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

When it happens

Trigger: jax.vmap over a lax.cond whose true/false branches call host_callback.io_callback, debug_callback, or any primitive carrying _IOEffect/_OrderedIOEffect.

Common situations: Mixing logging, printing, or host-side side effects with vmap inside conditionals; using io_callback for I/O in batched inference code.

Related errors

• State effect not supported in vmap-of-cond.
• 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
• {name} was requested to map a value of non-array type {core.

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