From 56c293cc6221d499306539cb74957a38772c67ea Mon Sep 17 00:00:00 2001
From: Masahiro Hiramori <mhg00g13@gmail.com>
Date: Tue, 22 Oct 2024 21:53:54 +0900
Subject: [PATCH] Revert part of  "[CI] Upgrade CI (#17425)" change the
 versions of jax, tensorflow, tflite back to what we've been using before

---
 docker/install/ubuntu_install_jax.sh          | 18 ++++-----
 docker/install/ubuntu_install_tensorflow.sh   |  4 +-
 .../ubuntu_install_tensorflow_aarch64.sh      |  4 +-
 docker/install/ubuntu_install_tflite.sh       | 40 +++++++++----------
 4 files changed, 32 insertions(+), 34 deletions(-)

diff --git a/docker/install/ubuntu_install_jax.sh b/docker/install/ubuntu_install_jax.sh
index 17114e0efce8..19149909161e 100644
--- a/docker/install/ubuntu_install_jax.sh
+++ b/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
diff --git a/docker/install/ubuntu_install_tensorflow.sh b/docker/install/ubuntu_install_tensorflow.sh
index 012b678916b3..2225b7aef3b8 100755
--- a/docker/install/ubuntu_install_tensorflow.sh
+++ b/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
diff --git a/docker/install/ubuntu_install_tensorflow_aarch64.sh b/docker/install/ubuntu_install_tensorflow_aarch64.sh
index 4b158948387b..fcd912a4478a 100755
--- a/docker/install/ubuntu_install_tensorflow_aarch64.sh
+++ b/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
diff --git a/docker/install/ubuntu_install_tflite.sh b/docker/install/ubuntu_install_tflite.sh
index 8faabc022640..36e6dfc42794 100755
--- a/docker/install/ubuntu_install_tflite.sh
+++ b/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