jax

dbc860fe - Cherry-pick fixes for JAX 0.9.2 release from JAX 0.9.1 (#740)

Cherry-pick fixes for JAX 0.9.2 release from JAX 0.9.1 (#740) Co-authored-by: Arati Ganesh <arati.ganesh@amd.com> Co-authored-by: shurale-nkn <shurale.nkn@gmail.com>