[CI] Revert jax, keras, tensorflow, and tflite upgrades introduced #17425

docker/install/ubuntu_install_jax.sh

@@ -20,18 +20,16 @@ set -e
 set -u
 set -o pipefail
 
-JAX_VERSION=0.4.30
-
-# Install jaxlib
+# Install jax and jaxlib
 if [ "$1" == "cuda" ]; then
-    pip install -U \
-        "jax[cuda12]~=${JAX_VERSION}" \
-        jaxlib~=${JAX_VERSION}
+    pip3 install --upgrade \
+        jaxlib~=0.4.9 \
+        "jax[cuda11_pip]~=0.4.9" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
 else
-    pip3 install -U \
-        jax~=${JAX_VERSION} \
-        jaxlib~=${JAX_VERSION}
+    pip3 install --upgrade \
+        jaxlib~=0.4.9 \
+        "jax[cpu]~=0.4.9"
 fi
 
 # Install flax
-pip3 install flax~=0.8.5
+pip3 install flax~=0.6.9

docker/install/ubuntu_install_tensorflow.sh

@@ -21,5 +21,5 @@ set -u
 set -o pipefail
 
 pip3 install \
-    keras==3.5 \
-    tensorflow==2.17.0
+    keras==2.9 \
+    tensorflow==2.9.1

docker/install/ubuntu_install_tensorflow_aarch64.sh

@@ -25,5 +25,5 @@ apt-install-and-clear -y --no-install-recommends libhdf5-dev
 # h5py wheel tries to use the wrong .so file
 pip3 install \
     numpy==1.23.5 \
-    keras==3.5 \
-    tensorflow-aarch64~=2.16.1
+    keras==2.9 \
+    tensorflow-aarch64~=2.9.3

docker/install/ubuntu_install_tflite.sh

@@ -26,11 +26,11 @@ set -o pipefail
 TENSORFLOW_VERSION=$(python3 -c "import tensorflow; print(tensorflow.__version__)" 2> /dev/null)
 
 # Download, build and install flatbuffers
-git clone --branch=v24.3.25 --depth=1 --recursive https://github.com/google/flatbuffers.git
-pushd flatbuffers
-  cmake -G Ninja -DCMAKE_BUILD_TYPE=Release -DCMAKE_CXX_FLAGS="-Wno-class-memaccess"
-  ninja install -j8
-popd
+git clone --branch=v1.12.0 --depth=1 --recursive https://github.com/google/flatbuffers.git
+cd flatbuffers
+cmake -G "Unix Makefiles" -DCMAKE_BUILD_TYPE=Release -DCMAKE_CXX_FLAGS="-Wno-class-memaccess"
+make install -j8
+cd ..
 
 # Install flatbuffers python packages.
 pip3 install flatbuffers
@@ -41,22 +41,22 @@ pip3 install flatbuffers
 git clone https://github.com/tensorflow/tensorflow --branch=v${TENSORFLOW_VERSION} --depth 1
 
 mkdir -p /opt/tflite
-pushd /opt/tflite
-  cmake -G Ninja \
-    -DTFLITE_ENABLE_XNNPACK=OFF \
-    /tensorflow/tensorflow/lite
+cd /opt/tflite
+cmake \
+  -DTFLITE_ENABLE_XNNPACK=OFF \
+  /tensorflow/tensorflow/lite
+
+cmake --build .
+cd -
 
-  cmake --build .
-popd
 
 # Setup tflite from schema
 mkdir tflite
-find / -name "schema.fbs"
-cp /tensorflow/tensorflow/lite/stablehlo/schema/schema.fbs tflite
-pushd tflite
-  flatc --python schema.fbs
+cp tensorflow/tensorflow/lite/schema/schema.fbs tflite
+cd tflite
+flatc --python schema.fbs
 
-  cat <<EOM >setup.py
+cat <<EOM >setup.py
 import setuptools
 
 setuptools.setup(
@@ -77,12 +77,12 @@ setuptools.setup(
 )
 EOM
 
-  cat <<EOM >__init__.py
+cat <<EOM >__init__.py
 name = "tflite"
 EOM
 
-  # Install tflite over python3
-  python3 setup.py install
+# Install tflite over python3
+python3 setup.py install
 
-popd
+cd ..
 rm -rf tflite