added requirements

Author: prishajain1

requirements_with_jax_ai_image.txt

@@ -1,7 +1,7 @@
 # Requirements for Building the MaxDifussion Docker Image
 # These requirements are additional to the dependencies present in the JAX AI base image.
 --extra-index-url https://download.pytorch.org/whl/cpu
-jax>=0.7.2
+jax[k8s]>=0.7.2
 jaxlib>=0.4.30
 grain
 google-cloud-storage>=2.17.0