jax-ml/jax · error · Exception

Ordered IO effects not supported in vmap.

Error message

Ordered IO effects not supported in vmap.

What it means

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.

Source

Thrown at jax/_src/lax/control_flow/loops.py:1808

  #   a1, a2 = next((a1, a2) for a1, a2 in zip(body_avals, body_jaxpr.in_avals)
  #                 if not core.typecompat(a1, a2))
  #   raise core.JaxprTypeError(f"while_loop body function input type error: {a1} != {a2}")


  joined_effects = _join_while_effects(body_jaxpr, cond_jaxpr, body_nconsts,
                                       cond_nconsts)
  disallowed_effects = effects.control_flow_allowed_effects.filter_not_in(joined_effects)
  if disallowed_effects:
    raise NotImplementedError(
        f'Effects not supported in `while`: {disallowed_effects}')
  return body_jaxpr.out_avals, joined_effects


def _while_loop_batching_rule(axis_data, args, dims, cond_nconsts, cond_jaxpr,
                              body_nconsts, body_jaxpr):
  from jax._src.callback import _IOEffect, _OrderedIOEffect
  if any(_OrderedIOEffect in fn.effects for fn in [body_jaxpr, cond_jaxpr]):
    raise Exception("Ordered IO effects not supported in vmap.")

  orig_batched = [d is not None for d in dims]
  cconst_bat, bconst_bat, init_bat = split_list(orig_batched, [cond_nconsts, body_nconsts])
  cconsts, bconsts, init = split_list(args, [cond_nconsts, body_nconsts])
  cconst_dims, bconst_dims, init_dims = split_list(dims, [cond_nconsts, body_nconsts])

  carry_bat = init_bat
  # Fixpoint computation of which carry are batched: either
  # batched from init, or the carry out is batched. Each iteration promotes
  # at least one carry to batched. We need at most len(carry) iterations to
  # reach a fixpoint.
  for _ in range(1 + len(carry_bat)):
    _, carry_bat_out = batching.batch_jaxpr(
        body_jaxpr, axis_data, bconst_bat + carry_bat, instantiate=carry_bat)
    if carry_bat == carry_bat_out:
      break
    carry_bat = safe_map(operator.or_, carry_bat, carry_bat_out)
  else:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove or gate the debug print/callback inside the while_loop body before vmap (e.g. only print in unbatched debug runs)
2. 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
3. Debug by running the unvmapped version, or with config.disable_jit, then vmap the clean version

Example fix

// before
def body(c, x):
    jax.debug.print('c {}', c, ordered=True)
    return c + x
jax.vmap(lambda xs: jax.lax.scan(body, 0, xs))(batch)

// after
def body(c, x):
    return c + x  # no ordered callback
out = jax.vmap(lambda xs: jax.lax.scan(body, 0, xs))(batch)
jax.debug.print('out {}', out)  # print outside

import jax
jaxpr = jax.make_jaxpr(lambda: batched_fn(sample_batch))()  # hard to check pre-hoc; simplest guard:
# grep-like check: ensure no jax.debug.print/ordered callbacks in the vmapped function source
import inspect
src = inspect.getsource(fn_to_vmap)
assert 'debug.print' not in src and 'debug.callback' not in src, 'ordered IO blocks vmap'

try:
    jax.vmap(fn)(batch)
except Exception as e:
    if 'Ordered IO effects not supported in vmap' in str(e):
        result = jax.lax.map(fn, batch)  # or strip debug callbacks and re-vmap
    else:
        raise

When it happens

Trigger: 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.

Common situations: 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.

Related errors

• Checkify does not support batched while-loops (checkify-of-v
• Currently only support batch_dim in [0, None], but got {dim=
• for vmap support, subclass {type(self)} must implement `batc
• under vmap, the {member_name} of {name} produced an output b
• Effects not supported in `while`: {}

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