jax

c7a70f49 - [ROCm] Cherry-pick: Fix inter-block write race in pallas test_non_range_while_loop (#764)