jax

a94e1bb9 - [ROCm] Manually bump of XLA pin.

[ROCm] Manually bump of XLA pin. This adds on top of upstream's XLA pin: - 8028c6c78ea42b7184be31bcbf9d1fe40809bde1 - d052572562cfe937f447a565712ebb6c59620945