Ref: mutable arrays for data plumbing and memory control — JAX documentation

Skip to main content

Back to top

Ctrl+K

Search<br>Ctrl+K

.ipynb

.pdf

Ref: mutable arrays for data plumbing and memory control

Contents

Ref: mutable arrays for data plumbing and memory control#

JAX Arrays are immutable, representing mathematical values. Immutability can<br>make code easier to reason about, and is useful for optimized compilation,<br>parallelization, rematerialization, and transformations like autodiff.

But immutability is constraining too:

expressiveness — plumbing out intermediate data or maintaining state,<br>e.g. for normalization statistics or metrics, can feel heavyweight;

performance — it’s more difficult to reason about performance, like<br>memory lifetimes and in-place updates.

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

import jax<br>import jax.numpy as jnp

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

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

print(x_ref) # Ref([0., 0., 0.])<br>f()<br>f()<br>print(x_ref) # Ref([0., 2., 0.])

Ref([0., 0., 0.], dtype=float32)<br>Ref([0., 2., 0.], dtype=float32)

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

def g(x):<br>x_ref = jax.new_ref(0.)<br>x_ref[...] = jnp.sin(x)<br>return x_ref[...]

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

0.5403023

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

x_ref = jax.new_ref(1.0)<br>try:<br>jnp.sin(x_ref) # error! can't do math on refs<br>except Exception as e:<br>print(e)

sin requires ndarray or scalar arguments, got at position 0.

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<br>recipes.

API#

If you’ve ever used<br>Pallas, then Ref<br>should look familiar. A big difference is that you can create new Refs<br>yourself directly using jax.new_ref:

from jax import Array, Ref

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

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

def freeze(ref: Ref) -> Array:<br>"""Invalidate given reference and produce its final value."""

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

import types<br>Index = int | slice | Array | types.EllipsisType<br>Indexer = Index | tuple[Index, ...]

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

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

Here, Indexer can be any NumPy indexing expression:

x_ref = jax.new_ref(jnp.arange(12.).reshape(3, 4))

# int indexing<br>row = x_ref[0]<br>x_ref[1] = row

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

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

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

As with Arrays, indexing mostly follows NumPy behavior, except for<br>out-of-bounds indexing which behaves in the usual way for JAX<br>Arrays.

Pure and impure functions#

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

# takes ref as an argument => impure<br>@jax.jit<br>def impure1(x_ref, y_ref):<br>x_ref[...] = y_ref[...]

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

@jax.jit<br>def impure2(x):<br>y_ref[...] = x

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

# internal refs => still pure<br>@jax.jit<br>def pure1(x):<br>ref = jax.new_ref(x)<br>ref[...] = ref[...] + ref[...]<br>return ref[...]

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

Impure functions are sequenced in Python program order.

Restrictions#

Refs are second-class, in the sense that there are restrictions on their<br>use:

Can’t return refs from jit-decorated functions or the bodies of<br>higher-order primitives like jax.lax.scan, jax.lax.while_loop, or<br>jax.lax.cond

Can’t pass a ref as an argument more than once to jit-decorated<br>functions or higher-order primitives

Can only freeze in creation scope

No...