{"record":{"id":"e68df60bfc644493","repo":"jax-ml/jax","slug":"cannot-use-partial-eval-jaxpr-custom-with-statef","errorCode":null,"errorMessage":"Cannot use `partial_eval_jaxpr_custom` with stateful jaxprs.","messagePattern":"Cannot use `partial_eval_jaxpr_custom` with stateful jaxprs\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/interpreters/partial_eval.py","lineNumber":799,"sourceCode":"  assert next(residuals_, None) is None\n  idxs = {id(x): i for i, x in enumerate((*consts, *dummys))}\n  fwds = [idxs.get(id(r)) for r in residuals]\n  residuals = [r for r in residuals if id(r) not in idxs]\n  return fwds, residuals\n\n\ndef partial_eval_jaxpr_custom(\n    jaxpr: Jaxpr,\n    in_unknowns: Sequence[bool],\n    in_inst: bool | Sequence[bool],\n    ensure_out_unknowns: bool | Sequence[bool],\n    ensure_out_inst: bool | Sequence[bool],\n    saveable: Callable[..., RematCases_],\n  ) -> tuple[Jaxpr, Jaxpr, list[bool], list[bool], int]:\n  *outs, num_res_ref = partial_eval_jaxpr_stateful(\n      jaxpr, in_unknowns, in_inst, ensure_out_unknowns, ensure_out_inst, saveable)\n  if num_res_ref:\n    raise ValueError(\"Cannot use `partial_eval_jaxpr_custom` with stateful jaxprs.\")\n  return *outs,  # pyrefly: ignore[bad-return]\n\ndef partial_eval_jaxpr_stateful(\n    jaxpr: Jaxpr,\n    in_unknowns: Sequence[bool],\n    in_inst: bool | Sequence[bool],\n    ensure_out_unknowns: bool | Sequence[bool],\n    ensure_out_inst: bool | Sequence[bool],\n    saveable: Callable[..., RematCases_] | None,\n  ) -> tuple[Jaxpr, Jaxpr, list[bool], list[bool], int, int]:\n  if type(in_inst) is bool:\n    in_inst = (in_inst,) * len(jaxpr.invars)\n  if type(ensure_out_unknowns) is bool:\n    ensure_out_unknowns = (ensure_out_unknowns,) * len(jaxpr.outvars)\n  if type(ensure_out_inst) is bool:\n    ensure_out_inst = (ensure_out_inst,) * len(jaxpr.outvars)\n  if saveable is None:\n    saveable = everything_saveable","sourceCodeStart":781,"sourceCodeEnd":817,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/interpreters/partial_eval.py#L781-L817","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","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"],"exampleFix":"# before\n@jax.remat\ndef f(ref, x):\n    ref[...] += x  # stateful inside remat\n\n# after\ndef f(ref_val, x):\n    return ref_val + x  # stateless; update ref outside","handlingStrategy":"validation","validationCode":"jaxpr = jax.make_jaxpr(fn_with_state)(*args).jaxpr\n# check statefulness indicators before applying remat/custom transforms\nhas_state = any('ref' in str(v.aval) for v in jaxpr.invars)\nif has_state:\n    raise RuntimeError('refactor to stateless before remat/custom rule')","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Keep remat regions purely functional","Thread state explicitly instead of mutating Refs inside transforms"],"tags":["jax","remat","partial-eval","stateful-jaxpr","control-flow"],"backgroundTag":"unsupported-operation-for-codepath","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}