jax

93170d9c - Add JAX version to TPU_ML_PLATFORM_VERSION environment variable.

Add JAX version to TPU_ML_PLATFORM_VERSION environment variable. This will allow us to track the JAX version that is being used on Cloud TPUs PiperOrigin-RevId: 637025132