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.

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.
• axis_index_groups not supported in vmap collectives. Please…
• batch_dims must be None or satisfy 0 < dim < n_batch. Got
• out of range for indices with

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

Appendix: 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)