OSSForks/tensorflow
1
0
Fork 0
mirror of https://github.com/zebrajr/tensorflow.git synced 2025-12-06 12:20:11 +01:00
Code Issues Actions 5 Packages Projects Releases Wiki Activity
5720ab7845
tensorflow/tensorflow/python/tpu/BUILD
Vadym Matsishevskyi 5720ab7845 Introduce pywrap bazel rules and migrate Tensorflow to it 
The gist of this change (the new rules implementation) is contained within `rules_pywrap` folder. The rules are generic and not tensorflow-specific

1) (internal-specific)

2) (internal-specific)

3) It provides same linking strategy of final artifacts on all 3 supported platforms (no major differences between Linux, Mac and Windows).

4) It makes it possible to abandon usage of header-only targets to prevent ODR violations. Simply speaking you can now depend on generated protobuf message classes normally, without need to worry how that is linked afterwards.

5) The current version is backward-compatible and unless explicitly enabled is a no-op. To enable the new rules pass `--repo_env=USE_PYWRAP_RULES=True` flag to build/test command.

6) The `if_pywrap` construct is temporary and will be removed once full migration is completed. Currently if_pywrap is mainly used to pass normal dependencies (instead of header-only). The header-only stuff is kept for backward compatibility and smoother migration but will be eventually removed.

7) This CL migrates TF and the most problematic among all google ML repositories. Once TF is sabilized the other repositories, such as JAX and XLA will be migrated too (which should be way easier than migrating TF anyways)

PiperOrigin-RevId: 684324990
2024-10-10 00:12:03 -07:00

1208 lines
40 KiB
Python
Raw Blame History

