{"record":{"id":"f15be6fb264407c1","repo":"jax-ml/jax","slug":"performing-a-set-swap-operation-with-a-differentia","errorCode":null,"errorMessage":"performing a set/swap operation with a differentiated value on a non-differentiated array reference of type {core.typeof(ref_primal)}. Move the array reference to be an argument of the differentiated function?","messagePattern":"performing a set/swap operation with a differentiated value on a non-differentiated array reference of type (.+?)\\. Move the array reference to be an argument of the differentiated function\\?","errorType":"exception","errorClass":"Exception","httpStatus":null,"severity":"error","filePath":"jax/_src/state/primitives.py","lineNumber":599,"sourceCode":"  out_primal = get_p.bind(ref_primal, *idx, **params)\n  if isinstance(ref_tangent, ad_util.Zero):\n    out_tangent = ad_util.Zero(core.typeof(out_primal).to_tangent_aval())\n  else:\n    out_tangent = get_p.bind(ref_tangent, *idx, **params)\n  return out_primal, out_tangent\nad.primitive_jvps[get_p] = _get_jvp\n\ndef _swap_jvp(primals: list[Any], tangents: list[Any], **params: Any):\n  ref_primal, x_primal, *idx = primals\n  ref_tangent, x_tangent, *_ = tangents\n  out_primal = swap_p.bind(ref_primal, x_primal, *idx, **params)\n  if isinstance(ref_tangent, ad_util.Zero) and isinstance(x_tangent, ad_util.Zero):\n    out_tangent = ad_util.Zero(core.typeof(out_primal).to_tangent_aval())\n  elif ref_tangent.aval.kind == \"no_grad_no_remat\":\n    out_tangent = ad_util.Zero(core.typeof(out_primal).to_tangent_aval())\n  else:\n    if isinstance(ref_tangent, ad_util.Zero):\n      raise Exception(\"performing a set/swap operation with a differentiated \"\n                      \"value on a non-differentiated array reference of type \"\n                      f\"{core.typeof(ref_primal)}. Move the array reference \"\n                      \"to be an argument of the differentiated function?\")\n    x_tangent = ad_util.instantiate(x_tangent)\n    out_tangent = swap_p.bind(ref_tangent, x_tangent, *idx, **params)\n  return out_primal, out_tangent\nad.primitive_jvps[swap_p] = _swap_jvp\n\ndef addupdate_jvp_rule(primals: list[Any], tangents: list[Any], **params: Any):\n  ref_primal, x_primal, *idx = primals\n  ref_tangent, x_tangent, *_ = tangents\n  x_tangent = ad_util.instantiate(x_tangent)\n  if ref_tangent.aval.kind != \"no_grad_no_remat\":\n    addupdate_p.bind(ref_primal, x_primal, *idx, **params)\n    addupdate_p.bind(ref_tangent, x_tangent, *idx, **params)\n  return [], []\nad.primitive_jvps[addupdate_p] = addupdate_jvp_rule\n","sourceCodeStart":581,"sourceCodeEnd":617,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/state/primitives.py#L581-L617","documentation":"In the JVP rule for `swap`, JAX requires that if the value written carries a tangent (is differentiated), the ref must also carry a tangent. When ref_tangent is Zero but x_tangent is not, differentiation cannot propagate and JAX raises this with a hint to make the ref an argument of the differentiated function.","triggerScenarios":"Running `jax.grad`/`jax.jvp`/`jax.vjp` over a function that swaps a differentiated value into a Ref that was not passed as an (differentiable) argument of the function — e.g. a ref captured from an enclosing scope or created via a non-differentiated path.","commonSituations":"Autodiff through ODE solvers / lax.while_loop-style in-place state updates where the state buffer isn't threaded as a function argument; using Refs created outside the traced function; grad of functions using `ref.set`/`swap` on closed-over buffers.","solutions":["Move the Ref inside the differentiated function's argument list so JAX threads a tangent through it.","Recompute the ref inside the traced function (allocate buffers within the jvp-traced region).","If gradients through the state aren't needed, stop-gradient the value before swapping: `lax.stop_gradient(x)`.","Use higher-level APIs (e.g. custom_vjp or structural state handling) instead of raw swap under autodiff."],"exampleFix":"// before\ndef loss(x):\n    return ref.swap(0, x)  # ref closed over\njax.grad(loss)(x)\n// after\ndef loss(x, ref):\n    return ref.swap(0, x)\njax.grad(lambda x, r: loss(x, r).sum())(x, ref)","handlingStrategy":"validation","validationCode":"x = jax.lax.stop_gradient(x) if not need_grad_through_state else x\n# and thread refs through the differentiated function's arguments:\ndef f(x, ref): ...","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Always pass Refs as explicit arguments to functions under jax.grad/jvp.","Never capture refs from enclosing scopes inside differentiated code.","Stop-gradient values written into state you don't need gradients through."],"tags":["jax","autodiff","state-primitives","jvp"],"backgroundTag":"autodiff-through-mutated-state","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}