---
jupytext:
  cell_metadata_filter: -all
  formats: ipynb,md:myst,py
  main_language: python
  text_representation:
    extension: .md
    format_name: myst
    format_version: 0.13
    jupytext_version: 1.16.4
---

```{raw-cell}

---
Copyright 2025 The JAX Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    https://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.

---
```

# `Ref`: mutable arrays for data plumbing and memory control

JAX `Array`s are immutable, representing mathematical values. Immutability can
make code easier to reason about, and is useful for optimized compilation,
parallelization, rematerialization, and transformations like autodiff.

But immutability is constraining too:
* **expressiveness** --- plumbing out intermediate data or maintaining state,
  e.g. for normalization statistics or metrics, can feel heavyweight;
* **performance** --- it's more difficult to reason about performance, like
  memory lifetimes and in-place updates.

`Ref`s can help! They represent mutable arrays that can be read and written
in-place. These array references are compatible with JAX transformations, like
`jax.jit` and `jax.grad`:

```{code-cell}
import jax
import jax.numpy as jnp

x_ref = jax.new_ref(jnp.zeros(3))  # new array ref, with initial value [0., 0., 0.]

@jax.jit
def f():
  x_ref[1] += 1.  # indexed add-update

print(x_ref)  # Ref([0., 0., 0.])
f()
f()
print(x_ref)  # Ref([0., 2., 0.])
```

The indexing syntax follows NumPy's. For a `Ref` called `x_ref`, we can
read its entire value into an `Array` by writing `x_ref[...]`, and write its
entire value using `x_ref[...] = A` for some `Array`-valued expression `A`:

```{code-cell}
def g(x):
  x_ref = jax.new_ref(0.)
  x_ref[...] = jnp.sin(x)
  return x_ref[...]

print(jax.grad(g)(1.0))  # 0.54
```

`Ref` is a distinct type from `Array`, and it comes with some important
constraints and limitations. In particular, indexed reading and writing is just
about the *only* thing you can do with an `Ref`. References can't be passed
where `Array`s are expected:

```{code-cell}
x_ref = jax.new_ref(1.0)
try:
  jnp.sin(x_ref)  # error! can't do math on refs
except Exception as e:
  print(e)
```

To do math, you need to read the ref's value first, like `jnp.sin(x_ref[...])`.

So what _can_ you do with `Ref`? Read on for the details, and some useful
recipes.

### API

