jax-ml/jax · error · Exception

performing a set/swap operation with a differentiated value

Error message

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?

What it means

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.

Source

Thrown at jax/_src/state/primitives.py:599

  out_primal = get_p.bind(ref_primal, *idx, **params)
  if isinstance(ref_tangent, ad_util.Zero):
    out_tangent = ad_util.Zero(core.typeof(out_primal).to_tangent_aval())
  else:
    out_tangent = get_p.bind(ref_tangent, *idx, **params)
  return out_primal, out_tangent
ad.primitive_jvps[get_p] = _get_jvp

def _swap_jvp(primals: list[Any], tangents: list[Any], **params: Any):
  ref_primal, x_primal, *idx = primals
  ref_tangent, x_tangent, *_ = tangents
  out_primal = swap_p.bind(ref_primal, x_primal, *idx, **params)
  if isinstance(ref_tangent, ad_util.Zero) and isinstance(x_tangent, ad_util.Zero):
    out_tangent = ad_util.Zero(core.typeof(out_primal).to_tangent_aval())
  elif ref_tangent.aval.kind == "no_grad_no_remat":
    out_tangent = ad_util.Zero(core.typeof(out_primal).to_tangent_aval())
  else:
    if isinstance(ref_tangent, ad_util.Zero):
      raise Exception("performing a set/swap operation with a differentiated "
                      "value on a non-differentiated array reference of type "
                      f"{core.typeof(ref_primal)}. Move the array reference "
                      "to be an argument of the differentiated function?")
    x_tangent = ad_util.instantiate(x_tangent)
    out_tangent = swap_p.bind(ref_tangent, x_tangent, *idx, **params)
  return out_primal, out_tangent
ad.primitive_jvps[swap_p] = _swap_jvp

def addupdate_jvp_rule(primals: list[Any], tangents: list[Any], **params: Any):
  ref_primal, x_primal, *idx = primals
  ref_tangent, x_tangent, *_ = tangents
  x_tangent = ad_util.instantiate(x_tangent)
  if ref_tangent.aval.kind != "no_grad_no_remat":
    addupdate_p.bind(ref_primal, x_primal, *idx, **params)
    addupdate_p.bind(ref_tangent, x_tangent, *idx, **params)
  return [], []
ad.primitive_jvps[addupdate_p] = addupdate_jvp_rule

View on GitHub (pinned to 1e1c6a8fc0)

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.

Example fix

// before
def loss(x):
    return ref.swap(0, x)  # ref closed over
jax.grad(loss)(x)
// after
def loss(x, ref):
    return ref.swap(0, x)
jax.grad(lambda x, r: loss(x, r).sum())(x, ref)

Defensive patterns

Strategy: validation

Validate before calling

x = jax.lax.stop_gradient(x) if not need_grad_through_state else x
# and thread refs through the differentiated function's arguments:
def f(x, ref): ...

Prevention

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.

When it happens

Trigger: 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.

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

Related errors

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