Merge branch '25.07-devel' into 25.07-devel-add-ngc-release-testing

Author: aybchan

.github/container/build-jax.sh

@@ -51,6 +51,7 @@ usage() {
     echo "    --clean-only                   Do not build, just cleanup"
     echo "    --cpu-arch                     Target CPU architecture, e.g. amd64, arm64, etc."
     echo "    --debug                        Build in debug mode"
+    echo "    -r, --release                  Build for release"
     echo "    -h, --help                     Print usage."
     echo "    --install                      Install the JAX wheels when build succeeds"
     echo "    --no-install                   Do not install the JAX wheels when build succeeds"
@@ -76,10 +77,11 @@ CPU_ARCH="$(dpkg --print-architecture)"
 CUDA_COMPUTE_CAPABILITIES="local"
 DEBUG=0
 INSTALL=1
+IS_RELEASE=0
 SRC_PATH_JAX="/opt/jax"
 SRC_PATH_XLA="/opt/xla"
 
-args=$(getopt -o h --long bazel-cache:,bazel-cache-namespace:,build-param:,build-path-jaxlib:,clean,cpu-arch:,debug,no-clean,clean-only,help,install,no-install,src-path-jax:,src-path-xla:,sm: -- "$@")
+args=$(getopt -o h,r --long bazel-cache:,bazel-cache-namespace:,build-param:,build-path-jaxlib:,clean,release,cpu-arch:,debug,no-clean,clean-only,help,install,no-install,src-path-jax:,src-path-xla:,sm: -- "$@")
 if [[ $? -ne 0 ]]; then
     exit 1
 fi
@@ -123,6 +125,10 @@ while [ : ]; do
             DEBUG=1
             shift 1
             ;;
+        -r | --release)
+            IS_RELEASE=1
+            shift 1
+            ;;
         -h | --help)
             usage 1
             ;;
@@ -206,6 +212,7 @@ print_var INSTALL
 print_var PYTHON_VERSION
 print_var SRC_PATH_JAX
 print_var SRC_PATH_XLA
+print_var IS_RELEASE
 
 echo "=================================================="
 
@@ -248,6 +255,12 @@ done
 bazel run --config=cuda_libraries_from_stubs --verbose_failures=true //build:requirements.update --repo_env=HERMETIC_PYTHON_VERSION="${PYTHON_VERSION}"
 popd
 
+
+if [[ "${IS_RELEASE}" == "1" ]]; then
+    jaxlib_version=$(pip show jaxlib | grep Version | tr ':' '\n' | tail -1)
+    sed -i "s|      f'jaxlib >={_minimum_jaxlib_version}, <={_jax_version}',|      f'jaxlib>=0.5.0',|" /opt/jax/setup.py
+fi
+
 ## Install the built packages
 
 if [[ "${INSTALL}" == "1" ]]; then