wehub-resource-sync
559 lines
20 KiB
Python
559 lines
20 KiB
Python
# Description:
|
|
# Tools for building the TensorFlow pip package.
|
|
load("@cuda_cudart//:version.bzl", _cudart_version = "VERSION")
|
|
load("@local_config_cuda//cuda:build_defs.bzl", "if_cuda")
|
|
load("@local_config_syslibs//:build_defs.bzl", "if_not_system_lib")
|
|
load("@nvidia_wheel_versions//:versions.bzl", "NVIDIA_WHEEL_VERSIONS")
|
|
load(
|
|
"@xla//third_party/py:py_import.bzl",
|
|
"py_import",
|
|
)
|
|
load(
|
|
"@xla//third_party/py:py_manylinux_compliance_test.bzl",
|
|
"verify_manylinux_compliance_test",
|
|
)
|
|
load("@xla//third_party/py:python_wheel.bzl", "collect_data_files", "transitive_py_deps")
|
|
load("@xla//third_party/rules_python/python:defs.bzl", "py_binary", "py_test")
|
|
load("@xla//xla/tsl/mkl:build_defs.bzl", "if_enable_mkl", "if_mkl", "if_mkl_ml")
|
|
load("//tensorflow:tensorflow.bzl", "if_wheel_dependency", "if_with_tpu_support", "transitive_hdrs")
|
|
load("//tensorflow:tf_version.bzl", "TF_SEMANTIC_VERSION_SUFFIX", "TF_VERSION")
|
|
load(
|
|
"//tensorflow/core/platform:build_config_root.bzl",
|
|
"if_pywrap",
|
|
"tf_additional_license_deps",
|
|
"tf_cuda_tests_tags",
|
|
"tf_exec_properties",
|
|
)
|
|
load("//tensorflow/tools/pip_package/utils:tf_wheel.bzl", "tf_wheel", "tf_wheel_dep")
|
|
|
|
package(default_visibility = ["//visibility:public"])
|
|
|
|
# This returns a list of headers of all public header libraries (e.g.,
|
|
# framework, lib), and all of the transitive dependencies of those
|
|
# public headers. Not all of the headers returned by the filegroup
|
|
# are public (e.g., internal headers that are included by public
|
|
# headers), but the internal headers need to be packaged in the
|
|
# pip_package for the public headers to be properly included.
|
|
#
|
|
# Public headers are therefore defined by those that are both:
|
|
#
|
|
# 1) "publicly visible" as defined by bazel
|
|
# 2) Have documentation.
|
|
#
|
|
# This matches the policy of "public" for our python API.
|
|
transitive_hdrs(
|
|
name = "included_headers",
|
|
deps = [
|
|
"//tensorflow/compiler/tf2xla:xla_compiled_cpu_function",
|
|
"//tensorflow/compiler/mlir:mlir_graph_optimization_pass",
|
|
"//tensorflow/core:core_cpu",
|
|
"//tensorflow/core:framework",
|
|
"//tensorflow/core:lib",
|
|
"//tensorflow/core:protos_all_cc",
|
|
"//tensorflow/core/platform:stream_executor",
|
|
"//tensorflow/cc/saved_model:loader",
|
|
"//tensorflow/cc/saved_model:reader",
|
|
"//tensorflow/cc/saved_model:bundle_v2",
|
|
"//tensorflow/c/experimental/filesystem:filesystem_interface",
|
|
"//tensorflow/c/experimental/next_pluggable_device:c_api_hdrs",
|
|
"//tensorflow/c/experimental/pluggable_profiler:pluggable_profiler_hdrs",
|
|
"//tensorflow/c/experimental/stream_executor:stream_executor_hdrs",
|
|
"//tensorflow/c/experimental/grappler:grappler_hdrs",
|
|
"//tensorflow/c:kernels_hdrs",
|
|
"//tensorflow/c:ops_hdrs",
|
|
"@xla//xla/pjrt/c:pjrt_c_api_hdrs",
|
|
"//tensorflow/core/common_runtime/next_pluggable_device/c:plugin_c_api_hdrs",
|
|
# TODO(rostam): Revisit these after cc_shared_library everywhere
|
|
"//tensorflow/cc:client_session",
|
|
"//tensorflow/cc:cc_ops",
|
|
# WARNING: None of the C/C++ code under python/ has any API guarantees, and TF team
|
|
# reserves the right to change APIs and other header-level interfaces. If your custom
|
|
# op uses these headers, it may break when users upgrade their version of tensorflow.
|
|
# NOTE(ebrevdo): None of the python symbols are exported by libtensorflow_framework.so.
|
|
# Code that relies on these headers should dynamically link to
|
|
# _pywrap_tensorflow_internal.so as well.
|
|
"//tensorflow/python/grappler:model_analyzer_lib",
|
|
"//tensorflow/python/lib/core:py_exception_registry",
|
|
"//tensorflow/python/lib/core:pybind11_proto",
|
|
"//tensorflow/python/lib/core:pybind11_absl",
|
|
"//tensorflow/python/lib/core:pybind11_lib",
|
|
"//tensorflow/python/lib/core:pybind11_status",
|
|
"//tensorflow/python/lib/core:py_func_lib",
|
|
"//tensorflow/python/lib/core:py_seq_tensor",
|
|
"//tensorflow/python/lib/core:py_util",
|
|
"//tensorflow/python/util:cpp_python_util",
|
|
"//tensorflow/python/util:kernel_registry",
|
|
"//tensorflow/python/framework:python_op_gen",
|
|
"//tensorflow/python/client:tf_session_helper",
|
|
] + if_cuda([
|
|
"@local_config_cuda//cuda:cuda_headers",
|
|
]),
|
|
)
|
|
|
|
filegroup(
|
|
name = "xla_cmake",
|
|
srcs = ["xla_build/CMakeLists.txt"],
|
|
)
|
|
|
|
filegroup(
|
|
name = "licenses",
|
|
srcs = [
|
|
"//:LICENSE",
|
|
"//third_party/icu/data:LICENSE",
|
|
"@XNNPACK//:LICENSE",
|
|
"@com_google_absl//:LICENSE",
|
|
"@com_google_protobuf//:LICENSE",
|
|
"@curl//:COPYING",
|
|
"@ducc//:LICENSE",
|
|
"@flatbuffers//:LICENSE",
|
|
"@gemmlowp//:LICENSE",
|
|
"@jpegxl//:LICENSE",
|
|
"@libjpeg_turbo//:LICENSE.md",
|
|
"@llvm-project//llvm:LICENSE.TXT",
|
|
"@llvm-project//mlir:LICENSE.TXT",
|
|
"@local_config_tensorrt//:LICENSE",
|
|
"@ml_dtypes_py//:LICENSE",
|
|
"@org_brotli//:LICENSE",
|
|
"@png//:LICENSE",
|
|
"@ruy//:LICENSE",
|
|
"@snappy//:COPYING",
|
|
"@stablehlo//:LICENSE",
|
|
"@tf_runtime//:LICENSE",
|
|
] + select({
|
|
"//tensorflow:android": [],
|
|
"//tensorflow:ios": [],
|
|
"//tensorflow:linux_s390x": [],
|
|
"//tensorflow:windows": [],
|
|
"//conditions:default": [
|
|
"@com_github_googlecloudplatform_google_cloud_cpp//:LICENSE",
|
|
],
|
|
}) + if_cuda([
|
|
"@local_config_nccl//:LICENSE",
|
|
]) + if_mkl([
|
|
"@xla//third_party/mkl_dnn:LICENSE",
|
|
]) + if_enable_mkl(["@xla//xla/tsl/mkl:LICENSE"]) + if_not_system_lib(
|
|
"absl_py",
|
|
[
|
|
"@absl_py//:LICENSE",
|
|
],
|
|
) + tf_additional_license_deps(),
|
|
)
|
|
|
|
# Add dynamic kernel dso files here.
|
|
DYNAMIC_LOADED_KERNELS = [
|
|
"//tensorflow/core/kernels:libtfkernel_sobol_op.so",
|
|
]
|
|
|
|
COMMON_PIP_DEPS = [
|
|
"//tensorflow/compiler/mlir/stablehlo:stablehlo",
|
|
"//tensorflow/compiler/mlir/tensorflow:gen_mlir_passthrough_op_py",
|
|
"//tensorflow/core/function/trace_type:serialization_test_proto_py",
|
|
"//tensorflow/core/function/trace_type:serialization",
|
|
"//tensorflow/dtensor/python/tests:multi_client_test_util",
|
|
"//tensorflow/dtensor/python/tests:test_util",
|
|
"//tensorflow/lite/python:tflite_convert",
|
|
"//tensorflow/lite/tools:visualize",
|
|
"//tensorflow/python/autograph/core:test_lib",
|
|
"//tensorflow/python/autograph/impl/testing:pybind_for_testing",
|
|
"//tensorflow/python/autograph/pyct/testing:basic_definitions",
|
|
"//tensorflow/python/autograph/pyct/testing:decorators",
|
|
"//tensorflow/python/autograph/pyct/common_transformers:common_transformers",
|
|
"//tensorflow/python/autograph/pyct/static_analysis:type_inference",
|
|
"//tensorflow/python/autograph/utils:context_managers",
|
|
"//tensorflow/python/autograph/utils:tensor_list",
|
|
"//tensorflow/python/compiler:compiler",
|
|
"//tensorflow/python/ops:cond_v2",
|
|
"//tensorflow/python:distributed_framework_test_lib",
|
|
"//tensorflow/python/distribute:distribute_test_lib_pip",
|
|
"//tensorflow/python/training/experimental:loss_scale",
|
|
"//tensorflow/python/training/experimental:loss_scale_optimizer",
|
|
"//tensorflow/python/data/benchmarks:benchmark_base",
|
|
"//tensorflow/python/data/experimental/kernel_tests/service:multi_process_cluster",
|
|
"//tensorflow/python/data/experimental/kernel_tests/service:test_base",
|
|
"//tensorflow/python/data/experimental/ops:testing",
|
|
"//tensorflow/python/data/experimental/service:server_lib",
|
|
"//tensorflow/python/data/kernel_tests:checkpoint_test_base",
|
|
"//tensorflow/python/data/kernel_tests:test_base",
|
|
"//tensorflow/python/data/kernel_tests:tf_record_test_base",
|
|
"//tensorflow/python/ops/ragged:ragged_tensor_test_ops",
|
|
"//tensorflow/python/debug:debug_pip",
|
|
"//tensorflow/python/distribute:combinations",
|
|
"//tensorflow/python/distribute/coordinator:fault_tolerance_test_base",
|
|
"//tensorflow/python/distribute/failure_handling:check_preemption_py",
|
|
"//tensorflow/python/distribute/failure_handling:failure_handling_lib",
|
|
"//tensorflow/python/distribute/failure_handling:failure_handling_util",
|
|
"//tensorflow/python/distribute:multi_process_runner",
|
|
"//tensorflow/python/eager:eager_pip",
|
|
"//tensorflow/python/keras:combinations",
|
|
"//tensorflow/python/keras/mixed_precision:test_util",
|
|
"//tensorflow/python/keras/utils:dataset_creator",
|
|
"//tensorflow/python/kernel_tests/nn_ops:bias_op_base",
|
|
"//tensorflow/python/kernel_tests/nn_ops:cudnn_deterministic_base",
|
|
"//tensorflow/python/kernel_tests/nn_ops:depthwise_conv_op_base",
|
|
"//tensorflow/python/kernel_tests/nn_ops:xent_op_test_base",
|
|
"//tensorflow/python/kernel_tests/random:util",
|
|
"//tensorflow/python/kernel_tests/signal:test_util",
|
|
"//tensorflow/python/kernel_tests/sparse_ops:sparse_xent_op_test_base",
|
|
"//tensorflow/python/lib:__init__",
|
|
"//tensorflow/python/ops/parallel_for:test_util",
|
|
"//tensorflow/python/ops/structured:structured_tensor_dynamic",
|
|
"//tensorflow/python/platform:resource_loader",
|
|
"//tensorflow/python/profiler:trace",
|
|
"//tensorflow/python/saved_model:saved_model",
|
|
"//tensorflow/python/summary:__init__",
|
|
"//tensorflow/python/summary:plugin_asset",
|
|
"//tensorflow/python/summary:summary_iterator",
|
|
"//tensorflow/python/summary:summary_py",
|
|
"//tensorflow/python/tools:tools_pip",
|
|
"//tensorflow/python/tools/api/generator:create_python_api",
|
|
"//tensorflow/python/tools/api/generator2/extractor:extractor",
|
|
"//tensorflow/python/tools/api/generator2/generator:generator",
|
|
"//tensorflow/python/tpu",
|
|
"//tensorflow/python/util:deprecated_module",
|
|
"//tensorflow/python/util:deprecated_module_new",
|
|
"//tensorflow/python/util:example_parser_configuration",
|
|
"//tensorflow/python/util:function_utils",
|
|
"//tensorflow/python/util:keyword_args",
|
|
"//tensorflow/python/util:lock_util",
|
|
"//tensorflow/python/util:module_wrapper",
|
|
"//tensorflow/python/util:serialization",
|
|
"//tensorflow/python/ops:image_grad_test_base",
|
|
"//tensorflow/python/framework:memory_checker",
|
|
"//tensorflow/python/framework:test_ops",
|
|
"//tensorflow/python/ops:while_v2",
|
|
"//tensorflow/tools/common:public_api",
|
|
"//tensorflow/tools/common:test_module1",
|
|
"//tensorflow/tools/common:traverse",
|
|
"//tensorflow/tools/docs:tf_doctest_lib",
|
|
"@tsl//tsl/profiler/protobuf:trace_events_proto_py",
|
|
"//tensorflow/python/distribute:parameter_server_strategy_v2",
|
|
"//tensorflow/python/distribute/coordinator:cluster_coordinator",
|
|
"//tensorflow/python/distribute/coordinator:remote_eager_lib",
|
|
"//tensorflow/python/distribute/coordinator:metric_utils",
|
|
"//tensorflow/python/distribute/experimental/rpc:rpc_ops",
|
|
"//tensorflow/python/util:pywrap_xla_ops",
|
|
"//tensorflow:tensorflow_py",
|
|
"//tensorflow/tools/compatibility:tf_upgrade_v2",
|
|
] + if_with_tpu_support([
|
|
"//tensorflow/python/tools:grpc_tpu_worker",
|
|
"//tensorflow/python/tools:grpc_tpu_worker_service",
|
|
])
|
|
|
|
py_binary(
|
|
name = "build_pip_package_py",
|
|
srcs = [
|
|
"build_pip_package.py",
|
|
],
|
|
main = "build_pip_package.py",
|
|
deps = [
|
|
"//tensorflow/tools/pip_package/utils:py_utils",
|
|
"@pypi//setuptools",
|
|
"@pypi//wheel",
|
|
],
|
|
)
|
|
|
|
collect_data_files(
|
|
name = "cc_deps",
|
|
deps = [
|
|
#`tensorflow_py` is the main target to collect all required `data` dependencies
|
|
"//tensorflow:tensorflow_py",
|
|
# utils below are not part of tensorflow_py DAG
|
|
"//tensorflow/compiler/mlir/stablehlo:stablehlo_extension",
|
|
"//tensorflow/python/autograph/impl/testing:pybind_for_testing",
|
|
"//tensorflow/python/framework:memory_checker",
|
|
"//tensorflow/python/util:pywrap_xla_ops",
|
|
],
|
|
)
|
|
|
|
transitive_py_deps(
|
|
name = "py_deps",
|
|
deps = COMMON_PIP_DEPS,
|
|
)
|
|
|
|
py_binary(
|
|
name = "modify_setup_py_binary",
|
|
srcs = [
|
|
"modify_setup_py.py",
|
|
],
|
|
main = "modify_setup_py.py",
|
|
deps = ["@xla//third_party/py:setup_py_nvidia_dependencies_util"],
|
|
)
|
|
|
|
_cuda_major_version_split = _cudart_version.split(".")[0] if _cudart_version else ""
|
|
|
|
cuda_major_version = _cuda_major_version_split if _cuda_major_version_split else "12"
|
|
|
|
cuda_suffix = "_cu12" if cuda_major_version == "12" else ""
|
|
|
|
genrule(
|
|
name = "setup_py",
|
|
srcs = ["setup.py.tpl"],
|
|
outs = ["setup.py"],
|
|
cmd = """
|
|
$(location :modify_setup_py_binary) \\
|
|
--template_file $(location setup.py.tpl) \\
|
|
--output_file $(OUTS) \\
|
|
--nvidia_wheel_versions_data '{nvidia_wheel_versions}' \\
|
|
--tf_version "{wheel_version}{wheel_version_suffix}" \\
|
|
--cuda_version {cuda_version}
|
|
""".format(
|
|
cuda_version = cuda_major_version,
|
|
nvidia_wheel_versions = NVIDIA_WHEEL_VERSIONS,
|
|
wheel_version = TF_VERSION,
|
|
wheel_version_suffix = TF_SEMANTIC_VERSION_SUFFIX,
|
|
),
|
|
tools = [":modify_setup_py_binary"],
|
|
)
|
|
|
|
py_binary(
|
|
name = "setup_py_binary",
|
|
srcs = [":setup_py"],
|
|
data = [
|
|
"//:AUTHORS",
|
|
"//:LICENSE",
|
|
],
|
|
main = "setup.py",
|
|
)
|
|
|
|
tf_wheel(
|
|
name = "wheel",
|
|
# Explicitly list all external repositories that are needed in build_pip_package.py
|
|
external_repos = [
|
|
"com_google_absl",
|
|
"eigen_archive",
|
|
"jsoncpp_git",
|
|
"com_google_protobuf",
|
|
"xla",
|
|
"tsl",
|
|
"llvm-project",
|
|
"cpuinfo",
|
|
"FXdiv",
|
|
"net_zstd",
|
|
"org_brotli",
|
|
"pthreadpool",
|
|
"riegeli",
|
|
"XNNPACK",
|
|
"local_config_cuda",
|
|
"local_config_tensorrt",
|
|
"cuda_cccl",
|
|
"cuda_cublas",
|
|
"cuda_cudart",
|
|
"cuda_cudnn",
|
|
"cuda_cufft",
|
|
"cuda_cupti",
|
|
"cuda_curand",
|
|
"cuda_cusolver",
|
|
"cuda_cusparse",
|
|
"cuda_nvcc",
|
|
"cuda_nvml",
|
|
"cuda_nvrtc",
|
|
"cuda_nvtx",
|
|
"cuda_nvjitlink",
|
|
],
|
|
headers = [
|
|
":included_headers",
|
|
":licenses",
|
|
"//tensorflow/core:protos_all_proto_srcs",
|
|
],
|
|
platform_name = select({
|
|
"@platforms//os:osx": "macosx",
|
|
"@platforms//os:macos": "macosx",
|
|
"@platforms//os:windows": "win",
|
|
"@platforms//os:linux": "linux",
|
|
}),
|
|
platform_tag = select({
|
|
"@platforms//cpu:aarch64": "arm64",
|
|
"@platforms//cpu:arm64": "arm64",
|
|
"@platforms//cpu:x86_64": "x86_64",
|
|
"@platforms//cpu:ppc": "ppc64le",
|
|
}),
|
|
source_files = [
|
|
"MANIFEST.in",
|
|
"//tensorflow/tools/pip_package:THIRD_PARTY_NOTICES.txt",
|
|
":setup_py",
|
|
":cc_deps",
|
|
":py_deps",
|
|
"//tensorflow:tf_python_api_gen_v2",
|
|
] + select({
|
|
"//tensorflow:dynamic_loaded_kernels": if_pywrap(
|
|
if_false = DYNAMIC_LOADED_KERNELS,
|
|
if_true = [],
|
|
),
|
|
"//conditions:default": [],
|
|
}) + select({
|
|
"//conditions:default": if_pywrap(
|
|
if_false = [
|
|
"//tensorflow:tensorflow_cc",
|
|
"//tensorflow:tensorflow_framework",
|
|
],
|
|
if_true = [
|
|
"//tensorflow/python:pywrap_tensorflow_binaries",
|
|
],
|
|
),
|
|
"//tensorflow:windows": if_pywrap(
|
|
if_false = [
|
|
"//tensorflow/python:pywrap_tensorflow_import_lib_file",
|
|
],
|
|
if_true = [
|
|
"//tensorflow/python:pywrap_tensorflow_binaries",
|
|
],
|
|
),
|
|
}) + if_mkl_ml(["@xla//xla/tsl/mkl:intel_binary_blob"]),
|
|
xla_aot_compiled = [
|
|
"//tensorflow/compiler/tf2xla:xla_compiled_cpu_runtime_srcs",
|
|
":xla_cmake",
|
|
],
|
|
)
|
|
|
|
genrule(
|
|
name = "empty_test",
|
|
outs = ["empty_test.py"],
|
|
cmd = "echo '' > $@",
|
|
)
|
|
|
|
py_test(
|
|
name = "prebuilt_wheel_import_api_packages_test_cpu",
|
|
srcs = if_wheel_dependency(
|
|
["import_api_packages_test.py"],
|
|
[":empty_test"],
|
|
),
|
|
main = if_wheel_dependency("import_api_packages_test.py", "empty_test.py"),
|
|
tags = [
|
|
"cpu",
|
|
"manual",
|
|
],
|
|
deps = if_wheel_dependency([":prebuilt_tf_py_import"]) + [
|
|
"@pypi//absl_py",
|
|
],
|
|
)
|
|
|
|
py_test(
|
|
name = "prebuilt_wheel_import_api_packages_test_gpu",
|
|
srcs = if_wheel_dependency(
|
|
["import_api_packages_test.py"],
|
|
[":empty_test"],
|
|
),
|
|
exec_properties = if_cuda(
|
|
tf_exec_properties({"tags": tf_cuda_tests_tags()}),
|
|
{},
|
|
),
|
|
main = if_wheel_dependency("import_api_packages_test.py", "empty_test.py"),
|
|
tags = [
|
|
"gpu",
|
|
"manual",
|
|
],
|
|
deps = if_wheel_dependency([":prebuilt_tf_py_import"]) + [
|
|
"@pypi//absl_py",
|
|
],
|
|
)
|
|
|
|
py_test(
|
|
name = "import_api_packages_test_cpu",
|
|
srcs = ["import_api_packages_test.py"],
|
|
main = "import_api_packages_test.py",
|
|
tags = [
|
|
"cpu",
|
|
"manual",
|
|
],
|
|
deps = [
|
|
":tf_py_import",
|
|
"@pybind11_abseil//pybind11_abseil:status_casters",
|
|
],
|
|
)
|
|
|
|
py_test(
|
|
name = "import_api_packages_test_gpu",
|
|
srcs = ["import_api_packages_test.py"],
|
|
exec_properties = if_cuda(
|
|
tf_exec_properties({"tags": tf_cuda_tests_tags()}),
|
|
{},
|
|
),
|
|
main = "import_api_packages_test.py",
|
|
tags = [
|
|
"gpu",
|
|
"manual",
|
|
],
|
|
deps = [
|
|
":tf_py_import",
|
|
"@pybind11_abseil//pybind11_abseil:status_casters",
|
|
],
|
|
)
|
|
|
|
verify_manylinux_compliance_test(
|
|
name = "manylinux_compliance_test",
|
|
aarch64_compliance_tag = "manylinux_2_27_aarch64",
|
|
ppc64le_compliance_tag = "manylinux_2_17_ppc64le",
|
|
test_tags = [
|
|
"manual",
|
|
],
|
|
wheel = ":wheel",
|
|
x86_64_compliance_tag = "manylinux_2_27_x86_64",
|
|
)
|
|
|
|
filegroup(
|
|
name = "nvidia_wheel_deps",
|
|
srcs = if_cuda([
|
|
"@pypi_nvidia_cublas{cuda}//:pkg".format(cuda = cuda_suffix),
|
|
"@pypi_nvidia_cuda_cupti{cuda}//:pkg".format(cuda = cuda_suffix),
|
|
"@pypi_nvidia_cuda_nvcc{cuda}//:pkg".format(cuda = cuda_suffix),
|
|
"@pypi_nvidia_cuda_nvrtc{cuda}//:pkg".format(cuda = cuda_suffix),
|
|
"@pypi_nvidia_cuda_runtime{cuda}//:pkg".format(cuda = cuda_suffix),
|
|
"@pypi_nvidia_cudnn_cu{cuda}//:pkg".format(cuda = cuda_major_version),
|
|
"@pypi_nvidia_cufft{cuda}//:pkg".format(cuda = cuda_suffix),
|
|
"@pypi_nvidia_curand{cuda}//:pkg".format(cuda = cuda_suffix),
|
|
"@pypi_nvidia_cusolver{cuda}//:pkg".format(cuda = cuda_suffix),
|
|
"@pypi_nvidia_cusparse{cuda}//:pkg".format(cuda = cuda_suffix),
|
|
"@pypi_nvidia_nccl_cu{cuda}//:pkg".format(cuda = cuda_major_version),
|
|
"@pypi_nvidia_nvjitlink{cuda}//:pkg".format(cuda = cuda_suffix),
|
|
"@pypi_nvidia_nvshmem_cu{cuda}//:pkg".format(cuda = cuda_major_version),
|
|
]) + if_cuda(([
|
|
"@pypi_nvidia_nvvm{cuda}//:pkg".format(cuda = cuda_suffix),
|
|
"@pypi_nvidia_cuda_crt{cuda}//:pkg".format(cuda = cuda_suffix),
|
|
] if cuda_major_version == "13" else [])),
|
|
)
|
|
|
|
py_import(
|
|
name = "tf_py_import",
|
|
wheel = ":wheel",
|
|
wheel_deps = if_cuda([":nvidia_wheel_deps"]),
|
|
deps = [
|
|
"@pypi//absl_py",
|
|
"@pypi//astunparse",
|
|
"@pypi//flatbuffers",
|
|
"@pypi//gast",
|
|
"@pypi//ml_dtypes",
|
|
"@pypi//numpy",
|
|
"@pypi//opt_einsum",
|
|
"@pypi//packaging",
|
|
"@pypi//protobuf",
|
|
"@pypi//requests",
|
|
"@pypi//termcolor",
|
|
"@pypi//typing_extensions",
|
|
"@pypi//wrapt",
|
|
],
|
|
)
|
|
|
|
py_import(
|
|
name = "prebuilt_tf_py_import",
|
|
wheel = if_wheel_dependency(
|
|
tf_wheel_dep(),
|
|
":wheel",
|
|
),
|
|
wheel_deps = if_cuda([":nvidia_wheel_deps"]),
|
|
deps = [
|
|
"@pypi//absl_py",
|
|
"@pypi//astunparse",
|
|
"@pypi//flatbuffers",
|
|
"@pypi//gast",
|
|
"@pypi//ml_dtypes",
|
|
"@pypi//numpy",
|
|
"@pypi//opt_einsum",
|
|
"@pypi//packaging",
|
|
"@pypi//protobuf",
|
|
"@pypi//requests",
|
|
"@pypi//termcolor",
|
|
"@pypi//typing_extensions",
|
|
"@pypi//wrapt",
|
|
],
|
|
)
|