jax-ml/jax · error · ValueError

the VJP function was applied before restoring its not-saveab

Error message

the VJP function was applied before restoring its not-saveable residuals.

Because `saveable_args` was passed to `jax.vjp`, some argument values that
would have been saved for the backward pass were instead replaced with
`NotSaveable()` sentinels. Before the VJP function can be applied, these
values must be restored, e.g. by assigning to the VJP function's `args_res`
attribute. The values not yet restored correspond to:

What it means

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.

Source

Thrown at jax/_src/api.py:1831

"""
  msg += '\n'.join(f"  * args{keystr(path)} was a {thing1} in the primal "
                   f"arguments, but a {thing2} in the `with_refs` arguments, "
                   f"so {explanation}."
                   for path, thing1, thing2, explanation
                   in equality_errors_pytreedef(in_tree, refs_tree))
  raise ValueError(msg)

def _vjp_not_saveable_error(jaxpr, in_tree, idxs):
  msg = """the VJP function was applied before restoring its not-saveable residuals.

Because `saveable_args` was passed to `jax.vjp`, some argument values that
would have been saved for the backward pass were instead replaced with
`NotSaveable()` sentinels. Before the VJP function can be applied, these
values must be restored, e.g. by assigning to the VJP function's `args_res`
attribute. The values not yet restored correspond to:
"""
  msg += '\n'.join(f"  * {_vjp_arg_name(jaxpr, in_tree, idx)};" for idx in idxs)
  raise ValueError(msg)

def check_accum(aval, acc):
  if not core.typecompat(acc.aval, aval):
    raise ValueError(f"Accumulator aval mismatch: expected {aval}, got {acc.aval}")
  return acc

def _vjp3_bwd(in_tree, out_tree, out_zeros, jaxpr, out_primal_avals, want_logs,
              residuals, structured_res, maybe_accums, out_ct):
  cts_flat, out_tree_ = tree_flatten(out_ct, is_leaf=lambda x: isinstance(x, ad.Zero))
  if out_tree != out_tree_:
    _vjp_ct_tree_error(jaxpr, out_tree, out_tree_)
  _vjp_check_ct_avals(cts_flat, out_primal_avals)
  cts_flat = [ct for ct, k in zip(cts_flat, out_zeros) if not k]
  primals_in = [*maybe_accums, *tree_leaves(structured_res)]
  logs = ad.backward_pass3(jaxpr, True, residuals, primals_in, cts_flat)
  arg_cts = [x.freeze() if isinstance(x, ad.ValAccum) else
             DidntWant() if isinstance(x, ad.NullAccum) else GradRef()
             for x in maybe_accums]

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Restore the values before applying: vjp_fn.args_res = (real_values,...) or assign the specific leaves
2. Only pass values in saveable_args that you can reliably re-supply before backward
3. Drop saveable_args if simpler memory strategies (checkpointing/remat) suffice

Example fix

# before
_, f_vjp = jax.vjp(f, x, saveable_args=(x,))
g = f_vjp(ct)
# after
_, f_vjp = jax.vjp(f, x, saveable_args=(x,))
f_vjp.args_res = (x,)
g = f_vjp(ct)

unrestored = [i for i, v in enumerate(f_vjp.args_res) if isinstance(v, ad.NotSaveable)]
assert not unrestored, f'restore args_res for indices {unrestored} before applying vjp'

def vjp_ready(f_vjp): return not any(isinstance(v, ad.NotSaveable) for v in tree_leaves(f_vjp.args_res))

When it happens

Trigger: jax.vjp(f, x, saveable_args=(x,)) followed by calling the vjp function without assigning the real values back to vjp_fn.args_res.

Common situations: 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.

Related errors

• reduce_axes argument to vjp is deprecated
• unexpected JAX type (e.g. shape/dtype) for gradient ref pass
• the gradient for {_vjp_arg_name(jaxpr, in_tree, idx)}, which
• {_vjp_arg_name(jaxpr, in_tree, idx)} is Ref-typed, so its gr
• But the tree structures differ:

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