Clarification Behavior of Donated Arguments

[cdelv]

Hi,

I'm writing because I have some ideas about using donate arguments to speed up computations, but I'm not sure how this works. So, I would appreciate any ideas and clarification.

The following is my understanding.

Let's say I have a Jax array x. If I want to mutate x, the Jax way would be to apply a transformation and then overwrite x with the new values:

import jax.numpy as jnp

def func(x):
    return 2*x + 2.

x = jnp.ones(10)

x = func(x)

This is where it gets tricky. What happens to the memory? Let's say we are on a CPU, so there is no need to consider where the data is. My impression from Jax's documentation is that we create a copy of x, and then replace the pointer of x so it now points to the new memory with the updated values. The async execution hides the latency of creating new memory by continuing with the computations. We can confirm this by printing the id() of x before and after the function is called:

714305696
714458352

This happened even inside JITted functions:

@jax.jit
def main():
    N = 10
    x = jnp.ones(N)
    print(id(x))
    x = func(x)
    print(id(x))

wich yields

139688742335696
139688742337040

I know the array type holds the pointer to the actual data buffer allocated on the device. One question is whether this buffer is reused. So even if we get a new object, it's just a tiny struct (in case the buffer is reused). Maybe JIT optimizes this?

Now, let's get into the donated arguments. The idea is that if we donate a buffer, we can't use it again. I like to think of this as the hoe rust arguments are consumed when they get inside a function:

@partial(jax.jit, donate_argnames=("y"))
def func(x, y):
    return 2 * x + 2.0 * y


N = 10
x = jnp.ones(N)
y = jnp.ones(N)
x = func(x, y)
print(x.is_deleted()) # False
print(y.is_deleted()) # True

However, what if we donate x? Well, that's weird, and this is why I would like some explanations. If I do this:

@partial(jax.jit, donate_argnames=("x", "y"))
def func(x, y):
    return 2 * x + 2.0


N = 10
x = jnp.ones(N)
y = jnp.ones(N)
x = func(x, y)
print(x.is_deleted()) # False
print(y.is_deleted()) # False

But if I do this

@partial(jax.jit, donate_argnames=("x", "y"))
def func(x, y):
    return 2 * x + 2.0 * y


N = 10
x = jnp.ones(N)
y = jnp.ones(N)
x = func(x, y)
print(x.is_deleted()) # False
print(y.is_deleted()) # False

and I get this warning

UserWarning: Some donated buffers were not usable: float32[10].

See an explanation at https://docs.jax.dev/en/latest/faq.html#buffer-donation. warnings.warn("Some donated buffers were not usable:"

Why does the first example work, but this one does not? Now, if func only takes x and I donate x, it does not show as deleted, but I also don't get the warning. This gives me the idea that every function with this pattern should have the arguments donated: x = func(x). But at this point, I have no idea if the donating argument in this case does anything, and it's weird how some cases work and some don't.

---

Now, I wanted to see what was going on with the data buffer pointer and found something very weird

import jax
jax.config.update("jax_platform_name", "cpu")
import jax.numpy as jnp
from functools import partial
import numpy as np

@partial(jax.jit, donate_argnames=())
def func(x, y):
    return 2 * x + 2.0 * y
.
.
.
ptr = np.asarray(x).ctypes.data
print(ptr) # 226193472

x = func(x, y)

ptr = np.asarray(x).ctypes.data
print(ptr) # 228708032

@partial(jax.jit, donate_argnames=("x", "y"))

238847680
/home/wind/PythonEnvs/python3/lib/python3.13/site-packages/jax/_src/interpreters/mlir.py:1316: UserWarning: Some donated buffers were not usable: float32[10].
See an explanation at https://docs.jax.dev/en/latest/faq.html#buffer-donation.
  warnings.warn("Some donated buffers were not usable:"
238847680

@partial(jax.jit, donate_argnames=("y"))

214935872
216714048

@partial(jax.jit, donate_argnames=("x")):

105575808
105575808

Note that donating x allows the buffer to be reused, even with the warning. However, it's weird that we get a warning when donating both arguments but not x.

I would appreciate it if anyone has insights into how this works and could clarify best practices and how this affects performance.