6 / 8 JAX & XLA Mental Model

Donate Refresh

Calling jax.jit(f, donate_argnums=(0,)) tells the compiler it is allowed to do what with the buffer backing argument 0?

A

Key the compilation cache on that argument's shape only, ignoring its dtype for retracing decisions

B

Treat the argument as static, forcing a full retrace whenever its concrete value changes

C

Alias that input buffer's device memory for reuse as an output buffer, avoiding an extra allocation, but the caller's original array becomes invalid (deleted) after the call

D

Exclude that argument from the autodiff graph, treating it as a compile-time constant

E

Asynchronously spill that argument to host RAM while compilation runs, to reduce peak device memory