jax

port rocm/jax 0.9.2 cherrypicks to amd-main branch (#740)

#741

Open