jax

a5634cb9 - [Pallas] In GPU Interpret Mode, replace ordered IO callbacks with unordered

[Pallas] In GPU Interpret Mode, replace ordered IO callbacks with unordered This change threads an `int32` "token" through all `callback.io_callback` calls in GPU Interpret Mode to enforce their ordering, so that we can use unordered IO callbacks instead of ordered IO callbacks. A single ordered IO callback is made at the start and at the end of interpreting a kernel, to prevent a jitted computation from overlapping the interpretation of two different GPU kernels (which is not currently supported). Using unordered callbacks enables an upcoming improvement to error handling in GPU Interpret Mode. PiperOrigin-RevId: 927799871