If you've ever used
[Pallas](https://docs.jax.dev/en/latest/pallas/quickstart.html), then `Ref`
should look familiar. A big difference is that you can create new `Ref`s
yourself directly using `jax.new_ref`:

```{code-cell}
from jax import Array, Ref

def array_ref(init_val: Array) -> Ref:
  """Introduce a new reference with given initial value."""
```

`jax.freeze` is its antithesis, invalidating the given ref (so that accessing it
afterwards is an error) and producing its final value:

```{code-cell}
def freeze(ref: Ref) -> Array:
  """Invalidate given reference and produce its final value."""
```

In between creating and destroying them, you can perform indexed reads and
writes on refs. You can read and write using the functions `jax.ref.get` and
`jax.ref.swap`, but usually you'd just use NumPy-style array indexing syntax:

```{code-cell}
import types
Index = int | slice | Array | types.EllipsisType
Indexer = Index | tuple[Index, ...]

def get(ref: Ref, idx: Indexer) -> Array:
  """Returns `ref[idx]` for NumPy-style indexer `idx`."""

def swap(ref: Ref, idx: Indexer, val: Array) -> Array:
  """Performs `newval, ref[idx] = ref[idx], val` and returns `newval`."""
```

Here, `Indexer` can be any NumPy indexing expression:

```{code-cell}
x_ref = jax.new_ref(jnp.arange(12.).reshape(3, 4))

# int indexing
row = x_ref[0]
x_ref[1] = row

# tuple indexing
val = x_ref[1, 2]
x_ref[2, 3] = val

# slice indexing
col = x_ref[:, 1]
x_ref[0, :3] = col

# advanced int array indexing
vals = x_ref[jnp.array([0, 0, 1]), jnp.array([1, 2, 3])]
x_ref[jnp.array([1, 2, 1]), jnp.array([0, 0, 1])] = vals
```

As with `Array`s, indexing mostly follows NumPy behavior, except for
out-of-bounds indexing which [behaves in the usual way for JAX
`Array`s](https://docs.jax.dev/en/latest/notebooks/Common_Gotchas_in_JAX.html#out-of-bounds-indexing).

### Pure and impure functions

A function that takes a ref as an argument (either explicitly or by lexical
closure) is considered _impure_. For example:

```{code-cell}
# takes ref as an argument => impure
@jax.jit
def impure1(x_ref, y_ref):
  x_ref[...] = y_ref[...]

# closes over ref => impure
y_ref = jax.new_ref(0)

@jax.jit
def impure2(x):
  y_ref[...] = x
```

If a function only uses refs internally, it is still considered _pure_. Purity
is in the eye of the caller. For example:

```{code-cell}
# internal refs => still pure
@jax.jit
def pure1(x):
  ref = jax.new_ref(x)
  ref[...] = ref[...] + ref[...]
  return ref[...]
```

Pure functions, even those that use refs internally, are familiar: for example,
they work with transformations like `jax.grad`, `jax.vmap`, `jax.shard_map`, and
others in the usual way.

Impure functions are sequenced in Python program order.

### Restrictions

`Ref`s are second-class, in the sense that there are restrictions on their
use:

* **Can't return refs** from `jit`\-decorated functions or the bodies of
  higher-order primitives like `jax.lax.scan`, `jax.lax.while_loop`, or
  `jax.lax.cond`
* **Can't pass a ref as an argument more than once** to `jit`\-decorated
  functions or higher-order primitives
* **Can only `freeze` in creation scope**
* **No higher-order refs** (refs-to-refs)

For example, these are errors:

```{code-cell}
x_ref = jax.new_ref(0.)

# can't return refs
@jax.jit
def err1(x_ref):
  x_ref[...] = 5.
  return x_ref  # error!
try:
  err1(x_ref)
except Exception as e:
  print(e)

# can't pass a ref as an argument more than once
@jax.jit
def err2(x_ref, y_ref):
  ...
try:
  err2(x_ref, x_ref)  # error!
except Exception as e:
  print(e)

# can't pass and close over the same ref
@jax.jit
def err3(y_ref):
  y_ref[...] = x_ref[...]
try:
  err3(x_ref)  # error!
except Exception as e:
  print(e)

# can only freeze in creation scope
@jax.jit
def err4(x_ref):
  jax.freeze(x_ref)
try:
  err4(x_ref)  # error!
except Exception as e:
  print(e)
```

These restrictions exist to rule out aliasing, where two refs might refer to the
same mutable memory, making programs harder to reason about and transform.
Weaker restrictions would also suffice, so some of these restrictions may be
lifted as we improve JAX's ability to verify that no aliasing is present.

There are also restrictions stemming from undefined semantics, e.g. in the
presence of parallelism or rematerialization:

* **Can't `vmap` or `shard_map` a function that closes over refs**
* **Can't apply `jax.remat`/`jax.checkpoint` to an impure function**

For example, here are ways you can and can't use `vmap` with impure functions:

```{code-cell}
# vmap over ref args is okay
def dist(x, y, out_ref):
  assert x.ndim == y.ndim == 1
  assert out_ref.ndim == 0
  out_ref[...] = jnp.sum((x - y) ** 2)

vecs = jnp.arange(12.).reshape(3, 4)
out_ref = jax.new_ref(jnp.zeros((3, 3)))
jax.vmap(jax.vmap(dist, (0, None, 0)), (None, 0, 0))(vecs, vecs, out_ref)  # ok!
print(out_ref)
```

```{code-cell}
# vmap with a closed-over ref is not
x_ref = jax.new_ref(0.)

def err5(x):
  x_ref[...] = x

try:
  jax.vmap(err5)(jnp.arange(3.))  # error!
except Exception as e:
  print(e)
```

The latter is an error because it's not clear which value `x_ref` should be
after we run `jax.vmap(err5)`.

### `Ref`s and automatic differentiation

Autodiff can be applied to pure functions as before, even if they use array refs
internally. For example:

```{code-cell}
@jax.jit
def pure2(x):
  ref = jax.new_ref(x)
  ref[...] = ref[...] + ref[...]
  return ref[...]

print(jax.grad(pure2)(3.0))  # 2.0
```

Autodiff can also be applied to functions that take array refs as arguments.
The simplest case is when those ref arguments are only used for plumbing, and
aren't involved in differentiation. For example, `jax.grad` differentiates
with respect to its function's first argument by default, so ref arguments in
other positions are just along for the ride. Only non-differentiated values
can be written into such plumbing refs:

```{code-cell}
# error
def err6(x, some_plumbing_ref):
  y = x + x
  some_plumbing_ref[...] += y
  return y

# fine
def foo(x, some_plumbing_ref):
  y = x + x
  some_plumbing_ref[...] += jax.lax.stop_gradient(y)
  return y
```

Differentiating _with respect to_ a ref-typed argument is another matter:
pointing `jax.grad` at one (e.g. via `argnums`) is an error, since the
gradient for a ref must itself live in a ref. Instead, use `jax.vjp` and
`with_refs`, described below.

You can combine plumbing refs with `custom_vjp` to plumb data out of the
backward pass of a differentiated function:

```{code-cell}
# First, define the helper `stash_grads`:

@jax.custom_vjp
def stash_grads(grads_ref, x):
  return x

def stash_grads_fwd(grads_ref, x):
  return x, grads_ref

def stash_grads_bwd(grads_ref, g):
  grads_ref[...] = g
  return None, g

stash_grads.defvjp(stash_grads_fwd, stash_grads_bwd)
```

```{code-cell}
# Now, use `stash_grads` to stash intermediate gradients:

def f(x, grads_ref):
  x = stash_grads(grads_ref, x)
  x = jnp.sin(x)
  return x

grads_ref = jax.new_ref(0.)
jax.grad(f)(1., grads_ref)
print(grads_ref)  # Ref(0.54), the gradient at the stash point: cos(1.)
```

Notice `stash_grads_fwd` is returning a `Ref` here. That's a special
allowance for `custom_vjp` fwd rules: it's really syntax for indicating which
ref arguments should be shared by both the fwd and bwd rules. So any refs
returned by a fwd rule must be arguments to that fwd rule.

#### Differentiating with respect to `Ref` arguments

The plumbing refs above are just passengers: they carry data out of the
computation, but no gradients flow through them. We can also differentiate
_with respect to_ a ref argument. Since the gradient for a `Ref`-typed input
is itself `Ref`-typed, `jax.grad` doesn't apply here. Instead we use
`jax.vjp`, and bind a gradient ref to the VJP function using its `with_refs`
method:

```{code-cell}
def f(x_ref):
  return x_ref[...] ** 2

x_ref = jax.new_ref(2.)
y, f_vjp = jax.vjp(f, x_ref)

x_grad_ref = jax.new_ref(0.)
f_vjp.with_refs(x_grad_ref)(1.0)  # bind the gradient ref, then apply the VJP
print(x_grad_ref)  # Ref(4.)
```

Here `with_refs` takes one entry per argument of the differentiated function
and returns a new VJP function with those gradient refs bound. When
differentiating with respect to a ref argument, using `with_refs` is
mandatory; the gradient needs a ref to be accumulated into, so calling the
VJP function without binding one is an error:

```{code-cell}
_, f_vjp = jax.vjp(f, jax.new_ref(2.))
try:
  f_vjp(1.0)  # error! no ref bound for the ref-typed argument's gradient
except Exception as e:
  print(e)  # ... gradient must be accumulated into a ref ... `with_refs` ...
```

The gradient is accumulated into the bound ref via `+=`, added to whatever
the ref already contains rather than overwriting it:

```{code-cell}
x_grad_ref = jax.new_ref(100.)
_, f_vjp = jax.vjp(f, jax.new_ref(2.))
f_vjp.with_refs(x_grad_ref)(1.0)
print(x_grad_ref)  # Ref(104.), i.e. 100. + 4.: accumulated, not set
```

Accumulating rather than setting might seem like an odd choice, but it means
one gradient ref can collect contributions from several backward passes with
no extra memory traffic, as in the examples below.

In-place updates to the ref argument inside the differentiated function are
differentiated too. The result is the gradient with respect to the ref's
initial value:

```{code-cell}
def g(x_ref):
  x_ref[...] = jnp.sin(x_ref[...])
  return x_ref[...] ** 2

x_ref = jax.new_ref(2.)
_, g_vjp = jax.vjp(g, x_ref)  # runs g, so x_ref is updated in-place here
g_grad_ref = jax.new_ref(0.)
g_vjp.with_refs(g_grad_ref)(1.0)
print(g_grad_ref)  # Ref(-0.757), i.e. 2*sin(2)*cos(2)
```

#### Gradient refs for value arguments

When differentiating with respect to an ordinary `Array` argument,
`with_refs` is optional: we can call the VJP function directly and get the
gradient back as a value in the usual way, or we can bind a ref and have the
gradient accumulated into it in-place:

```{code-cell}
_, sin_vjp = jax.vjp(jnp.sin, 1.0)
x_bar, = sin_vjp(1.0)  # the usual way: gradient returned as a value
print(x_bar)  # 0.54

grad_ref = jax.new_ref(0.)
_, sin_vjp = jax.vjp(jnp.sin, 1.0)
sin_vjp.with_refs(grad_ref)(1.0)  # gradient accumulated into grad_ref
print(grad_ref)  # Ref(0.54)
```

We can mix and match. Each entry of `with_refs` can be:

* a `Ref`, meaning accumulate this argument's gradient into the ref in-place
  (the VJP call then returns a `jax.ad.GradRef()` placeholder in that
  position);
* `jax.ad.GradValue()`, meaning return this argument's gradient as a value in
  the usual way (the default); or
* `jax.ad.DontWant()`, meaning don't compute this argument's gradient at all
  (the VJP call returns a `jax.ad.DidntWant()` placeholder in that position —
  more on this below).

One reason to use a gradient ref here is to exploit sparsity. Consider
differentiating a function that slices its input:

```{code-cell}
@jax.jit
def take(x, i):
  return x[i]

x = jnp.arange(10.)

_, take_vjp = jax.vjp(take, x, 3)
x_bar, _ = take_vjp(1.0)
print(x_bar)  # [0., 0., 0., 1., 0., 0., 0., 0., 0., 0.]
```

The gradient with respect to `x` is one-hot: the backward pass materializes a
dense array of zeros and sets a single element of it. If we compute many such
gradients and sum them, say over a loop of sparse accesses, we pay for a
dense array's worth of memory traffic on each one, even though each
contribution only touches one element.

If we instead bind a gradient ref with `with_refs`, each backward pass
performs a sparse in-place add-update, writing only where it needs to:

```{code-cell}
grad_ref = jax.new_ref(jnp.zeros(10))

for i in [3, 5, 3]:
  _, take_vjp = jax.vjp(take, x, i)
  take_vjp.with_refs(grad_ref, jax.ad.GradValue())(1.0)  # no gradient ref for i

print(grad_ref)  # Ref([0., 0., 0., 2., 0., 1., 0., 0., 0., 0.])
```

We can check that the update is sparse by inspecting the jaxpr of a single
VJP application:

```{code-cell}
@jax.make_jaxpr
def take_vjp_jaxpr():
  _, take_vjp = jax.vjp(take, x, 3)
  take_vjp.with_refs(grad_ref, jax.ad.GradValue())(1.0)

print(take_vjp_jaxpr())
```

The backward pass boils down to `b[c:c+1] += j`, an in-place add-update of
one element of the gradient ref, with no dense one-hot array in sight.

#### `DontWant`: skipping unneeded gradients

The VJP function returned by `jax.vjp` computes gradients for all the
arguments of the differentiated function. But sometimes we don't need all of
them, like when differentiating with respect to parameters but not data.
Passing `jax.ad.DontWant()` for an argument tells the backward pass not to
compute its gradient at all, playing the same role for `jax.vjp` that
`argnums` plays for `jax.grad`:

```{code-cell}
def predict(W, x):
  return W @ x

W = jnp.ones((4, 4))
x = jnp.ones(4)

_, f_vjp = jax.vjp(predict, W, x)
W_bar, x_bar = f_vjp.with_refs(jax.ad.GradValue(), jax.ad.DontWant())(jnp.ones(4))
print(W_bar[0])  # [1., 1., 1., 1.]
print(x_bar)  # DidntWant()
```

This is more than a convenience: it can save real work in the backward pass,
like an eager form of dead code elimination. Transpose rules can check for
`DontWant` and skip computing the corresponding cotangents. For example, the
transpose of matrix multiplication usually computes two dot products, one for
each operand's gradient, but with `DontWant` it computes only one. We can see
that by counting the `dot_general` operations in the jaxpr of each VJP
application:

```{code-cell}
_, f_vjp = jax.vjp(predict, W, x)
both = jax.make_jaxpr(lambda: f_vjp(jnp.ones(4)))()
only_W = jax.make_jaxpr(
    lambda: f_vjp.with_refs(jax.ad.GradValue(), jax.ad.DontWant())(jnp.ones(4)))()

print(str(both).count('dot_general'))    # 2, one dot for each gradient
print(str(only_W).count('dot_general'))  # 1, the dot for x_bar is skipped
```

#### Example: gradient accumulation over microbatches

Here's a more realistic recipe that puts these pieces together. In pipelined
training, we often split a batch into microbatches, run a forward and
backward pass one microbatch at a time, and accumulate weight gradients as we
go. Using `with_refs`, each microbatch's backward pass accumulates directly
into a single fixed gradient buffer:

```{code-cell}
NUM_LAYERS = 3
NUM_MUBATCHES = 5
MUBATCH_SIZE = 7

def mubatch_loss(Ws, xs):
  # inner loop over layers
  act, _ = jax.lax.scan(lambda x, W: (jnp.dot(x, W), None), xs, Ws)
  return jnp.mean(act)

def process_batch(Ws, xs_batch):
  grad_acc = jax.new_ref(jnp.zeros_like(Ws))

  def process_mubatch(_, xs):
    loss, f_vjp = jax.vjp(lambda Ws: mubatch_loss(Ws, xs), Ws)
    f_vjp.with_refs(grad_acc)(jnp.ones_like(loss))  # accumulate in-place
    return (), loss

  xs_mubatches = xs_batch.reshape(NUM_MUBATCHES, MUBATCH_SIZE, -1)
  # outer loop over microbatches
  (), losses = jax.lax.scan(process_mubatch, (), xs_mubatches)
  return jax.freeze(grad_acc), losses

Ws = jnp.ones((NUM_LAYERS, 4, 4))
xs_batch = jnp.ones((NUM_MUBATCHES * MUBATCH_SIZE, 4))
grads, losses = process_batch(Ws, xs_batch)
```

Each iteration of the outer scan runs a forward and backward pass for one
microbatch, and `with_refs(grad_acc)` makes the backward pass add that
microbatch's gradient contribution directly into `grad_acc`. Note that the
scan body closes over `grad_acc`, which is fine for `scan` (though it
wouldn't be for `vmap` or `shard_map`, as discussed above). Once all the
microbatches are processed, we `freeze` the accumulator to get the total
batch gradient as an immutable `Array`.

The result matches what we'd get by differentiating the summed loss directly:

```{code-cell}
xs_mubatches = xs_batch.reshape(NUM_MUBATCHES, MUBATCH_SIZE, -1)
grads_expected = jax.grad(
    lambda Ws: sum(mubatch_loss(Ws, xs) for xs in xs_mubatches))(Ws)
print(jnp.allclose(grads, grads_expected, atol=1e-3, rtol=1e-3))  # True
```

But unlike that version, the ref-based version never materializes
per-microbatch gradients as separate arrays: there's one gradient buffer,
allocated once, no matter how many microbatches we process.

### `Ref`s and performance

At the top level, when calling `jit`\-decorated functions, `Ref`s obviate
the need for donation, since they are effectively always donated:

```{code-cell}
@jax.jit
def sin_inplace(x_ref):
  x_ref[...] = jnp.sin(x_ref[...])

x_ref = jax.new_ref(jnp.arange(3.))
print(x_ref.unsafe_buffer_pointer(), x_ref)
sin_inplace(x_ref)
print(x_ref.unsafe_buffer_pointer(), x_ref)
```

Here `sin_inplace` operates in-place, updating the buffer backing `x_ref` so
that its address stays the same.

Under a `jit`, you should expect array references to point to fixed buffer
addresses, and for indexed updates to be performed in-place.

**Temporary caveat:** dispatch from Python to impure `jit`\-compiled functions
that take `Ref` inputs is currently slower than dispatch to pure
`jit`\-compiled functions, since it takes a less optimized path.

### `foreach`, a new way to write `scan`

As you may know, `jax.lax.scan` is a loop construct with a built-in fixed access
pattern for scanned-over inputs and outputs. The access pattern is built in for
autodiff reasons: if we were instead to slice into immutable inputs directly,
reverse-mode autodiff would end up creating one-hot gradients and summing them
up, which can be asymptotically inefficient. See [Sec 5.3.3 of the Dex
paper](https://arxiv.org/pdf/2104.05372).

But reading slices of `Ref`s doesn't have this efficiency problem: when we
apply reverse-mode autodiff, we always generate in-place accumulation
operations. As a result, we no longer need to be constrained by `scan`'s fixed
access pattern. We can write more flexible loops, e.g. with non-sequential
access.

Moreover, having mutation available allows for some syntax tricks, like in this
recipe for a `foreach` decorator:

```{code-cell}
import jax
import jax.numpy as jnp
from jax.lax import scan

def foreach(*args):
  def decorator(body):
    return scan(lambda _, elts: (None, body(*elts)), None, args)[1]
  return decorator
```

```{code-cell}
r = jax.new_ref(0)
xs = jnp.arange(10)

@foreach(xs)
def ys(x):
  r[...] += x
  return x * 2

print(r)   # Ref(45, dtype=int32)
print(ys)  # [ 0  2  4  6  8 10 12 14 16 18]
```

Here, the loop runs immediately, updating `r` in-place and binding `ys` to be
the mapped result.
