jax-ml/jax · error · Exception

Unordered IO effects not supported in while_loop with batche

Error message

Unordered IO effects not supported in while_loop with batched predicate

What it means

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.

Source

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

  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:
    assert False, "Fixpoint not reached"

  # Knowing how the carry is batched now, we can determine if the predicate is
  # batched.
  _, (pred_bat,) = batching.batch_jaxpr(
      cond_jaxpr, axis_data, cconst_bat + carry_bat, instantiate=False)

  if pred_bat:
    # If the predicate is batched, we have to batch *all* of the carry
    # regardless of if the body needs it.
    if any(_IOEffect in fn.effects for fn in [body_jaxpr, cond_jaxpr]):
      raise Exception("Unordered IO effects not supported in while_loop "
                      "with batched predicate")
    carry_bat = [True] * len(carry_bat)
    carry_dims = [0] * len(carry_bat)
    body_jaxpr_batched, _ = batching.batch_jaxpr_axes(
        body_jaxpr, axis_data, bconst_dims + carry_dims, carry_dims)
    cond_jaxpr_batched, _ = batching.batch_jaxpr_axes(
        cond_jaxpr, axis_data, cconst_dims + carry_dims, [0])
  else:
    # If the predicate is not batched, we can look at the `cond_jaxpr`'s out
    # shape to determine the rank of the predicate. From this rank we pick the
    # dims of the carry to be batched to ensure that the predicate shape is a
    # prefix of the carry in and out shapes. We can then batch the `body_jaxpr`
    # according to these new batch dims.
    cond_rank = len(cond_jaxpr.out_avals[0].shape)
    carry_dims = [cond_rank if b else None for b in carry_bat]
    body_jaxpr_batched, _ = batching.batch_jaxpr_axes(
        body_jaxpr, axis_data, bconst_dims + carry_dims, carry_dims)
    # Now we need to rebatch the `cond_jaxpr` according to the new dims of the

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. 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
2. 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)
3. Collect values in the loop carry and print/inspect them after the vmap call outside of the loop
4. If a fixed number of iterations is possible, replace while_loop with lax.scan, which supports batching with effects

Example fix

// before
jax.vmap(lambda x: lax.while_loop(lambda c: c[0] < 10,
                                   lambda c: (jax.debug.print('{}', c[0]), (c[0]+1, c[1]*x))[1],
                                   (0, 1.0)))
// after (no IO effect inside loop)
jit_fn = jax.vmap(lambda x: lax.fori_loop(0, 10, lambda i, acc: acc * x, 1.0))

from jax._src import effects
import jax
# before vmap, check for IO effects by tracing
def has_io_effects(fun, *args):
    jaxpr = jax.make_jaxpr(fun)(*args)
    return any(getattr(e, 'name', '') == 'IO' or 'IO' in type(e).__name__ for e in jaxpr.effects)

def loop_is_vmappable_with_batched_pred(cond_fn, body_fn, example_carry):
    jaxpr_c = jax.make_jaxpr(cond_fn)(example_carry)
    jaxpr_b = jax.make_jaxpr(body_fn)(example_carry)
    return not any('IO' in type(e).__name__ for e in (*jaxpr_c.effects, *jaxpr_b.effects))

When it happens

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

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

Related errors

• Checkify does not support batched while-loops (checkify-of-v
• Ordered IO effects not supported in vmap.
• 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

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