jax-ml/jax · error · ValueError

Cannot use `partial_eval_jaxpr_custom` with stateful jaxprs.

Error message

Cannot use `partial_eval_jaxpr_custom` with stateful jaxprs.

What it means

partial_eval_jaxpr_custom splits a jaxpr into known/unknown parts for custom rules (remat, cond, scan, while). The stateful variant additionally returns residual count; if that count is nonzero the jaxpr carries state (mutable arrays / state effects), which the non-stateful custom API cannot represent, so it refuses.

Source

Thrown at jax/_src/interpreters/partial_eval.py:799

  assert next(residuals_, None) is None
  idxs = {id(x): i for i, x in enumerate((*consts, *dummys))}
  fwds = [idxs.get(id(r)) for r in residuals]
  residuals = [r for r in residuals if id(r) not in idxs]
  return fwds, residuals


def partial_eval_jaxpr_custom(
    jaxpr: Jaxpr,
    in_unknowns: Sequence[bool],
    in_inst: bool | Sequence[bool],
    ensure_out_unknowns: bool | Sequence[bool],
    ensure_out_inst: bool | Sequence[bool],
    saveable: Callable[..., RematCases_],
  ) -> tuple[Jaxpr, Jaxpr, list[bool], list[bool], int]:
  *outs, num_res_ref = partial_eval_jaxpr_stateful(
      jaxpr, in_unknowns, in_inst, ensure_out_unknowns, ensure_out_inst, saveable)
  if num_res_ref:
    raise ValueError("Cannot use `partial_eval_jaxpr_custom` with stateful jaxprs.")
  return *outs,  # pyrefly: ignore[bad-return]

def partial_eval_jaxpr_stateful(
    jaxpr: Jaxpr,
    in_unknowns: Sequence[bool],
    in_inst: bool | Sequence[bool],
    ensure_out_unknowns: bool | Sequence[bool],
    ensure_out_inst: bool | Sequence[bool],
    saveable: Callable[..., RematCases_] | None,
  ) -> tuple[Jaxpr, Jaxpr, list[bool], list[bool], int, int]:
  if type(in_inst) is bool:
    in_inst = (in_inst,) * len(jaxpr.invars)
  if type(ensure_out_unknowns) is bool:
    ensure_out_unknowns = (ensure_out_unknowns,) * len(jaxpr.outvars)
  if type(ensure_out_inst) is bool:
    ensure_out_inst = (ensure_out_inst,) * len(jaxpr.outvars)
  if saveable is None:
    saveable = everything_saveable

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Move the stateful operations (Ref get/put, mutable state) outside the remat/custom-transformed region
Re-express the computation functionally: pass state in and out as explicit arguments instead of mutating Refs
If you implement the custom rule, switch to partial_eval_jaxpr_stateful which supports state

Example fix

# before
@jax.remat
def f(ref, x):
    ref[...] += x  # stateful inside remat

# after
def f(ref_val, x):
    return ref_val + x  # stateless; update ref outside

Defensive patterns

Strategy: validation

Validate before calling

jaxpr = jax.make_jaxpr(fn_with_state)(*args).jaxpr
# check statefulness indicators before applying remat/custom transforms
has_state = any('ref' in str(v.aval) for v in jaxpr.invars)
if has_state:
    raise RuntimeError('refactor to stateless before remat/custom rule')

Prevention

Keep remat regions purely functional
Thread state explicitly instead of mutating Refs inside transforms

When it happens

Trigger: A custom rule (e.g. remat_partial_eval or a custom bind_with state primitives like get/put/ref) is invoked on a jaxpr containing stateful operations. Typically triggered when jax.remat or control-flow transforms encounter jax.experimental.array_stats style state or new state primitives.

Common situations: Using jax.remat (or custom partitioning/_find_downstream paths) around code that mutates Refs / uses experimental state APIs, where the codepath only supports stateless jaxprs.

Related errors

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