libcudnn8=8.9.6.50-1+cuda12.2
libcudnn8-dev=8.9.6.50-1+cuda12.2

# This can be removed once NVIDIA publishes a cuda-12.3.2 Docker image.
# For now it ensures that we install at least version 12.3.107 of PTXAS,
# since 12.3.103 has a bug.

source ci/official/envs/linux_x86
TFCI_BAZEL_COMMON_ARGS="--repo_env=TF_PYTHON_VERSION=$TFCI_PYTHON_VERSION --config release_gpu_linux"
TFCI_BAZEL_TARGET_SELECTING_CONFIG_PREFIX=linux_cuda
TFCI_BUILD_PIP_PACKAGE_ARGS="--repo_env=WHEEL_NAME=tensorflow"
TFCI_DOCKER_ARGS="--gpus all"
TFCI_LIB_SUFFIX="-gpu-linux-x86_64"

def set_tf_cuda_version(environ_cp):
  """Set TF_CUDA_VERSION."""
  ask_cuda_version = (
      'Please specify the CUDA SDK version you want to use. '
      '[Leave empty to default to CUDA %s]: ') % _DEFAULT_CUDA_VERSION
  tf_cuda_version = get_from_env_or_user_or_default(environ_cp,
                                                    'TF_CUDA_VERSION',

    --@local_config_cuda//:enable_cuda \
    --define framework_shared_object=false \
    "somepath(//tensorflow/tools/pip_package:wheel, " \
    "@local_config_cuda//cuda:cudart + "\
    "@local_config_cuda//cuda:cudart + "\
    "@local_config_cuda//cuda:cuda_driver + "\
    "@local_config_cuda//cuda:cudnn + "\
    "@local_config_cuda//cuda:curand + "\
    "@local_config_cuda//cuda:cusolver + "\

build:mkl_aarch64_threadpool -c opt

# CUDA: This config refers to building CUDA op kernels with nvcc.
build:cuda --repo_env TF_NEED_CUDA=1
build:cuda --crosstool_top=@local_config_cuda//crosstool:toolchain
build:cuda --@local_config_cuda//:enable_cuda

# CUDA: This config refers to building CUDA op kernels with clang.
build:cuda_clang --config=cuda
# Enable TensorRT optimizations https://developer.nvidia.com/tensorrt

        "//tensorflow/core:test",
        "//tensorflow/core:test_main",
    ],
)

tf_cuda_cc_test(
    name = "nn_grad_test",
    size = "small",
    srcs = [
        "nn_grad_test.cc",
    ],
    args = ["--heap_check="],  # TODO(b/174752220): Remove
    tags = tf_cuda_tests_tags() + ["no_cuda_asan"],  # b/173654156,
    deps = [
        ":grad_test_helper",
        ":nn_grad",

load(
    "//tensorflow:tensorflow.bzl",
    "check_deps",
    "if_google",
    "if_not_mobile",
    "tf_cc_test",
    "tf_copts",
    "tf_cuda_library",
    "tf_custom_op_library",
    "tf_kernel_library",
)
load("//tensorflow:tensorflow.default.bzl", "filegroup", "tf_cuda_cc_test")
load(
    "//tensorflow/core/tpu:build_defs.bzl",
    "if_libtpu_tf_status",
)

package(

#    build. This should also match the system you're using--you cannot build
#    the TF MacOS package from Linux.
#      Ex. linux_x86        -- x86_64 Linux platform
#      Ex. linux_x86_cuda   -- x86_64 Linux platform, with Nvidia CUDA support
#      Ex. macos_arm64      -- arm64 MacOS platform
# 3. Add modifiers. Some modifiers for local execution are:
#      Ex. disk_cache -- Use a local cache

# Install devtoolset devel dependencies
COPY setup.sources.sh /setup.sources.sh
COPY setup.packages.sh /setup.packages.sh
COPY devel.packages.txt /devel.packages.txt
COPY cuda.packages.txt /cuda.packages.txt
RUN /setup.sources.sh && /setup.packages.sh /devel.packages.txt

# Install various tools.
# - bats: bash unit testing framework
#         NOTE: v1.6.0 seems to have a bug that made "git" in setup_file break

fi
# VERY basic check to ensure the [and-cuda] package variant is installable.
# Checks TFCI_BAZEL_COMMON_ARGS for "gpu" or "cuda", implying that the test is
# relevant. All of the GPU test machines have CUDA installed via other means,
# so I am not sure how to verify that the dependencies themselves are valid for
# the moment.
if [[ "$TFCI_BAZEL_COMMON_ARGS" =~ gpu|cuda ]]; then