jax-ml/jax · error · NotImplementedError

State effect not supported in cond partial-eval.

Error message

State effect not supported in cond partial-eval.

What it means

Raised during partial evaluation of lax.cond when a branch contains RefEffects (state mutation). Partial eval cannot stage half of a stateful cond, because unknown values would leave the ref updates undefined, so it errors.

Source

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

  index, *ops = primals
  _, *ops_dot = tangents
  ops_dot = _prune_zeros(ops_dot)

  out = cond_p.bind(index, *ops, *ops_dot, branches=branches_jvp,
                    **params)
  out_primals, out_tangents = split_list(out, [len(out_nz)])
  out_tangents_iter = iter(out_tangents)
  out_tangents = [next(out_tangents_iter) if nz else
                  ad_util.p2tz(p)
                  for p, nz in zip(out_primals, out_nz)]
  return out_primals, out_tangents

def _cond_partial_eval(trace, *tracers, branches, **params):
  in_unknowns = [not t.pval.is_known() for t in tracers]
  index_uk, *ops_uk = in_unknowns
  if any(isinstance(eff, RefEffect) for branch in branches for eff in
      branch.effects):
    raise NotImplementedError(
        "State effect not supported in cond partial-eval.")

  if index_uk:
    # When the branch index is unknown, we stage out the whole cond.
    # TODO(mattjj): remove this path when old remat is removed
    params = dict(branches=branches, **params)
    return trace.default_process_primitive(cond_p, tracers, params)

  branches_out_uks = []
  for branch_jaxpr in branches:
    _, _, out_uks, _ = pe.partial_eval_jaxpr_nounits(
        branch_jaxpr, ops_uk, instantiate=False)
    branches_out_uks.append(out_uks)
  out_uks = [any(uks) for uks in zip(*branches_out_uks)]

  branches_known, branches_unknown, branch_res_avals = [], [], []
  for branch_jaxpr in branches:
    branch_jaxpr_known, branch_jaxpr_unknown, _, res_avals = \

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Make branches pure: read refs before the cond and write after it, selecting values with jnp.where
2. Avoid autodiff over state-mutating conds; restructure to update state outside conditional logic
3. Upgrade JAX: newer versions progressively support state effects in more transforms
4. Fall back to classic lax.fori_loop/scan with explicit carry instead of refs

Example fix

// before
val = ref_get(ref)
out = lax.cond(flag, lambda: val + 1, lambda: val)
ref_set(ref, out)
# inside grad-traced code with refs in cond branches
// after
val = ref_get(ref)
out = jnp.where(flag, val + 1, val)  # autodiff-safe
ref_set(ref, out)

null

null

try: jax.grad(f)(x)\nexcept NotImplementedError as e:\n    if 'State effect not supported in cond partial-eval' in str(e):\n        return functional_version(f)(x)  # ref-free reimplementation\n    raise

When it happens

Trigger: Running gradient (grad), jit with unknown-constant inputs, or any transform that triggers pe.trace_to_jaxpr over a cond whose branches mutate refs.

Common situations: Differentiating or rematerializing through experimental stateful code that uses lax.cond; combining new-style state with jax.grad or custom JVPs.

Related errors

• State effect not supported in vmap-of-cond.
• primal and tangent arguments to jax.jvp must be tuples or li
• primal and tangent arguments to jax.jvp must have the same t
• primal and tangent arguments to jax.jvp do not match; dtypes
• jvp called with different primal and tangent shapes;Got prim

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