jax-ml/jax · error · NotImplementedError

State effect not supported in vmap-of-cond.

Error message

State effect not supported in vmap-of-cond.

What it means

Raised when jax.lax.cond is used under vmap (batching) and one of the cond branches contains a RefEffect (i.e., mutates a state ref via experimental state primitives). The batching rule for cond has no way to soundly batch state mutations across branches, so it refuses.

Source

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

def _bcast_select(pred, on_true, on_false):
  if np.ndim(pred) != np.ndim(on_true):
    idx = list(range(np.ndim(pred)))
    pred = lax.broadcast_in_dim(pred, np.shape(on_true), idx)
  return lax.select(pred, on_true, on_false)

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.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Hoist the state mutation out of the cond so branches are pure and return values instead
2. Use jax.lax.switch-free restructuring: compute both branch results and select with jnp.where
3. Drop vmap and use an explicit batched loop (lax.map or manual axis handling)
4. Check jax release notes for newer state-effect batching support and upgrade

Example fix

// before
@jax.vmap
def f(x, ref):
  return lax.cond(x > 0, lambda: ref_set(ref, 1), lambda: ref_set(ref, 0))
// after
@jax.vmap
def f(x, ref):
  ref_set(ref, (x > 0).astype(jnp.int32))
  return None

from jax._src.effects import control_flow_allowed_effects
import jax
# trace branches' jaxprs and check effects before vmap
jaxprs = [jax.make_jaxpr(branch)(*example_args) for branch in branches]
bad = [j.effects for j in jaxprs if control_flow_allowed_effects.filter_not_in(j.effects)]
assert not bad, f'RefEffects in branches: {bad}'

null

try: jax.vmap(f)(x)\nexcept NotImplementedError as e:\n    if 'State effect' in str(e): restructure without refs in cond\n    else: raise

When it happens

Trigger: Calling jax.vmap over a function that uses lax.cond where either branch reads/writes a jax.experimental.ref or otherwise produces RefEffects in its jaxpr.

Common situations: Using experimental stateful code (ref_get/ref_set, while_state, new-style RNG or mutable arrays) inside a conditional that is then batched with vmap or vjp-of-vmap.

Related errors

• IO effect not supported in vmap-of-cond.
• State effect not supported in cond partial-eval.
• 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/edb440f2797b63ac. Report an issue: GitHub.