jax

82ed029f - relax version requirements for ROCm Jax Plugin wheels, allow overriding using env var