export DEBIAN_FRONTEND=noninteractive

# Fetch the NVIDIA key.
apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/3bf863cc.pub;

# Set up sources for NVIDIA CUDNN.
cat >/etc/apt/sources.list.d/nvidia.list <<SOURCES
# NVIDIA
deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/ /

nvidia-cuda-nvcc-cu12>=12.5.82,<13.0
nvidia-cuda-nvrtc-cu12>=12.5.82,<13.0
nvidia-cuda-runtime-cu12>=12.5.82,<13.0
# The upper bound is set for the CUDNN API compatibility.
# See
# https://docs.nvidia.com/deeplearning/cudnn/backend/latest/developer/forward-compatibility.html#cudnn-api-compatibility
nvidia-cudnn-cu12>=9.3.0.75,<10.0
nvidia-cufft-cu12>=11.2.3.61,<12.0
nvidia-curand-cu12>=10.3.6.82,<11.0
nvidia-cusolver-cu12>=11.6.3.83,<12.0

COPY setup.sources.cudnn.sh /setup.sources.cudnn.sh
ARG INSTALL_CUDNN=false
ARG CUDA_PACKAGES_FILE=cuda12.1_cudnn9.1.packages.txt
COPY ${CUDA_PACKAGES_FILE} /cudnn.packages.txt

RUN if [ "${INSTALL_CUDNN}" = "true" ]; then \
      echo "Installing cudnn"; \
      /setup.sources.cudnn.sh && /setup.packages.sh /cudnn.packages.txt; \
    else \
      echo "cudnn is not installed"; \
    fi

cuda-nvml-dev-12-1
libcufft-12-1
libcurand-12-1
libcusolver-dev-12-1
libcusparse-dev-12-1
libcublas-12-1
libcublas-dev-12-1
libnccl-dev=2.18.3-1+cuda12.1
libnccl2=2.18.3-1+cuda12.1
# CuDNN: https://docs.nvidia.com/deeplearning/sdk/cudnn-install/index.html#ubuntu-network-installation
libcudnn9-dev-cuda-12=9.8.0.87-1

cuda-nvml-dev-12-3
libcufft-12-3
libcurand-12-3
libcusolver-dev-12-3
libcusparse-dev-12-3
libcublas-12-3
libcublas-dev-12-3
libnccl-dev=2.19.3-1+cuda12.3
libnccl2=2.19.3-1+cuda12.3
# CuDNN: https://docs.nvidia.com/deeplearning/sdk/cudnn-install/index.html#ubuntu-network-installation
libcudnn9-dev-cuda-12=9.1.1.17-1

cuda-nvml-dev-12-8
libcufft-12-8
libcurand-12-8
libcusolver-dev-12-8
libcusparse-dev-12-8
libcublas-12-8
libcublas-dev-12-8
libnccl-dev=2.25.1-1+cuda12.8
libnccl2=2.25.1-1+cuda12.8
# CuDNN: https://docs.nvidia.com/deeplearning/sdk/cudnn-install/index.html#ubuntu-network-installation
libcudnn9-dev-cuda-12=9.8.0.87-1

cuda-nvml-dev-13-0
libcufft-13-0
libcurand-13-0
libcusolver-dev-13-0
libcusparse-dev-13-0
libcublas-13-0
libcublas-dev-13-0
libnccl-dev=2.27.7-1+cuda13.0
libnccl2=2.27.7-1+cuda13.0
# CuDNN: https://docs.nvidia.com/deeplearning/sdk/cudnn-install/index.html#ubuntu-network-installation
libcudnn9-dev-cuda-13=9.12.0.46-1

cuda-nvml-dev-12-1
libcufft-12-1
libcurand-12-1
libcusolver-dev-12-1
libcusparse-dev-12-1
libcublas-12-1
libcublas-dev-12-1
libnccl-dev=2.18.3-1+cuda12.1
libnccl2=2.18.3-1+cuda12.1
# CuDNN: https://docs.nvidia.com/deeplearning/sdk/cudnn-install/index.html#ubuntu-network-installation
libcudnn9-dev-cuda-12=9.1.1.17-1

       * Try adding the CUDA, CUPTI, and cuDNN installation directories to the %PATH% environment variable.
       * Refer [windows setup guide](https://www.tensorflow.org/install/gpu#windows_setup).
     * If you have above configuration and using _**Ubuntu/Linux**_ platform -
       * Try adding the CUDA, CUPTI, and cuDNN installation directories to the $LD_LIBRARY_PATH environment variable.

    # via
    #   -c ci/official/requirements_updater/nvidia-constraints.txt
    #   -r ci/official/requirements_updater/nvidia-constraints.txt
    #   -r ci/official/requirements_updater/nvidia-requirements.txt
    #   nvidia-cudnn-cu12
    #   nvidia-cusolver-cu12
nvidia-cuda-cupti-cu12==12.5.82 \
    --hash=sha256:4f835281cf492e2bedd153f5c3de9da8f1d775a419468305e64ce73b3b0c6dc3 \