{"record":{"id":"c11bf58feb3163b6","repo":"jax-ml/jax","slug":"the-vjp-function-was-applied-before-restoring-its","errorCode":null,"errorMessage":"the VJP function was applied before restoring its not-saveable residuals.\n\nBecause `saveable_args` was passed to `jax.vjp`, some argument values that\nwould have been saved for the backward pass were instead replaced with\n`NotSaveable()` sentinels. Before the VJP function can be applied, these\nvalues must be restored, e.g. by assigning to the VJP function's `args_res`\nattribute. The values not yet restored correspond to:","messagePattern":"the VJP function was applied before restoring its not-saveable residuals\\.\n\nBecause `saveable_args` was passed to `jax\\.vjp`, some argument values that\nwould have been saved for the backward pass were instead replaced with\n`NotSaveable\\(\\)` sentinels\\. Before the VJP function can be applied, these\nvalues must be restored, e\\.g\\. by assigning to the VJP function's `args_res`\nattribute\\. The values not yet restored correspond to:","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/api.py","lineNumber":1831,"sourceCode":"\"\"\"\n  msg += '\\n'.join(f\"  * args{keystr(path)} was a {thing1} in the primal \"\n                   f\"arguments, but a {thing2} in the `with_refs` arguments, \"\n                   f\"so {explanation}.\"\n                   for path, thing1, thing2, explanation\n                   in equality_errors_pytreedef(in_tree, refs_tree))\n  raise ValueError(msg)\n\ndef _vjp_not_saveable_error(jaxpr, in_tree, idxs):\n  msg = \"\"\"the VJP function was applied before restoring its not-saveable residuals.\n\nBecause `saveable_args` was passed to `jax.vjp`, some argument values that\nwould have been saved for the backward pass were instead replaced with\n`NotSaveable()` sentinels. Before the VJP function can be applied, these\nvalues must be restored, e.g. by assigning to the VJP function's `args_res`\nattribute. The values not yet restored correspond to:\n\"\"\"\n  msg += '\\n'.join(f\"  * {_vjp_arg_name(jaxpr, in_tree, idx)};\" for idx in idxs)\n  raise ValueError(msg)\n\ndef check_accum(aval, acc):\n  if not core.typecompat(acc.aval, aval):\n    raise ValueError(f\"Accumulator aval mismatch: expected {aval}, got {acc.aval}\")\n  return acc\n\ndef _vjp3_bwd(in_tree, out_tree, out_zeros, jaxpr, out_primal_avals, want_logs,\n              residuals, structured_res, maybe_accums, out_ct):\n  cts_flat, out_tree_ = tree_flatten(out_ct, is_leaf=lambda x: isinstance(x, ad.Zero))\n  if out_tree != out_tree_:\n    _vjp_ct_tree_error(jaxpr, out_tree, out_tree_)\n  _vjp_check_ct_avals(cts_flat, out_primal_avals)\n  cts_flat = [ct for ct, k in zip(cts_flat, out_zeros) if not k]\n  primals_in = [*maybe_accums, *tree_leaves(structured_res)]\n  logs = ad.backward_pass3(jaxpr, True, residuals, primals_in, cts_flat)\n  arg_cts = [x.freeze() if isinstance(x, ad.ValAccum) else\n             DidntWant() if isinstance(x, ad.NullAccum) else GradRef()\n             for x in maybe_accums]","sourceCodeStart":1813,"sourceCodeEnd":1849,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/api.py#L1813-L1849","documentation":"jax.vjp can be called with saveable_args to replace some saved-for-backward values with NotSaveable() sentinels (to avoid holding them in memory). The returned vjp function must have those values restored (via its args_res attribute) before being applied; this error lists the argument paths still unrestored.","triggerScenarios":"jax.vjp(f, x, saveable_args=(x,)) followed by calling the vjp function without assigning the real values back to vjp_fn.args_res.","commonSituations":"Memory-optimization pipelines where large activations/inputs are dropped and re-supplied later (e.g. checkpointing across steps or processes); forgetting the restore step after adding saveable_args for memory savings.","solutions":["Restore the values before applying: vjp_fn.args_res = (real_values,...) or assign the specific leaves","Only pass values in saveable_args that you can reliably re-supply before backward","Drop saveable_args if simpler memory strategies (checkpointing/remat) suffice"],"exampleFix":"# before\n_, f_vjp = jax.vjp(f, x, saveable_args=(x,))\ng = f_vjp(ct)\n# after\n_, f_vjp = jax.vjp(f, x, saveable_args=(x,))\nf_vjp.args_res = (x,)\ng = f_vjp(ct)","handlingStrategy":"validation","validationCode":"unrestored = [i for i, v in enumerate(f_vjp.args_res) if isinstance(v, ad.NotSaveable)]\nassert not unrestored, f'restore args_res for indices {unrestored} before applying vjp'","typeGuard":"def vjp_ready(f_vjp): return not any(isinstance(v, ad.NotSaveable) for v in tree_leaves(f_vjp.args_res))","tryCatchPattern":null,"preventionTips":["Restore vjp_fn.args_res immediately after the forward pass when using saveable_args","Keep the values needed for restoration alive in the caller","Prefer jax.checkpoint/remat unless saveable_args is required"],"tags":["jax","vjp","memory","checkpointing","saveable-args"],"backgroundTag":"uninitialized-state","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}