Just another casual day of refactoring MaxText PR#2541 (#1768)

AI-Hypercomputer/maxtext#2541

Author: gpupuck

.github/container/Dockerfile.maxtext

@@ -18,19 +18,17 @@ EOF
 
 RUN <<"EOF" bash -ex -o pipefail
 echo "-e file://${SRC_PATH_MAXTEXT}" >> /opt/pip-tools.d/requirements-maxtext.in
-echo "-r ${SRC_PATH_MAXTEXT}/base_requirements/requirements.txt" >> /opt/pip-tools.d/requirements-maxtext.in
+echo "-r ${SRC_PATH_MAXTEXT}/dependencies/requirements/base_requirements/requirements.txt" >> /opt/pip-tools.d/requirements-maxtext.in
 EOF
 
 # add version constraints to avoid eternal dependency resolution
 RUN <<"EOF" bash -ex -o pipefail
 for pattern in \
     "s|^tensorflow$|tensorflow==2.18.1|g" \
     "s|^tensorflow-text$|tensorflow-text==2.18.1|g" \
-    "s|^jax!=.*|jax|g" \
-    "s|^jaxlib!=.*|jaxlib|g" \
   ; do
     # tensorflow-cpu,tensorboard,tensorflow-text>=2.19.0 is incompatible with tensorflow==2.18.1
-    sed -i "${pattern}" ${SRC_PATH_MAXTEXT}/base_requirements/requirements.txt
+    sed -i "${pattern}" ${SRC_PATH_MAXTEXT}/dependencies/requirements/base_requirements/requirements.txt
 done
 EOF