Document disabled donation under `debug_nans`

[Gattocrucco]

Description

I had a buffer donation not working and it took me a while to find out that it's disabled in debug nans mode. It also took me a while to find out this is intentional. Could you document this in the config reference? It's missing in all of

• https://docs.jax.dev/en/latest/_autosummary/jax.jit.html#jax.jit
• https://docs.jax.dev/en/latest/buffer_donation.html#buffer-donation
• https://docs.jax.dev/en/latest/config_options.html#jax_debug_nans

The first issue on debug_nans-donation conflict is #12514, pr #13144 disabled donation in debug_nans mode altogether.

Minimal reproduction:

from functools import partial

from jax import debug_nans, jit
from jax import numpy as jnp


# check jitted function does not donate, ok
@jit
def f(x):
    return x


x = jnp.arange(10)
xp = x.unsafe_buffer_pointer()
fx = f(x)
assert xp != fx.unsafe_buffer_pointer()


# check donation works, ok
@partial(jit, donate_argnums=(0,))
def fd(x):
    return x


x = jnp.arange(10)
xp = x.unsafe_buffer_pointer()
fx = f(x)
fdx = fd(x)
assert xp == fdx.unsafe_buffer_pointer()


# donation disabled if debug_nans enabled, surprising but intentional
with debug_nans(True):

    @partial(jit, donate_argnums=(0,))
    def fdn(x):
        return x

    x = jnp.arange(10)
    xp = x.unsafe_buffer_pointer()
    fdnx = fdn(x)
    assert xp != fdnx.unsafe_buffer_pointer()

System info (python version, jaxlib version, accelerator, etc.)

jax:    0.8.0
jaxlib: 0.8.0
numpy:  2.3.4
python: 3.14.0 (main, Oct 14 2025, 21:10:22) [Clang 20.1.4 ]
device info: cpu-1, 1 local devices"
process_count: 1
platform: uname_result(system='Darwin', node='MacBook-Pro-2.local', release='25.1.0', version='Darwin Kernel Version 25.1.0: Mon Oct 20 19:32:41 PDT 2025; root:xnu-12377.41.6~2/RELEASE_ARM64_T6000', machine='arm64')

[Gattocrucco]

Note: debug_nans(True) disables donation at the point where a function is compiled, example that pins this down:

"""Check that debug_nans disables donation at the point where a function is compiled."""

from functools import partial

from jax import debug_nans, jit
from jax import numpy as jnp


@partial(jit, donate_argnums=(0,))
def f(x):
    return x


x = jnp.arange(10)
xp = x.unsafe_buffer_pointer()

with debug_nans(True):
    fx = f(x)
    assert xp != fx.unsafe_buffer_pointer()  # not donated, ok

fx2 = f(x)
assert xp != fx2.unsafe_buffer_pointer()  # still not donated, a bit unexpected

f.clear_cache()
fx3 = f(x)
assert xp == fx3.unsafe_buffer_pointer()  # donated because re-compiled