jax

[ROCm] Cherry-pick: Fix inter-block write race in pallas test_non_range_while_loop

#764

Merged