Skip to content

← JAX Performance & Distributed step 3 of 25

Medium Primitives

jit with donate_argnums

Jump to the task

Let an update reuse an old buffer

JAX arrays are immutable from Python's point of view: an update returns a new value rather than changing the array you supplied. If you no longer need an input, buffer donation lets the compiler reuse its storage for a compatible output.

donate_argnums=(0,) offers the first argument's buffer for reuse. This can reduce peak memory use, but it does not guarantee zero allocations. The compiler decides whether a donated buffer is usable, and intermediate results or other outputs may still need storage.

Worked example

import jax
import jax.numpy as jnp

update = jax.jit(lambda p, g: p - g, donate_argnums=(0,))
params = jnp.ones(4, dtype=jnp.float32)
grads = jnp.full(4, 0.01, dtype=jnp.float32)
params = update(params, grads)
# params now refers to the result, approximately [0.99, 0.99, 0.99, 0.99]

Rebinding params is intentional: later steps use the returned array. Treat the old input as consumed after the call. When JAX reuses its buffer, accessing the old array raises a deleted-buffer error. Another Python name for the same array does not preserve a usable copy of its data.

Problem

Implement update_with_donate(state, delta) to return state + delta. Compile the addition with jax.jit(..., donate_argnums=(0,)), offering state for reuse while retaining delta.

  • Inputs are one-dimensional float32 JAX arrays with the same shape.
  • The two arguments use distinct buffers; the caller no longer needs the old state after the call.
  • Return the sum with the same shape. Empty input returns an empty array.
  • Construct the jitted helper once at module scope, then call it from update_with_donate so repeated updates can reuse the compiled function.

The tests check the returned values. They do not prove a particular allocation count, memory saving, or backend buffer assignment.

Respect the lifetime of the old value

Do not inspect or reuse an input after offering it for donation. This applies to all aliases of that array, including an array stored in a list or dictionary. Python references do not become unsafe raw pointers; they refer to an array whose storage may have been invalidated.

If you need the original values later, keep an independent copy before the call or do not donate. Similarly, do not pass the same buffer as both the donated state and an argument you intend to retain. Repeated calls should use the previous result as the next state.

Donation is permission, not a guarantee

A donated buffer must be compatible with an output for reuse, including its shape and dtype. If no output can use it, JAX may warn that a donated buffer was not usable. Do not base correctness on donation happening or on the old input remaining accessible when it does not happen.

With multiple outputs, do not assume the first output receives a particular donated buffer. Matching and allocation are compiler decisions. For this exercise there is one array output matching the state shape and dtype.

The diagram below illustrates compilation reuse for a stable function with lengths 32 and 64. It does not visualize buffer ownership or prove memory reuse. Its x32/x64 labels describe lengths, not floating-point dtypes.

Loading visualization…