What the process is still holding

Two arrays, one of them derived from the other, and a question about what that costs. jax.live_arrays() answers it directly: it returns every array currently alive on the devices, so the whole ledger is a sum of nbytes over that list. One 2048 by 2048 float32 array is 16 MiB, and after y = x * 2 the ledger reads two arrays and 32 MiB.

del y takes it straight back to one array and 16 MiB. The device buffer is released when the last Python reference to the array goes, which is what makes a ledger printed at the top of a training loop a fair account of what that loop is holding at that moment.

One thing the list does not include is the memory a call needs while it is running. The temporary arena the compiled program reserves is not a jax.Array and never appears here, and lesson three reads that number off the compiled program instead. Two instruments, two questions, and a figure from one is not an answer to the other.

import jax
import jax.numpy as jnp

def live():
    arrays = jax.live_arrays()
    return len(arrays), sum(a.nbytes for a in arrays) // 1024 ** 2

x = jnp.ones((2048, 2048)).block_until_ready()
print("x alive          ", live())
y = (x * 2.0).block_until_ready()
print("y = x * 2 alive  ", live())
del y
print("after del y      ", live())

# x alive           (1, 16)
# y = x * 2 alive   (2, 32)
# after del y       (1, 16)

Close a jitted function over an array and that array gets a second holder your code never names. Add a 16 MiB constant to a 16 MiB input, call the function once, and the ledger reads three arrays and 48 MiB while your program has two names.

Now del c, and nothing moves. The count stays at three and the total stays at 48 MiB, because the jit cache is holding both the original array and a copy of it, and neither is reachable from a name in your program. jax.clear_caches() is what returns the ledger to a single array and 16 MiB.

Pass the same array as an argument instead of capturing it and the ledger matches what the code says: two arrays after the call, one once the name is gone. So when a long-running process holds more than the arrays you can point at, cached programs and the constants they captured are the first place to look.

import gc

import jax
import jax.numpy as jnp

def live():
    arrays = jax.live_arrays()
    return len(arrays), sum(a.nbytes for a in arrays) // 1024 ** 2

x = jnp.ones((2048, 2048)).block_until_ready()

c = jnp.full((2048, 2048), 7.0).block_until_ready()
captured = jax.jit(lambda p: p + c)
captured(x).block_until_ready()
gc.collect()
print("captured, after one call", live())
del c
gc.collect()
print("captured, after del c   ", live())
jax.clear_caches()
gc.collect()
print("captured, after clear   ", live())

k = jnp.full((2048, 2048), 7.0).block_until_ready()
passed = jax.jit(lambda p, q: p + q)
passed(x, k).block_until_ready()
gc.collect()
print("argument, after one call", live())
del k
gc.collect()
print("argument, after del k   ", live())

# captured, after one call (3, 48)
# captured, after del c    (3, 48)
# captured, after clear    (1, 16)
# argument, after one call (2, 32)
# argument, after del k    (1, 16)

One keyword changes that arithmetic, and its own lesson is elsewhere. donate_argnums is told at /jax/jit/giving-an-argument-away: what the promise says, how to prove before anything runs that the compiler accepted it, what happens to the array you gave away, and the two mismatches that get a donation refused. What belongs in a performance chapter is the measurement.

Run the same update twice, plain and donating, counting either side of the call. Without the keyword one 16 MiB array goes in and two come out, because your name still holds the old parameters and the new ones are a fresh allocation. With it, one goes in and one comes out, so the peak stays at 16 MiB rather than doubling.

Then time the two loops and nothing separates them. Fifty rotations at this size measured 3508 microseconds per call plain against 3447 donating in one run, and across five runs the plain figure landed between 3184 and 4021 while the donated one landed between 3185 and 4561. The ranges overlap, so the honest report is that this machine cannot tell them apart. What donation saves is bytes; turning bytes into time takes a device where the second allocation and the larger footprint cost something.

import jax
import jax.numpy as jnp

def live():
    arrays = jax.live_arrays()
    return len(arrays), sum(a.nbytes for a in arrays) // 1024 ** 2

for donate in (False, True):
    kwargs = {"donate_argnums": 0} if donate else {}
    update = jax.jit(lambda p: p * 0.99, **kwargs)
    p = jnp.ones((2048, 2048))
    p.block_until_ready()
    before = live()
    new = update(p).block_until_ready()
    print(f"donate={donate!s:5s} before {before}  after {live()}")
    del p, new

# donate=False before (1, 16)  after (2, 32)
# donate=True  before (1, 16)  after (1, 16)

Put a donating function inside the harness from lesson one and the second iteration fails. The first call consumes the input, and every call after it hands the runtime a buffer that no longer exists, so the rejection arrives at the argument, before any of your program runs.

The message names the buffer and not the line, and it comes back as a ValueError from the call itself: INVALID_ARGUMENT: Invalid buffer passed: buffer has been deleted or donated. A harness that catches broadly will log that as a bad measurement when it is really a statement about the loop's shape.

The fix is not a flag. Feed the loop the value the previous call produced, p = update(p), and the promise is true every time, because the new parameters really did replace the old ones. A training loop has that shape already, which is why donation fits there and not in a microbenchmark that calls f(x) twenty times with the same x.

Donation and a repeat-the-same-argument benchmark are incompatible by construction, not by accident.

import jax
import jax.numpy as jnp

update = jax.jit(lambda p: p * 0.99, donate_argnums=0)

p = jnp.ones((4,))
for i in range(3):
    try:
        update(p).block_until_ready()
        print("call", i, "ok")
    except ValueError as err:
        print("call", i, err)

p = jnp.ones((4,))
for _ in range(3):
    p = update(p)
print(p.block_until_ready())

# call 0 ok
# call 1 INVALID_ARGUMENT: Invalid buffer passed: buffer has been deleted or donated.
# call 2 INVALID_ARGUMENT: Invalid buffer passed: buffer has been deleted or donated.
# [0.97029907 0.97029907 0.97029907 0.97029907]