{"record":{"id":"edb440f2797b63ac","repo":"jax-ml/jax","slug":"state-effect-not-supported-in-vmap-of-cond","errorCode":null,"errorMessage":"State effect not supported in vmap-of-cond.","messagePattern":"State effect not supported in vmap-of-cond\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/control_flow/conditionals.py","lineNumber":467,"sourceCode":"def _bcast_select(pred, on_true, on_false):\n  if np.ndim(pred) != np.ndim(on_true):\n    idx = list(range(np.ndim(pred)))\n    pred = lax.broadcast_in_dim(pred, np.shape(on_true), idx)\n  return lax.select(pred, on_true, on_false)\n\ndef _bcast_select_n(pred, *cases):\n  if np.ndim(pred) != np.ndim(cases[0]):\n    idx = list(range(np.ndim(pred)))\n    pred = lax.broadcast_in_dim(pred, np.shape(cases[0]), idx)\n  return lax.select_n(pred, *cases)\n\ndef _cond_batching_rule(axis_data, args, dims, *, branches, **params):\n  index, *ops = args\n  index_dim, *op_dims = dims\n  # TODO(sharadmv): clean this up by adding a specific blocklist\n  if any(isinstance(eff, RefEffect) for branch in branches for eff in\n      branch.effects):\n    raise NotImplementedError(\n        \"State effect not supported in vmap-of-cond.\")\n  from jax._src.callback import _IOEffect, _OrderedIOEffect\n  if any(eff in branch.effects for eff in [_IOEffect, _OrderedIOEffect]\n      for branch in branches):\n    raise NotImplementedError(\n        \"IO effect not supported in vmap-of-cond.\")\n\n  if \"branches_platforms\" in params and (index_dim is not None):\n    # If we end up with a mapped index for a platform_dependent cond, we can\n    # replace the index with a fresh call to platform_index. See #29329.\n    index = platform_index_p.bind(platforms=params[\"branches_platforms\"])\n    index_dim = None\n\n  if index_dim is not None:\n    # Convert to a lax.select. While we could get away with not broadcasting\n    # some operands yet, because all outputs must be broadcast together anyway\n    # for the select we broadcast the input operands for simplicity and leave\n    # optimizations to XLA.","sourceCodeStart":449,"sourceCodeEnd":485,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/control_flow/conditionals.py#L449-L485","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Hoist the state mutation out of the cond so branches are pure and return values instead","Use jax.lax.switch-free restructuring: compute both branch results and select with jnp.where","Drop vmap and use an explicit batched loop (lax.map or manual axis handling)","Check jax release notes for newer state-effect batching support and upgrade"],"exampleFix":"// before\n@jax.vmap\ndef f(x, ref):\n  return lax.cond(x > 0, lambda: ref_set(ref, 1), lambda: ref_set(ref, 0))\n// after\n@jax.vmap\ndef f(x, ref):\n  ref_set(ref, (x > 0).astype(jnp.int32))\n  return None","handlingStrategy":"validation","validationCode":"from jax._src.effects import control_flow_allowed_effects\nimport jax\n# trace branches' jaxprs and check effects before vmap\njaxprs = [jax.make_jaxpr(branch)(*example_args) for branch in branches]\nbad = [j.effects for j in jaxprs if control_flow_allowed_effects.filter_not_in(j.effects)]\nassert not bad, f'RefEffects in branches: {bad}'","typeGuard":"null","tryCatchPattern":"try: jax.vmap(f)(x)\\nexcept NotImplementedError as e:\\n    if 'State effect' in str(e): restructure without refs in cond\\n    else: raise","preventionTips":["Keep cond branches pure; do all ref updates outside conditionals","Never place ref_get/ref_set inside code that will be vmap'd","Wrap stateful experiment code behind a pure functional interface"],"tags":["jax","vmap","cond","state-effects","experimental"],"backgroundTag":"jax-vmap-unsupported-effect","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}