load("@bazel_skylib//:bzl_library.bzl", "bzl_library")
# Description: Operations defined for Cloud TPUs
load("//tensorflow:pytype.default.bzl", "pytype_strict_library")
load("//tensorflow:strict.default.bzl", "py_strict_library", "py_strict_test")
# Placeholder: load py_proto_library
load("//tensorflow:tensorflow.default.bzl", "tf_py_strict_test", "tf_python_pybind_extension")
load("//tensorflow/core/platform:build_config.bzl", "tf_proto_library")
load("//tensorflow/core/platform:build_config_root.bzl", "if_pywrap")
load("//tensorflow/python/tpu:tpu.bzl", "internal_create_sanitizer_settings", "tpu_py_strict_test")
# Do not add anymore paths here. You do not need to be in the visibility list
# to use TPU symbols. They are accessible from tf.contrib.tpu in TF 1.x and
# tf.tpu and tf.compat.v1.tpu in TF 2.x.
visibility = [
"//learning/brain:__subpackages__",
"//learning/deepmind:__subpackages__",
"//learning/serving:__subpackages__",
"//research/graph:__subpackages__",
"//smartass/brain/configure:__subpackages__",
"//third_party/py/jax_tpu_embedding:__subpackages__",
"//third_party/py/lingvo:__subpackages__",
"//third_party/py/medical_research_foundations:__subpackages__",
"//tensorflow:__subpackages__",
"//waymo/ml/deploy/sync_test/tools:__subpackages__",
]
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = visibility + ["//tensorflow/tools/pip_package:__subpackages__"],
licenses = ["notice"],
)
exports_files(["tpu_test_wrapper.py"])
bzl_library(
name = "tpu_test_wrapper_bzl",
srcs = ["tpu_test_wrapper.bzl"],
)
py_strict_library(
name = "tpu_test_wrapper",
srcs = ["tpu_test_wrapper.py"],
srcs_version = "PY3",
deps = [
"//tensorflow/python/platform:flags",
"//tensorflow/python/util:tf_inspect",
],
)
py_strict_test(
name = "tpu_test_wrapper_test",
srcs = ["tpu_test_wrapper_test.py"],
main = "tpu_test_wrapper_test.py",
python_version = "PY3",
srcs_version = "PY3",
tags = [
"no_oss_py35",
"no_pip",
],
deps = [
":tpu_test_wrapper",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/platform:flags",
"@absl_py//absl/testing:flagsaver",
],
)
tf_py_strict_test(
name = "tpu_embedding_v3_cpu_ops_test",
srcs = ["tpu_embedding_v3_cpu_ops_test.py"],
shard_count = 8,
tags = ["no_oss"],
deps = [
# copybara:uncomment "//tensorflow/core/tpu/kernels",
"//tensorflow/python/compat:v2_compat",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:sparse_tensor",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:bitwise_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:sparse_ops",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/tpu/ops:gen_xla_ops",
"//third_party/py/numpy",
"@absl_py//absl/testing:parameterized",
],
)
alias(
name = "tpu_ops",
actual = "//tensorflow/python/tpu/ops",
)
pytype_strict_library(
name = "async_checkpoint",
srcs = ["async_checkpoint.py"],
srcs_version = "PY3",
deps = [
"//tensorflow/core:protos_all_py",
"//tensorflow/python/client:session",
"//tensorflow/python/framework:meta_graph",
"//tensorflow/python/framework:ops",
"//tensorflow/python/platform:tf_logging",
"//tensorflow/python/saved_model:pywrap_saved_model",
"//tensorflow/python/training",
"//tensorflow/python/training:basic_session_run_hooks",
"//tensorflow/python/training:monitored_session",
"//tensorflow/python/training:saver",
"//tensorflow/python/training:session_run_hook",
"//tensorflow/python/training:summary_io",
"//tensorflow/python/training:training_util",
],
)
tpu_py_strict_test(
name = "async_checkpoint_test",
size = "medium",
srcs = ["async_checkpoint_test.py"],
disable_experimental = True,
disable_mlir_bridge = False,
tags = ["no_oss"],
deps = [
":async_checkpoint",
":tpu_lib",
"//tensorflow/core:protos_all_py",
"//tensorflow/python/compat:v2_compat",
"//tensorflow/python/distribute/cluster_resolver:tpu_cluster_resolver_py",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/lib/io:file_io",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:metrics",
"//tensorflow/python/ops:variable_scope",
"//tensorflow/python/ops/losses",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/platform:flags",
"//tensorflow/python/platform:tf_logging",
"//tensorflow/python/saved_model:pywrap_saved_model",
"//tensorflow/python/training:basic_session_run_hooks",
"//tensorflow/python/training:training_lib",
"//third_party/py/numpy",
],
)
pytype_strict_library(
name = "device_assignment",
srcs = ["device_assignment.py"],
visibility = [
"//tensorflow:internal",
],
deps = [
":topology",
"//tensorflow/python/platform:tf_logging",
"//tensorflow/python/util:numpy_compat",
"//tensorflow/python/util:tf_export",
"//third_party/py/numpy",
],
)
pytype_strict_library(
name = "preempted_hook_py",
srcs = ["preempted_hook.py"],
srcs_version = "PY3",
deps = [
"//tensorflow/python/distribute/cluster_resolver:tpu_cluster_resolver_py",
"//tensorflow/python/platform:tf_logging",
"//tensorflow/python/training:session_run_hook",
],
)
py_strict_library(
name = "tpu_replication",
srcs = ["tpu_replication.py"],
deps = [
":device_assignment",
"//tensorflow/core:protos_all_py",
"//tensorflow/python/distribute:device_util",
"//tensorflow/python/distribute:distribute_lib",
"//tensorflow/python/framework:device",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:func_graph",
"//tensorflow/python/framework:ops",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:control_flow_ops",
"//tensorflow/python/ops:variables",
"//tensorflow/python/tpu/ops",
"//tensorflow/python/types:core",
"//tensorflow/python/util:compat",
"//tensorflow/python/util:tf_export",
"@absl_py//absl/logging",
],
)
py_strict_library(
name = "functional",
srcs = ["functional.py"],
srcs_version = "PY3",
visibility = [
"//visibility:public",
],
deps = [
"//tensorflow/python/tpu/ops",
],
)
pytype_strict_library(
name = "topology",
srcs = ["topology.py"],
srcs_version = "PY3",
deps = [
"//tensorflow/core/protobuf/tpu:topology_proto_py",
"//tensorflow/python/util:numpy_compat",
"//tensorflow/python/util:tf_export",
"//third_party/py/numpy",
],
)
py_strict_library(
name = "tpu",
srcs = [
"__init__.py",
],
srcs_version = "PY3",
deps = ["//tensorflow/python/framework:versions"],
)
py_strict_library(
name = "tpu_noestimator",
srcs = [
"__init__.py",
"api.py",
],
srcs_version = "PY3",
deps = [
":feature_column_v2",
":tpu_embedding_for_serving",
":tpu_embedding_v1",
":tpu_embedding_v2",
":tpu_embedding_v2_utils",
":tpu_embedding_v3",
":tpu_hardware_feature",
":tpu_lib",
":tpu_py",
"//tensorflow/python/framework:versions",
],
)
pytype_strict_library(
name = "tensor_tracer",
srcs = ["tensor_tracer.py"],
deps = [
":tensor_tracer_flags",
":tensor_tracer_report",
":tpu_replication",
"//tensorflow/core:protos_all_py",
"//tensorflow/python/eager:monitoring",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:func_graph",
"//tensorflow/python/framework:function",
"//tensorflow/python/framework:graph_io",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor",
"//tensorflow/python/framework:tensor_util",
"//tensorflow/python/lib/io:file_io",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:array_ops_stack",
"//tensorflow/python/ops:cond",
"//tensorflow/python/ops:control_flow_case",
"//tensorflow/python/ops:control_flow_ops",
"//tensorflow/python/ops:control_flow_util",
"//tensorflow/python/ops:init_ops",
"//tensorflow/python/ops:linalg_ops",
"//tensorflow/python/ops:logging_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:math_ops_gen",
"//tensorflow/python/ops:nn_impl",
"//tensorflow/python/ops:state_ops",
"//tensorflow/python/ops:string_ops",
"//tensorflow/python/ops:summary_ops_v2",
"//tensorflow/python/ops:variable_scope",
"//tensorflow/python/platform:analytics",
"//tensorflow/python/platform:gfile",
"//tensorflow/python/platform:remote_utils",
"//tensorflow/python/platform:tf_logging",
"//tensorflow/python/summary:summary_iterator",
"//tensorflow/python/tpu/ops",
"//tensorflow/python/training:training_util",
"//third_party/py/numpy",
],
)
pytype_strict_library(
name = "tensor_tracer_report",
srcs = ["tensor_tracer_report.py"],
deps = [
":tensor_tracer_proto_py",
"//tensorflow/python/platform:gfile",
"//tensorflow/python/platform:tf_logging",
],
)
pytype_strict_library(
name = "tensor_tracer_flags",
srcs = ["tensor_tracer_flags.py"],
deps = [
"//tensorflow/python/ops:linalg_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/platform:tf_logging",
"@absl_py//absl/flags",
],
)
pytype_strict_library(
name = "tpu_lib",
srcs = [
"__init__.py",
"bfloat16.py",
"tpu_optimizer.py",
"tpu_strategy_util.py",
"training_loop.py",
],
srcs_version = "PY3",
deps = [
":tensor_tracer",
":topology",
":tpu_feed",
":tpu_function",
":tpu_py",
"//tensorflow/core:protos_all_py",
"//tensorflow/python/client:session",
"//tensorflow/python/compiler/xla",
"//tensorflow/python/distribute/cluster_resolver:base_cluster_resolver_py",
"//tensorflow/python/eager:context",
"//tensorflow/python/eager:def_function",
"//tensorflow/python/eager:monitoring",
"//tensorflow/python/framework:device",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:versions",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:control_flow_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:variable_scope",
"//tensorflow/python/ops:while_loop",
"//tensorflow/python/ops/losses",
"//tensorflow/python/platform:tf_logging",
"//tensorflow/python/tpu/ops",
"//tensorflow/python/training:optimizer",
"//tensorflow/python/training:session_run_hook",
"//tensorflow/python/training:training_util",
"//tensorflow/python/types:core",
"//tensorflow/python/util:compat",
"//tensorflow/python/util:tf_contextlib",
"//tensorflow/python/util:tf_export",
],
)
pytype_strict_library(
name = "tpu_py",
srcs = ["tpu.py"],
deps = [
":device_assignment",
":tensor_tracer",
":tpu_feed",
":tpu_function",
":tpu_name_util",
":tpu_replication",
"//tensorflow/compiler/tf2xla/python:xla",
"//tensorflow/core:protos_all_py",
"//tensorflow/core/protobuf/tpu:compilation_result_proto_py",
"//tensorflow/core/protobuf/tpu:dynamic_padding_proto_py",
"//tensorflow/core/protobuf/tpu:tpu_embedding_configuration_proto_py",
"//tensorflow/python:tf2",
"//tensorflow/python/compiler/xla",
"//tensorflow/python/framework:auto_control_deps",
"//tensorflow/python/framework:composite_tensor",
"//tensorflow/python/framework:config",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:func_graph",
"//tensorflow/python/framework:function",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor_shape",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:array_ops_stack",
"//tensorflow/python/ops:cond",
"//tensorflow/python/ops:control_flow_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:variable_scope",
"//tensorflow/python/tpu/ops",
"//tensorflow/python/types:core",
"//tensorflow/python/util:compat",
"//tensorflow/python/util:nest",
"//tensorflow/python/util:object_identity",
"//tensorflow/python/util:tf_export",
"//tensorflow/python/util:traceback_utils",
"//tensorflow/python/util:variable_utils",
"//third_party/py/numpy",
"@absl_py//absl/logging",
],
)
pytype_strict_library(
name = "tpu_feed",
srcs = ["tpu_feed.py"],
deps = [
":tpu_name_util",
":tpu_sharding",
"//tensorflow/python/compiler/xla/experimental:xla_sharding",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor_shape",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/tpu/ops",
"//tensorflow/python/util:nest",
"//third_party/py/numpy",
],
)
pytype_strict_library(
name = "tpu_function",
srcs = ["tpu_function.py"],
)
pytype_strict_library(
name = "tpu_sharding",
srcs = ["tpu_sharding.py"],
deps = [
"//tensorflow/python/framework:tensor_shape",
],
)
pytype_strict_library(
name = "tpu_system_metadata",
srcs = ["tpu_system_metadata.py"],
deps = [
":tpu_py",
"//tensorflow/core:protos_all_py",
"//tensorflow/python/client:session",
"//tensorflow/python/distribute:device_util",
"//tensorflow/python/eager:context",
"//tensorflow/python/framework:config",
"//tensorflow/python/framework:device",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:ops",
"//tensorflow/python/platform:tf_logging",
"//tensorflow/python/util:tf_export",
],
)
pytype_strict_library(
name = "datasets",
srcs = [
"datasets.py",
],
srcs_version = "PY3",
visibility = visibility,
deps = [
"//tensorflow/python/data/experimental/ops:interleave_ops",
"//tensorflow/python/data/ops:dataset_ops",
"//tensorflow/python/data/ops:iterator_ops",
"//tensorflow/python/data/ops:readers",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:function",
"//tensorflow/python/framework:ops",
"//tensorflow/python/ops:functional_ops",
"//tensorflow/python/types:data",
],
)
tf_py_strict_test(
name = "datasets_test",
size = "medium",
srcs = ["datasets_test.py"],
grpc_enabled = True,
shard_count = 4,
tags = ["no_oss"],
deps = [
":datasets",
"//tensorflow/core:protos_all_py",
"//tensorflow/python/client:session",
"//tensorflow/python/data/ops:dataset_ops",
"//tensorflow/python/data/ops:readers",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor_shape",
"//tensorflow/python/lib/io:python_io",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/training:server_lib",
"//tensorflow/python/util:compat",
],
)
tf_py_strict_test(
name = "tpu_test",
size = "small",
srcs = ["tpu_test.py"],
tags = [
"no_oss", # TODO(b/131157871): Reenable in OSS when fixed
"no_windows", # TODO: needs investigation on Windows
],
deps = [
":tpu_feed",
":tpu_lib",
":tpu_py",
":tpu_replication",
"//tensorflow/python/eager:def_function",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:importer",
"//tensorflow/python/framework:ops",
"//tensorflow/python/layers",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:control_flow_ops",
"//tensorflow/python/ops:control_flow_util",
"//tensorflow/python/ops:init_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:special_math_ops",
"//tensorflow/python/ops:variable_scope",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/tpu/ops",
],
)
tf_py_strict_test(
name = "tpu_sharding_test",
size = "small",
srcs = ["tpu_sharding_test.py"],
deps = [
":tpu_sharding",
"//tensorflow/python/framework:tensor_shape",
"//tensorflow/python/platform:client_testlib",
],
)
tf_py_strict_test(
name = "bfloat16_test",
size = "small",
srcs = ["bfloat16_test.py"],
deps = [
":tpu_lib",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:variable_scope",
"//tensorflow/python/ops:variables",
"//tensorflow/python/platform:client_testlib",
],
)
tf_py_strict_test(
name = "tpu_infeed_test",
size = "small",
srcs = ["tpu_infeed_test.py"],
deps = [
":tpu_feed",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/platform:client_testlib",
],
)
tf_py_strict_test(
name = "topology_test",
size = "medium",
srcs = ["topology_test.py"],
deps = [
":topology",
"//tensorflow/python/platform:client_testlib",
],
)
pytype_strict_library(
name = "tpu_strategy_util",
srcs = ["tpu_strategy_util.py"],
visibility = [
"//tensorflow:__subpackages__",
"//third_party/py/tensorflow_numerics/extensions:__pkg__",
],
deps = [
":topology",
":tpu_py",
"//tensorflow/core:protos_all_py",
"//tensorflow/python/client:session",
"//tensorflow/python/distribute/cluster_resolver:base_cluster_resolver_py",
"//tensorflow/python/eager:context",
"//tensorflow/python/eager:def_function",
"//tensorflow/python/eager:monitoring",
"//tensorflow/python/framework:device",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:ops",
"//tensorflow/python/platform:tf_logging",
"//tensorflow/python/util:compat",
],
)
pytype_strict_library(
name = "tpu_hardware_feature",
srcs = ["tpu_hardware_feature.py"],
deps = [
"//tensorflow/core/protobuf/tpu:topology_proto_py",
"//tensorflow/python/util:tf_export",
],
)
py_strict_library(
name = "tpu_name_util",
srcs = ["tpu_name_util.py"],
srcs_version = "PY3",
deps = [
"//tensorflow/python/util:tf_export",
],
)
pytype_strict_library(
name = "feature_column",
srcs = ["feature_column.py"],
deps = [
":tpu_function",
":tpu_py",
":tpu_replication",
"//tensorflow/python/feature_column",
"//tensorflow/python/feature_column:feature_column_py",
"//tensorflow/python/framework:ops",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:init_ops",
"//tensorflow/python/ops:variable_scope",
],
)
pytype_strict_library(
name = "feature_column_v2",
srcs = ["feature_column_v2.py"],
deps = [
":feature_column",
":tpu_py",
":tpu_replication",
"//tensorflow/python/feature_column",
"//tensorflow/python/feature_column:feature_column_py",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:embedding_ops",
"//tensorflow/python/ops:init_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:sparse_ops",
"//tensorflow/python/ops:variable_scope",
"//tensorflow/python/util:tf_export",
],
)
tf_py_strict_test(
name = "feature_column_test",
srcs = [
"feature_column_test.py",
],
main = "feature_column_test.py",
deps = [
":feature_column",
"//tensorflow/python/client:session",
"//tensorflow/python/feature_column",
"//tensorflow/python/feature_column:feature_column_py",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:sparse_tensor",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:lookup_ops",
"//tensorflow/python/ops:parsing_ops",
"//tensorflow/python/ops:variables",
"//tensorflow/python/platform:client_testlib",
"//third_party/py/numpy",
],
)
tf_py_strict_test(
name = "feature_column_v2_test",
srcs = [
"feature_column_v2_test.py",
],
main = "feature_column_v2_test.py",
tags = ["no_oss"], # Due to the usage of keras component.
deps = [
":feature_column_v2",
":tpu_function",
":tpu_py",
"//tensorflow/python/client:session",
"//tensorflow/python/feature_column:feature_column_py",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:sparse_tensor",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:init_ops",
"//tensorflow/python/ops:lookup_ops",
"//tensorflow/python/ops:parsing_ops",
"//tensorflow/python/ops:variable_scope",
"//tensorflow/python/ops:variables",
"//tensorflow/python/platform:client_testlib",
"@absl_py//absl/testing:parameterized",
],
)
pytype_strict_library(
name = "tpu_embedding_v2_utils",
srcs = ["tpu_embedding_v2_utils.py"],
srcs_version = "PY3",
deps = [
"//tensorflow/core/protobuf/tpu:optimization_parameters_proto_py",
"//tensorflow/core/protobuf/tpu:tpu_embedding_configuration_proto_py",
"//tensorflow/python/distribute:device_util",
"//tensorflow/python/distribute:sharded_variable",
"//tensorflow/python/distribute:tpu_strategy",
"//tensorflow/python/framework:device_spec",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor_shape",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:init_ops_v2",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:variables",
"//tensorflow/python/tpu/ops",
"//tensorflow/python/types:core",
"//tensorflow/python/util:tf_export",
"@absl_py//absl/logging",
],
)
pytype_strict_library(
name = "tpu_embedding_v2",
srcs = ["tpu_embedding_v2.py"],
srcs_version = "PY3",
deps = [
":tpu_embedding_v2_utils",
":tpu_py",
":tpu_replication",
"//tensorflow/core:protos_all_py",
"//tensorflow/core/protobuf/tpu:tpu_embedding_configuration_proto_py",
"//tensorflow/python/distribute:device_util",
"//tensorflow/python/distribute:distribute_lib",
"//tensorflow/python/distribute:distribute_utils",
"//tensorflow/python/distribute:sharded_variable",
"//tensorflow/python/distribute:tpu_strategy",
"//tensorflow/python/eager:context",
"//tensorflow/python/eager:def_function",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:device",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:sparse_tensor",
"//tensorflow/python/framework:tensor",
"//tensorflow/python/framework:tensor_shape",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:sparse_ops",
"//tensorflow/python/ops:variable_scope",
"//tensorflow/python/ops:variables",
"//tensorflow/python/ops/ragged:ragged_tensor",
"//tensorflow/python/saved_model:save_context",
"//tensorflow/python/saved_model/registration",
"//tensorflow/python/tpu/ops",
"//tensorflow/python/trackable:autotrackable",
"//tensorflow/python/trackable:base",
"//tensorflow/python/types:internal",
"//tensorflow/python/util:compat",
"//tensorflow/python/util:nest",
"//tensorflow/python/util:tf_export",
"//tensorflow/python/util:tf_inspect",
"@absl_py//absl/logging",
],
)
pytype_strict_library(
name = "tpu_embedding_base",
srcs = ["tpu_embedding_base.py"],
srcs_version = "PY3",
deps = [
":tpu_embedding_v2_utils",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/ops:variables",
"//tensorflow/python/trackable:autotrackable",
"//tensorflow/python/util:nest",
],
)
pytype_strict_library(
name = "tpu_embedding_v3_utils",
srcs = ["tpu_embedding_v3_utils.py"],
deps = [
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:tensor",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:control_flow_ops_gen",
"//tensorflow/python/ops:manip_ops",
"//tensorflow/python/ops:variables",
"//tensorflow/python/trackable:base",
],
)
tf_py_strict_test(
name = "tpu_embedding_v3_utils_test",
srcs = ["tpu_embedding_v3_utils_test.py"],
deps = [
":tpu_embedding_v3_utils",
"//tensorflow/python/compat:v2_compat",
"//tensorflow/python/eager:test",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/platform:client_testlib",
"@absl_py//absl/testing:parameterized",
],
)
pytype_strict_library(
name = "tpu_embedding_v3_checkpoint_adapter",
srcs = ["tpu_embedding_v3_checkpoint_adapter.py"],
deps = [
":tpu_embedding_v3_utils",
"//tensorflow/core/tpu/kernels:sparse_core_layout_proto_py",
"//tensorflow/python/checkpoint:checkpoint_adapter",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:tensor",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:manip_ops",
"//tensorflow/python/ops:variables",
"//tensorflow/python/trackable:base",
"//tensorflow/python/training:py_checkpoint_reader",
"//tensorflow/python/util/protobuf",
"@absl_py//absl/logging",
],
)
tf_py_strict_test(
name = "tpu_embedding_v3_checkpoint_adapter_test",
srcs = ["tpu_embedding_v3_checkpoint_adapter_test.py"],
deps = [
":tpu_embedding_v3_checkpoint_adapter",
"//tensorflow/core/tpu/kernels:sparse_core_layout_proto_py",
"//tensorflow/python/compat:v2_compat",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/platform:client_testlib",
"@absl_py//absl/logging",
],
)
pytype_strict_library(
name = "tpu_embedding_for_serving",
srcs = ["tpu_embedding_for_serving.py"],
srcs_version = "PY3",
deps = [
":tpu_embedding_base",
":tpu_embedding_v2_utils",
":tpu_embedding_v3_utils",
"//tensorflow/core/tpu/kernels:sparse_core_layout_proto_py",
"//tensorflow/python/distribute:distribute_lib",
"//tensorflow/python/distribute:tpu_strategy",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:sparse_tensor",
"//tensorflow/python/framework:tensor",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:array_ops_stack",
"//tensorflow/python/ops:embedding_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:sparse_ops",
"//tensorflow/python/ops:variables",
"//tensorflow/python/ops/ragged:ragged_tensor",
"//tensorflow/python/types:core",
"//tensorflow/python/util:nest",
"//tensorflow/python/util:tf_export",
"@absl_py//absl/logging",
],
)
tf_py_strict_test(
name = "tpu_embedding_for_serving_test",
srcs = [
"tpu_embedding_for_serving_test.py",
],
python_version = "PY3",
srcs_version = "PY3",
deps = [
":tpu_embedding_for_serving",
":tpu_embedding_v2_utils",
"//tensorflow/python/compat:v2_compat",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:sparse_tensor",
"//tensorflow/python/ops:init_ops_v2",
"//tensorflow/python/ops/ragged:ragged_tensor",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/util:nest",
"//third_party/py/numpy",
],
)
pytype_strict_library(
name = "tpu_embedding_v1",
srcs = ["tpu_embedding_v1.py"],
srcs_version = "PY3",
deps = [
":tpu_embedding_base",
":tpu_embedding_v2_utils",
":tpu_replication",
"//tensorflow/python/distribute:distribute_lib",
"//tensorflow/python/distribute:tpu_strategy",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:sparse_tensor",
"//tensorflow/python/framework:tensor",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:embedding_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:sparse_ops",
"//tensorflow/python/ops:variables",
"//tensorflow/python/ops/ragged:ragged_tensor",
"//tensorflow/python/util:nest",
"//tensorflow/python/util:tf_export",
],
)
tf_py_strict_test(
name = "tpu_embedding_v2_utils_test",
srcs = [
"tpu_embedding_v2_utils_test.py",
],
python_version = "PY3",
srcs_version = "PY3",
deps = [
":tpu_embedding_v2_utils",
"//tensorflow/core/protobuf/tpu:tpu_embedding_configuration_proto_py",
"//tensorflow/python/compat:v2_compat",
"//tensorflow/python/platform:client_testlib",
"@absl_py//absl/testing:parameterized",
],
)
tpu_py_strict_test(
name = "tpu_outside_compilation_test",
srcs = [
"tpu_outside_compilation_test.py",
],
disable_experimental = True,
disable_mlir_bridge = False,
python_version = "PY3",
deps = [
":functional",
":tpu_py",
":tpu_replication",
"//tensorflow/core:protos_all_py",
"//tensorflow/python/distribute:tpu_strategy",
"//tensorflow/python/distribute/cluster_resolver:tpu_cluster_resolver_py",
"//tensorflow/python/eager:def_function",
"//tensorflow/python/eager:remote",
"//tensorflow/python/eager:test",
"//tensorflow/python/framework:config",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor",
"//tensorflow/python/lib/io:tf_record",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:cond",
"//tensorflow/python/ops:gradients_impl",
"//tensorflow/python/ops:image_ops",
"//tensorflow/python/ops:logging_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:random_ops",
"//tensorflow/python/ops:string_ops",
"//tensorflow/python/ops:summary_ops_v2",
"//tensorflow/python/ops:tensor_array_ops",
"//tensorflow/python/ops:while_loop",
"//tensorflow/python/platform:flags",
"//tensorflow/python/platform:gfile",
"//tensorflow/python/tpu/ops",
"//third_party/py/numpy",
"@absl_py//absl/testing:parameterized",
],
)
# NOTE this target should only be depended on by the tpu_test_wrapper macro.
py_strict_library(
name = "tpu_test_deps",
srcs_version = "PY3",
visibility = ["//visibility:public"],
deps = [
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/platform:flags",
"//tensorflow/python/util:tf_inspect",
],
)
tf_proto_library(
name = "tensor_tracer_proto",
srcs = ["tensor_tracer.proto"],
protodeps = [
"//tensorflow/core:protos_all",
],
visibility = ["//visibility:public"],
)
# copybara:uncomment_begin(google-only)
# py_proto_library(
# name = "tensor_tracer_py_pb2",
# visibility = ["//visibility:public"],
# deps = [":tensor_tracer_proto"],
# )
# copybara:uncomment_end
tf_python_pybind_extension(
name = "_pywrap_sparse_core_layout",
srcs = ["pywrap_sparse_core_layout.cc"],
enable_stub_generation = True,
features = ["-layering_check"],
pytype_srcs = [
"_pywrap_sparse_core_layout.pyi",
],
deps = [
"//tensorflow/python/lib/core:pybind11_lib",
"//tensorflow/python/lib/core:pybind11_status",
"//tensorflow/python/lib/core:pybind11_status_headers",
"//third_party/python_runtime:headers", # buildcleaner: keep
"@pybind11",
"@pybind11_abseil//pybind11_abseil:status_casters",
"@pybind11_protobuf//pybind11_protobuf:native_proto_caster",
] + if_pywrap(
if_false = [
"//tensorflow/core/tpu/kernels:_pywrap_sparse_core_layout_header_only",
],
if_true = [
"//tensorflow/core/tpu/kernels:sparse_core_layout",
],
),
)
tf_python_pybind_extension(
name = "_pywrap_tpu_embedding",
srcs = ["pywrap_tpu_embedding.cc"],
enable_stub_generation = True,
features = ["-layering_check"],
pytype_srcs = [
"_pywrap_tpu_embedding.pyi",
],
deps = [
"//tensorflow/core/tpu/kernels:sparse_core_ops_utils",
"//tensorflow/python/lib/core:pybind11_lib",
"//third_party/python_runtime:headers", # buildcleaner: keep
"@pybind11",
],
)
py_strict_library(
name = "tpu_embedding_v3",
srcs = ["tpu_embedding_v3.py"],
deps = [
":_pywrap_sparse_core_layout",
":_pywrap_tpu_embedding",
":tpu_embedding_base",
":tpu_embedding_v2_utils",
":tpu_embedding_v3_checkpoint_adapter",
":tpu_embedding_v3_utils",
":tpu_replication",
"//tensorflow/core:protos_all_py",
"//tensorflow/core/tpu/kernels:sparse_core_layout_proto_py",
"//tensorflow/python/checkpoint:saveable_compat",
"//tensorflow/python/distribute:device_util",
"//tensorflow/python/distribute:distribute_lib",
"//tensorflow/python/distribute:tpu_strategy",
"//tensorflow/python/distribute:tpu_util",
"//tensorflow/python/distribute:tpu_values",
"//tensorflow/python/distribute:values",
"//tensorflow/python/distribute:values_util",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:device",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:sparse_tensor",
"//tensorflow/python/framework:tensor",
"//tensorflow/python/framework:tensor_shape",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:control_flow_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:resource_variable_ops_gen",
"//tensorflow/python/ops:variable_scope",
"//tensorflow/python/ops:variables",
"//tensorflow/python/ops/ragged:ragged_tensor",
"//tensorflow/python/tpu/ops:gen_xla_ops",
"//tensorflow/python/trackable:base",
"//tensorflow/python/training/saving:saveable_object",
"//tensorflow/python/util:compat",
"//tensorflow/python/util:nest",
"//tensorflow/python/util:tf_export",
"//tensorflow/python/util:tf_inspect",
"@absl_py//absl/logging",
],
)
tpu_py_strict_test(
name = "tpu_embedding_v3_test",
srcs = ["tpu_embedding_v3_test.py"],
disable_experimental = True,
disable_tfrt = False,
disable_v2 = True,
disable_v3 = True,
# copybara:uncomment disable_v4i = True,
# copybara:uncomment disable_v5 = False,
# copybara:uncomment disable_v5_grm = True,
# copybara:uncomment disable_v6e = True,
python_version = "PY3",
shard_count = 4,
srcs_version = "PY3",
# copybara:uncomment use_tpu_plugin_and_capi = True,
deps = [
":tpu_embedding_for_serving",
":tpu_embedding_v2_utils",
":tpu_embedding_v3",
":tpu_replication",
"//tensorflow/python/compat:v2_compat",
"//tensorflow/python/data/ops:dataset_ops",
"//tensorflow/python/distribute:distribute_lib",
"//tensorflow/python/distribute:tpu_strategy",
"//tensorflow/python/distribute/cluster_resolver:tpu_cluster_resolver_py",
"//tensorflow/python/eager:def_function",
"//tensorflow/python/eager:remote",
"//tensorflow/python/framework:config",
"//tensorflow/python/framework:sparse_tensor",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:init_ops_v2",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:sparse_ops",
"//tensorflow/python/ops:variables",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/util:nest",
"//third_party/py/numpy",
"@absl_py//absl/testing:parameterized",
],
)
tpu_py_strict_test(
name = "tpu_embedding_v3_additional_test",
srcs = ["tpu_embedding_v3_additional_test.py"],
args = [
"--tpu=''",
],
disable_experimental = True,
disable_tfrt = False,
disable_v2 = True,
disable_v3 = True,
# copybara:uncomment disable_v4i = True,
# copybara:uncomment disable_v5 = False,
# copybara:uncomment disable_v5_grm = True,
env = {
"TF_XLA_FLAGS": (
"--tf_mlir_enable_mlir_bridge=true " +
"--tf_mlir_enable_convert_control_to_data_outputs_pass=true " +
"--tf_mlir_enable_merge_control_flow_pass=true " +
"--tf_mlir_enable_strict_clusters=true " +
# These tests run lookup without while loops.
"--tf_xla_disable_full_embedding_pipelining=true"
),
},
python_version = "PY3",
# Make it so there are enough shards so each test runs in separate shards.
shard_count = 12,
srcs_version = "PY3",
visibility = ["//visibility:private"],
deps = [
":device_assignment",
":tpu_embedding_for_serving",
":tpu_embedding_v2_utils",
":tpu_embedding_v3",
"//tensorflow/python/checkpoint",
"//tensorflow/python/compat:v2_compat",
"//tensorflow/python/distribute:distribute_lib",
"//tensorflow/python/distribute:tpu_strategy",
"//tensorflow/python/distribute:values",
"//tensorflow/python/distribute/cluster_resolver:tpu_cluster_resolver_py",
"//tensorflow/python/eager:context",
"//tensorflow/python/eager:def_function",
"//tensorflow/python/eager:remote",
"//tensorflow/python/framework:combinations",
"//tensorflow/python/framework:config",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:sparse_tensor",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/platform:client_testlib",
"//third_party/py/numpy",
"@absl_py//absl/flags",
"@absl_py//absl/logging",
"@absl_py//absl/testing:parameterized",
],
)
internal_create_sanitizer_settings()
Reference in New Issue View Git Blame Copy Permalink