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

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(

        "//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",

#    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

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

# Config setting that is satisfied when TensorFlow is being built with CUDA
# support through e.g. `--config=cuda` (or `--config=cuda_clang` in OSS).
alias(
    name = "is_cuda_enabled",
    actual = if_oss(
        "@local_config_cuda//:is_cuda_enabled",
        "@local_config_cuda//cuda:using_clang",
    ),
)

# Config setting that is satisfied when CUDA device code should be compiled

load(
    "//tensorflow:tensorflow.bzl",
    "tf_cc_test",
    "tf_copts",
    "tf_cuda_cc_test",
    "tf_cuda_library",
)
load("//tensorflow:tensorflow.default.bzl", "cc_header_only_library", "filegroup")
load(
    "//tensorflow/core/platform:build_config_root.bzl",
    "tf_cuda_tests_tags",
)
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")

package(

    export flags="--config=opt -k"
    ```

    If the tests are to be run on the GPU, add CUDA paths to LD_LIBRARY_PATH and
    add the `cuda` option flag

    ```bash
    export LD_LIBRARY_PATH="${LD_LIBRARY_PATH}:/usr/local/cuda/lib64:/usr/local/cuda/extras/CUPTI/lib64:$LD_LIBRARY_PATH"
    export flags="--config=opt --config=cuda -k"
    ```

    For example, to run all tests under tensorflow/python, do:

    ```bash