Removing tunix condition in JAII build

Author: parambole

maxtext_jax_ai_image.Dockerfile

@@ -47,6 +47,11 @@ RUN if [ "$DEVICE" = "tpu" ] && [ "$JAX_STABLE_STACK_BASEIMAGE" = "us-docker.pkg
         python3 -m pip install -r /deps/requirements_with_jax_ai_image.txt; \
   fi
 
+# Install google-tunix for TPU devices, skip for GPU
+RUN if [ "$DEVICE" = "tpu" ]; then \
+        python3 -m pip install 'google-tunix>=0.1.2'; \
+  fi
+
 # Now copy the remaining code (source files that may change frequently)
 COPY . .