update python packages, onnx and pytorch

Author: mshr-h

docker/install/ubuntu_install_onnx.sh

@@ -20,70 +20,25 @@ set -e
 set -u
 set -o pipefail
 
-# We need to fix the onnx version because changing versions tends to break tests
-# TODO(mbrookhart): periodically update
-
-# onnx 1.9 removed onnx optimizer from the main repo (see
-# https://github.com/onnx/onnx/pull/2834).  When updating the CI image
-# to onnx>=1.9, onnxoptimizer should also be installed.
-
 # Get the Python version
 PYTHON_VERSION=$(python3 -c "import sys; print(f'{sys.version_info.major}.{sys.version_info.minor}')")
 
 # Set default value for first argument
 DEVICE=${1:-cpu}
 
 # Install the onnx package
-pip3 install future
-
-if [ "$PYTHON_VERSION" == "3.9" ]; then
-    pip3 install \
-        onnx==1.16.0 \
-        onnxruntime==1.19.2 \
-        onnxoptimizer==0.2.7
+pip3 install \
+    onnx==1.20.1 \
+    onnxruntime==1.24.1 \
+    onnxoptimizer==0.4.2
 
-    if [ "$DEVICE" == "cuda" ]; then
-        pip3 install \
-            torch==2.7.0 \
-            torchvision==0.22.0 \
-            --index-url https://download.pytorch.org/whl/cu118
-    else
-        pip3 install \
-            torch==2.7.0 \
-            torchvision==0.22.0 \
-            --extra-index-url https://download.pytorch.org/whl/cpu
-    fi
-elif [ "$PYTHON_VERSION" == "3.11" ]; then
+if [ "$DEVICE" == "cuda" ]; then
     pip3 install \
-        onnx==1.17.0 \
-        onnxruntime==1.20.1 \
-        onnxoptimizer==0.2.7
-
-    if [ "$DEVICE" == "cuda" ]; then
-        pip3 install \
-            torch==2.7.0 \
-            torchvision==0.22.0 \
-            --index-url https://download.pytorch.org/whl/cu118
-    else
-        pip3 install \
-            torch==2.7.0 \
-            torchvision==0.22.0 \
-            --extra-index-url https://download.pytorch.org/whl/cpu
-    fi
+        torch==2.10.0 \
+        torchvision==0.25.0
 else
     pip3 install \
-        onnx==1.12.0 \
-        onnxruntime==1.12.1 \
-        onnxoptimizer==0.2.7
-
-    if [ "$DEVICE" == "cuda" ]; then
-        pip3 install \
-            torch==2.4.1 \
-            torchvision==0.19.1
-    else
-        pip3 install \
-            torch==2.4.1 \
-            torchvision==0.19.1 \
-            --extra-index-url https://download.pytorch.org/whl/cpu
-    fi
+        torch==2.10.0 \
+        torchvision==0.25.0 \
+        --extra-index-url https://download.pytorch.org/whl/cpu
 fi

docker/install/ubuntu_install_python_package.sh

@@ -29,17 +29,17 @@ pip3 install --upgrade \
     numpy==1.21.* \
     orderedset \
     packaging \
-    Pillow==9.1.0 \
+    Pillow==12.1.1 \
     psutil \
     pytest \
     git+https://github.com/tlc-pack/tlcpack-sphinx-addon.git@768ec1dce349fe4708f6ad68be1ebb3f3dabafa1 \
     pytest-profiling \
     pytest-xdist \
-    pytest-rerunfailures==10.2 \
+    pytest-rerunfailures==16.1 \
     requests \
     scipy \
     Jinja2 \
-    junitparser==2.4.2 \
+    junitparser==4.0.2 \
     six \
     tornado \
     "ml_dtypes>=0.5.1"