708 lines
24 KiB
Python
708 lines
24 KiB
Python
load("@rules_cc//cc:cc_library.bzl", "cc_library")
|
|
load("@xla//third_party/rules_python/python:py_binary.bzl", "py_binary")
|
|
load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
|
|
load("//tensorflow:tensorflow.bzl", "if_not_windows")
|
|
load("//tensorflow:tensorflow.default.bzl", "cuda_py_strict_test", "get_compatible_with_portable", "tf_py_strict_test", "tf_pybind_cc_library_wrapper", "tf_python_pybind_extension")
|
|
load("//tensorflow/core/platform:build_config.bzl", "tf_protos_grappler")
|
|
load("//tensorflow/core/platform:build_config_root.bzl", "if_pywrap")
|
|
|
|
package(
|
|
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
|
|
default_visibility = ["//tensorflow:internal"],
|
|
licenses = ["notice"],
|
|
)
|
|
|
|
# TODO(gunan): Investigate making this action hermetic so we do not need
|
|
# to run it locally.
|
|
cc_library(
|
|
name = "cost_analyzer_lib",
|
|
srcs = ["cost_analyzer.cc"],
|
|
hdrs = ["cost_analyzer.h"],
|
|
compatible_with = get_compatible_with_portable(),
|
|
deps = [
|
|
"//tensorflow/core:lib",
|
|
"//tensorflow/core:protos_all_cc",
|
|
"//tensorflow/core/grappler:grappler_item",
|
|
"//tensorflow/core/grappler/clusters:cluster",
|
|
"//tensorflow/core/grappler/costs:analytical_cost_estimator",
|
|
"//tensorflow/core/grappler/costs:cost_estimator",
|
|
"//tensorflow/core/grappler/costs:measuring_cost_estimator",
|
|
"//tensorflow/core/grappler/costs:utils",
|
|
"@com_google_absl//absl/log",
|
|
"@com_google_absl//absl/status",
|
|
] + tf_protos_grappler(),
|
|
alwayslink = 1,
|
|
)
|
|
|
|
# Necessary for the pywrap inclusion below. Combining targets does not work
|
|
# properly.
|
|
tf_pybind_cc_library_wrapper(
|
|
name = "cost_analyzer_headers",
|
|
deps = [
|
|
":cost_analyzer_lib",
|
|
],
|
|
)
|
|
|
|
tf_python_pybind_extension(
|
|
name = "_pywrap_cost_analyzer",
|
|
srcs = ["cost_analyzer_wrapper.cc"],
|
|
hdrs = [
|
|
"cost_analyzer.h",
|
|
"//tensorflow/cc:pywrap_required_hdrs",
|
|
"//tensorflow/core/grappler:pywrap_required_hdrs",
|
|
"//tensorflow/core/grappler/clusters:pywrap_required_hdrs",
|
|
"//tensorflow/core/grappler/costs:pywrap_required_hdrs",
|
|
"//tensorflow/core/public:session.h",
|
|
"//tensorflow/core/public:session_options.h",
|
|
],
|
|
enable_stub_generation = True,
|
|
pytype_srcs = [
|
|
"_pywrap_cost_analyzer.pyi",
|
|
],
|
|
starlark_only = True,
|
|
deps = [
|
|
":cost_analyzer_headers",
|
|
"//tensorflow/core:core_cpu_base",
|
|
"//tensorflow/core:framework",
|
|
"//tensorflow/core:framework_headers_lib",
|
|
"//tensorflow/core:framework_lite",
|
|
"//tensorflow/core:lib",
|
|
"//tensorflow/core:lib_headers_for_pybind",
|
|
"//tensorflow/core:lib_proto_parsing",
|
|
"//tensorflow/core:portable_gif_internal",
|
|
"//tensorflow/core:protos_all_cc",
|
|
"//tensorflow/core/common_runtime:core_cpu_headers_lib",
|
|
"//tensorflow/core/common_runtime:device_set",
|
|
"//tensorflow/core/common_runtime/gpu:gpu_id",
|
|
"//tensorflow/core/framework:allocator",
|
|
"//tensorflow/core/framework:tensor",
|
|
"//tensorflow/core/platform:threadpool_options",
|
|
"//tensorflow/python/lib/core:pybind11_status",
|
|
"@com_google_absl//absl/container:flat_hash_map",
|
|
"@com_google_absl//absl/container:flat_hash_set",
|
|
"@com_google_absl//absl/status",
|
|
"@com_google_absl//absl/status:statusor",
|
|
"@pybind11",
|
|
"@tsl//tsl/platform:platform_port",
|
|
"@tsl//tsl/platform:thread_annotations",
|
|
"@xla//xla/tsl/protobuf:error_codes_proto_impl_cc",
|
|
],
|
|
)
|
|
|
|
cc_library(
|
|
name = "model_analyzer_lib",
|
|
srcs = ["model_analyzer.cc"],
|
|
hdrs = ["model_analyzer.h"],
|
|
deps = [
|
|
"//tensorflow/core:framework",
|
|
"//tensorflow/core:lib",
|
|
"//tensorflow/core:protos_all_cc",
|
|
"//tensorflow/core/grappler:grappler_item",
|
|
"//tensorflow/core/grappler/costs:graph_properties",
|
|
"@com_google_absl//absl/status",
|
|
],
|
|
)
|
|
|
|
tf_python_pybind_extension(
|
|
name = "_pywrap_model_analyzer",
|
|
srcs = ["model_analyzer_wrapper.cc"],
|
|
hdrs = [
|
|
"model_analyzer.h",
|
|
"//tensorflow/core/grappler:pywrap_required_hdrs",
|
|
],
|
|
enable_stub_generation = True,
|
|
pytype_srcs = [
|
|
"_pywrap_model_analyzer.pyi",
|
|
],
|
|
starlark_only = True,
|
|
deps = [
|
|
"//tensorflow/core:framework_headers_lib",
|
|
"//tensorflow/core:lib",
|
|
"//tensorflow/core:lib_headers_for_pybind",
|
|
"//tensorflow/core:lib_proto_parsing",
|
|
"//tensorflow/core:portable_gif_internal",
|
|
"//tensorflow/core:protos_all_cc",
|
|
"//tensorflow/core/framework:tensor",
|
|
"//tensorflow/python/lib/core:pybind11_status",
|
|
"@com_google_absl//absl/status",
|
|
"@pybind11",
|
|
"@tsl//tsl/platform:platform_port",
|
|
] + if_pywrap(["//tensorflow/python/grappler:model_analyzer_lib"]),
|
|
)
|
|
|
|
py_library(
|
|
name = "tf_item",
|
|
srcs = ["item.py"],
|
|
strict_deps = True,
|
|
visibility = ["//visibility:public"],
|
|
deps = [
|
|
":_pywrap_tf_item",
|
|
"//tensorflow/core:protos_all_py",
|
|
"//tensorflow/core/grappler/costs:op_performance_data_py",
|
|
],
|
|
)
|
|
|
|
tf_python_pybind_extension(
|
|
name = "_pywrap_tf_item",
|
|
srcs = ["item_wrapper.cc"],
|
|
hdrs = [
|
|
"//tensorflow/cc:pywrap_required_hdrs",
|
|
"//tensorflow/core/grappler:pywrap_required_hdrs",
|
|
"//tensorflow/core/grappler/clusters:pywrap_required_hdrs",
|
|
"//tensorflow/core/grappler/costs:pywrap_required_hdrs",
|
|
"//tensorflow/core/grappler/utils:pywrap_required_hdrs",
|
|
],
|
|
enable_stub_generation = True,
|
|
pytype_srcs = [
|
|
"_pywrap_tf_item.pyi",
|
|
],
|
|
deps = [
|
|
"//tensorflow/core:core_cpu_base",
|
|
"//tensorflow/core:framework",
|
|
"//tensorflow/core:framework_headers_lib",
|
|
"//tensorflow/core:framework_lite",
|
|
"//tensorflow/core:lib",
|
|
"//tensorflow/core:lib_proto_parsing",
|
|
"//tensorflow/core:portable_gif_internal",
|
|
"//tensorflow/core:protos_all_cc",
|
|
"//tensorflow/core:session_options",
|
|
"//tensorflow/core/common_runtime:core_cpu_headers_lib",
|
|
"//tensorflow/core/common_runtime:device_set",
|
|
"//tensorflow/core/common_runtime/gpu:gpu_id",
|
|
"//tensorflow/core/framework:allocator",
|
|
"//tensorflow/core/framework:tensor",
|
|
"//tensorflow/core/grappler:utils",
|
|
"//tensorflow/python/lib/core:pybind11_status",
|
|
"@com_google_absl//absl/container:flat_hash_map",
|
|
"@com_google_absl//absl/container:flat_hash_set",
|
|
"@com_google_absl//absl/status",
|
|
"@com_google_absl//absl/status:statusor",
|
|
"@com_google_absl//absl/types:span",
|
|
"@pybind11",
|
|
"@tsl//tsl/platform:platform_port",
|
|
"@tsl//tsl/platform:thread_annotations",
|
|
"@xla//xla/tsl/protobuf:error_codes_proto_impl_cc",
|
|
] + if_not_windows(["//tensorflow/core/grappler/costs:graph_properties"]), # b/148556093,
|
|
)
|
|
|
|
tf_py_strict_test(
|
|
name = "item_test",
|
|
size = "small",
|
|
srcs = ["item_test.py"],
|
|
tags = [
|
|
"grappler",
|
|
"no_pip", # tf_optimizer is not available in pip.
|
|
],
|
|
deps = [
|
|
":tf_item",
|
|
"//tensorflow/python/framework:constant_op",
|
|
"//tensorflow/python/framework:dtypes",
|
|
"//tensorflow/python/framework:errors",
|
|
"//tensorflow/python/framework:meta_graph",
|
|
"//tensorflow/python/framework:ops",
|
|
"//tensorflow/python/framework:tensor_shape",
|
|
"//tensorflow/python/framework:test_lib",
|
|
"//tensorflow/python/ops:array_ops_gen",
|
|
"//tensorflow/python/ops:control_flow_ops",
|
|
"//tensorflow/python/ops:state_ops",
|
|
"//tensorflow/python/ops:variable_v1",
|
|
"//tensorflow/python/platform:client_testlib",
|
|
],
|
|
)
|
|
|
|
tf_py_strict_test(
|
|
name = "datasets_test",
|
|
size = "medium",
|
|
srcs = ["datasets_test.py"],
|
|
tags = [
|
|
"grappler",
|
|
"no_pip", # tf_optimizer is not available in pip.
|
|
],
|
|
deps = [
|
|
":tf_item",
|
|
"//tensorflow/python/data/ops:dataset_ops",
|
|
"//tensorflow/python/data/ops:iterator_ops",
|
|
"//tensorflow/python/framework:dtypes",
|
|
"//tensorflow/python/framework:meta_graph",
|
|
"//tensorflow/python/framework:ops",
|
|
"//tensorflow/python/framework:tensor_shape",
|
|
"//tensorflow/python/ops:array_ops",
|
|
"//tensorflow/python/platform:client_testlib",
|
|
"//third_party/py/numpy",
|
|
],
|
|
)
|
|
|
|
py_library(
|
|
name = "tf_cluster",
|
|
srcs = ["cluster.py"],
|
|
strict_deps = True,
|
|
visibility = ["//visibility:public"],
|
|
deps = [
|
|
":_pywrap_tf_cluster",
|
|
"//tensorflow/core:protos_all_py",
|
|
"//tensorflow/core/grappler/costs:op_performance_data_py",
|
|
],
|
|
)
|
|
|
|
tf_python_pybind_extension(
|
|
name = "_pywrap_tf_cluster",
|
|
srcs = ["cluster_wrapper.cc"],
|
|
hdrs = [
|
|
"//tensorflow/cc:pywrap_required_hdrs",
|
|
] + if_pywrap(
|
|
if_false = [
|
|
"//tensorflow/core/grappler:pywrap_required_hdrs",
|
|
"//tensorflow/core/grappler/clusters:pywrap_required_hdrs",
|
|
"//tensorflow/core/grappler/costs:pywrap_required_hdrs",
|
|
"//tensorflow/core/grappler/utils:pywrap_required_hdrs",
|
|
],
|
|
),
|
|
enable_stub_generation = True,
|
|
pytype_srcs = [
|
|
"_pywrap_tf_cluster.pyi",
|
|
],
|
|
deps = [
|
|
"//tensorflow/core:core_cpu_base",
|
|
"//tensorflow/core:framework",
|
|
"//tensorflow/core:framework_headers_lib",
|
|
"//tensorflow/core:framework_lite",
|
|
"//tensorflow/core:lib",
|
|
"//tensorflow/core:lib_headers_for_pybind",
|
|
"//tensorflow/core:lib_proto_parsing",
|
|
"//tensorflow/core:portable_gif_internal",
|
|
"//tensorflow/core:protos_all_cc",
|
|
"//tensorflow/core:session_options",
|
|
"//tensorflow/core/common_runtime:core_cpu_headers_lib",
|
|
"//tensorflow/core/common_runtime:device_set",
|
|
"//tensorflow/core/common_runtime/gpu:gpu_id",
|
|
"//tensorflow/core/framework:allocator",
|
|
"//tensorflow/core/framework:tensor",
|
|
"//tensorflow/core/grappler:utils",
|
|
"//tensorflow/core/platform:status",
|
|
"//tensorflow/python/lib/core:pybind11_status",
|
|
"@com_google_absl//absl/container:flat_hash_map",
|
|
"@com_google_absl//absl/container:flat_hash_set",
|
|
"@com_google_absl//absl/log:check",
|
|
"@com_google_absl//absl/status",
|
|
"@com_google_absl//absl/status:statusor",
|
|
"@com_google_absl//absl/types:span",
|
|
"@pybind11",
|
|
"@tsl//tsl/platform:platform_port",
|
|
"@tsl//tsl/platform:thread_annotations",
|
|
"@xla//xla/tsl/protobuf:error_codes_proto_impl_cc",
|
|
] + if_pywrap(
|
|
if_true = [
|
|
"//tensorflow/core/grappler/costs:measuring_cost_estimator",
|
|
"//tensorflow/core/grappler/clusters:single_machine",
|
|
],
|
|
),
|
|
)
|
|
|
|
cuda_py_strict_test(
|
|
name = "cluster_test",
|
|
size = "small",
|
|
srcs = ["cluster_test.py"],
|
|
shard_count = 5,
|
|
tags = [
|
|
"grappler",
|
|
"no_pip", # tf_optimizer is not available in pip.
|
|
"no_windows", # b/173520599
|
|
"notap", # TODO(b/135924227): Re-enable after fixing flakiness.
|
|
],
|
|
# This test will not run on XLA because it primarily tests the TF Classic flow.
|
|
xla_enable_strict_auto_jit = False,
|
|
deps = [
|
|
":tf_cluster",
|
|
":tf_item",
|
|
"//tensorflow/core:protos_all_py",
|
|
"//tensorflow/python/framework:meta_graph",
|
|
"//tensorflow/python/framework:ops",
|
|
"//tensorflow/python/ops:array_ops",
|
|
"//tensorflow/python/ops:math_ops",
|
|
"//tensorflow/python/ops:random_ops",
|
|
"//tensorflow/python/platform:client_testlib",
|
|
],
|
|
)
|
|
|
|
py_library(
|
|
name = "tf_optimizer",
|
|
srcs = ["tf_optimizer.py"],
|
|
strict_deps = True,
|
|
visibility = ["//visibility:public"],
|
|
deps = [
|
|
":_pywrap_tf_optimizer",
|
|
":tf_cluster",
|
|
"//tensorflow/core:protos_all_py",
|
|
"@absl_py//absl/logging",
|
|
],
|
|
)
|
|
|
|
tf_python_pybind_extension(
|
|
name = "_pywrap_tf_optimizer",
|
|
srcs = ["tf_optimizer_wrapper.cc"],
|
|
hdrs = if_pywrap(
|
|
if_false = [
|
|
"//tensorflow/cc:pywrap_required_hdrs",
|
|
"//tensorflow/core/grappler:pywrap_required_hdrs",
|
|
"//tensorflow/core/grappler/clusters:pywrap_required_hdrs",
|
|
"//tensorflow/core/grappler/costs:pywrap_required_hdrs",
|
|
"//tensorflow/core/grappler/optimizers:pywrap_required_hdrs",
|
|
"//tensorflow/core/grappler/verifiers:pywrap_required_hdrs",
|
|
],
|
|
),
|
|
enable_stub_generation = True,
|
|
pytype_srcs = [
|
|
"_pywrap_tf_optimizer.pyi",
|
|
],
|
|
# This fails Windows builds. Please check b/266870200 for details.
|
|
# dynamic_deps = ["//tensorflow/python:_pywrap_tensorflow_internal.so"] + select({
|
|
# "//tensorflow:macos": ["//tensorflow:libtensorflow_framework.%s.dylib" % VERSION],
|
|
# "//conditions:default": ["//tensorflow:libtensorflow_framework.so.%s" % VERSION],
|
|
# "//tensorflow:windows": [],
|
|
# }),
|
|
# static_deps = tf_python_pybind_static_deps(),
|
|
deps = [
|
|
"//tensorflow/core:core_cpu_base",
|
|
"//tensorflow/core:framework",
|
|
"//tensorflow/core:framework_headers_lib",
|
|
"//tensorflow/core:framework_lite",
|
|
"//tensorflow/core:lib",
|
|
"//tensorflow/core:lib_headers_for_pybind",
|
|
"//tensorflow/core:lib_proto_parsing",
|
|
"//tensorflow/core:portable_gif_internal",
|
|
"//tensorflow/core:protos_all_cc",
|
|
"//tensorflow/core:session_options",
|
|
"//tensorflow/core/common_runtime:core_cpu_headers_lib",
|
|
"//tensorflow/core/common_runtime:device",
|
|
"//tensorflow/core/common_runtime:device_factory",
|
|
"//tensorflow/core/common_runtime:device_set",
|
|
"//tensorflow/core/common_runtime/gpu:gpu_id",
|
|
"//tensorflow/core/framework:allocator",
|
|
"//tensorflow/core/framework:tensor",
|
|
"//tensorflow/core/platform:errors",
|
|
"//tensorflow/core/platform:status",
|
|
"//tensorflow/python/lib/core:pybind11_status",
|
|
"@com_google_absl//absl/container:flat_hash_map",
|
|
"@com_google_absl//absl/container:flat_hash_set",
|
|
"@com_google_absl//absl/status",
|
|
"@com_google_absl//absl/status:statusor",
|
|
"@pybind11",
|
|
"@pybind11_protobuf//pybind11_protobuf:native_proto_caster",
|
|
"@tsl//tsl/platform:platform_port",
|
|
"@tsl//tsl/platform:thread_annotations",
|
|
"@xla//xla/tsl/protobuf:error_codes_proto_impl_cc",
|
|
] + if_pywrap(
|
|
if_true = [
|
|
"//tensorflow/core/grappler/clusters:cluster",
|
|
"//tensorflow/core/grappler/clusters:utils",
|
|
"//tensorflow/core/grappler:grappler_item_builder",
|
|
"//tensorflow/core/grappler/optimizers:meta_optimizer",
|
|
"//tensorflow/core/grappler/optimizers:graph_optimizer",
|
|
"//tensorflow/core/grappler/verifiers:graph_verifier",
|
|
],
|
|
),
|
|
)
|
|
|
|
tf_py_strict_test(
|
|
name = "tf_optimizer_test",
|
|
size = "small",
|
|
srcs = ["tf_optimizer_test.py"],
|
|
tags = [
|
|
"grappler",
|
|
"no_pip", # tf_optimizer is not available in pip.
|
|
],
|
|
deps = [
|
|
":tf_item",
|
|
":tf_optimizer",
|
|
"//tensorflow/core:protos_all_py",
|
|
"//tensorflow/python/framework:constant_op",
|
|
"//tensorflow/python/framework:dtypes",
|
|
"//tensorflow/python/framework:meta_graph",
|
|
"//tensorflow/python/framework:ops",
|
|
"//tensorflow/python/framework:tensor_shape",
|
|
"//tensorflow/python/framework:test_lib",
|
|
"//tensorflow/python/ops:array_ops",
|
|
"//tensorflow/python/ops:math_ops",
|
|
"//tensorflow/python/ops:variable_v1",
|
|
"//tensorflow/python/ops:while_loop",
|
|
"//tensorflow/python/platform:client_testlib",
|
|
],
|
|
)
|
|
|
|
tf_py_strict_test(
|
|
name = "memory_optimizer_test",
|
|
size = "medium",
|
|
srcs = ["memory_optimizer_test.py"],
|
|
tags = [
|
|
"grappler",
|
|
],
|
|
deps = [
|
|
":tf_optimizer",
|
|
"//tensorflow/core:protos_all_py",
|
|
"//tensorflow/python/client:session",
|
|
"//tensorflow/python/framework:meta_graph",
|
|
"//tensorflow/python/framework:ops",
|
|
"//tensorflow/python/framework:random_seed",
|
|
"//tensorflow/python/framework:test_lib",
|
|
"//tensorflow/python/ops:math_ops",
|
|
"//tensorflow/python/ops:nn",
|
|
"//tensorflow/python/ops:variable_scope",
|
|
"//tensorflow/python/ops:variable_v1",
|
|
"//tensorflow/python/ops:variables",
|
|
"//tensorflow/python/platform:client_testlib",
|
|
"//tensorflow/python/training:training_lib",
|
|
],
|
|
)
|
|
|
|
cuda_py_strict_test(
|
|
name = "constant_folding_test",
|
|
size = "medium",
|
|
srcs = ["constant_folding_test.py"],
|
|
tags = [
|
|
"grappler",
|
|
],
|
|
deps = [
|
|
"//tensorflow/python/eager:backprop",
|
|
"//tensorflow/python/eager:context",
|
|
"//tensorflow/python/eager:def_function",
|
|
"//tensorflow/python/framework:dtypes",
|
|
"//tensorflow/python/framework:ops",
|
|
"//tensorflow/python/framework:test_lib",
|
|
"//tensorflow/python/ops:array_ops",
|
|
"//tensorflow/python/ops:functional_ops",
|
|
"//tensorflow/python/ops:math_ops",
|
|
"//tensorflow/python/ops:resource_variable_ops",
|
|
"//tensorflow/python/ops:while_loop",
|
|
"//tensorflow/python/platform:client_testlib",
|
|
"//third_party/py/numpy",
|
|
],
|
|
)
|
|
|
|
cuda_py_strict_test(
|
|
name = "arithmetic_optimizer_test",
|
|
size = "small",
|
|
srcs = ["arithmetic_optimizer_test.py"],
|
|
tags = [
|
|
"grappler",
|
|
],
|
|
xla_enable_strict_auto_jit = False,
|
|
deps = [
|
|
"//tensorflow/python/eager:context",
|
|
"//tensorflow/python/eager:def_function",
|
|
"//tensorflow/python/ops:array_ops",
|
|
"//tensorflow/python/ops:math_ops",
|
|
"//tensorflow/python/platform:client_testlib",
|
|
],
|
|
)
|
|
|
|
# TODO(b/131764887) Remove once LayoutOptimizer is swapped out with GenericLayoutOptimizer.
|
|
#
|
|
# cuda_py_test(
|
|
# name = "layout_optimizer_test",
|
|
# size = "medium",
|
|
# srcs = [
|
|
# "layout_optimizer_test.py",
|
|
# ],
|
|
# deps = [
|
|
# "//tensorflow/python/platform:client_testlib",
|
|
# "//tensorflow/python/framework:for_generated_wrappers",
|
|
# "//tensorflow/python/ops:array_ops",
|
|
# "//tensorflow/python/ops:functional_ops",
|
|
# "//tensorflow/python/ops:math_ops",
|
|
# "//tensorflow/python:nn",
|
|
# "//tensorflow/python/user_ops:ops",
|
|
# "//tensorflow/python/ops:random_ops",
|
|
# "//tensorflow/python/ops:state_ops",
|
|
# ":tf_cluster",
|
|
# ":tf_optimizer",
|
|
# "//tensorflow/python/training:training",
|
|
# "//third_party/py/numpy",
|
|
# "//tensorflow/core:protos_all_py",
|
|
# "//tensorflow/python/framework:constant_op",
|
|
# "//tensorflow/python/framework:dtypes",
|
|
# ],
|
|
# shard_count = 10,
|
|
# tags = [
|
|
# "grappler",
|
|
# ],
|
|
# # This test will not run on XLA because it primarily tests the TF Classic flow.
|
|
# xla_enable_strict_auto_jit = False,
|
|
# )
|
|
|
|
py_library(
|
|
name = "cost_analyzer",
|
|
srcs = ["cost_analyzer.py"],
|
|
strict_deps = True,
|
|
deps = [
|
|
":_pywrap_cost_analyzer",
|
|
":tf_cluster",
|
|
":tf_item",
|
|
],
|
|
)
|
|
|
|
py_binary(
|
|
name = "cost_analyzer_tool",
|
|
srcs = ["cost_analyzer_tool.py"],
|
|
strict_deps = True,
|
|
deps = [
|
|
":cost_analyzer",
|
|
":tf_optimizer",
|
|
"//tensorflow/core:protos_all_py",
|
|
"//tensorflow/python/framework:importer",
|
|
"//tensorflow/python/framework:ops",
|
|
"//tensorflow/python/platform:gfile",
|
|
"//tensorflow/python/training:saver",
|
|
"@absl_py//absl:app",
|
|
],
|
|
)
|
|
|
|
tf_py_strict_test(
|
|
name = "cost_analyzer_test",
|
|
size = "small",
|
|
srcs = ["cost_analyzer_test.py"],
|
|
tags = [
|
|
"grappler",
|
|
"no_cuda_on_cpu_tap",
|
|
"no_mac",
|
|
"no_pip",
|
|
"no_windows", # TODO(b/151942037)
|
|
],
|
|
deps = [
|
|
":cost_analyzer",
|
|
"//tensorflow/python/framework:constant_op",
|
|
"//tensorflow/python/framework:dtypes",
|
|
"//tensorflow/python/framework:meta_graph",
|
|
"//tensorflow/python/framework:ops",
|
|
"//tensorflow/python/framework:test_lib",
|
|
"//tensorflow/python/ops:array_ops",
|
|
"//tensorflow/python/ops:math_ops",
|
|
"//tensorflow/python/ops:nn_grad",
|
|
"//tensorflow/python/ops:nn_ops",
|
|
"//tensorflow/python/ops:random_ops",
|
|
"//tensorflow/python/ops:variables",
|
|
"//tensorflow/python/platform:client_testlib",
|
|
"//tensorflow/python/training:adam",
|
|
],
|
|
)
|
|
|
|
py_library(
|
|
name = "model_analyzer",
|
|
srcs = [
|
|
"model_analyzer.py",
|
|
],
|
|
strict_deps = True,
|
|
deps = [":_pywrap_model_analyzer"],
|
|
)
|
|
|
|
tf_py_strict_test(
|
|
name = "model_analyzer_test",
|
|
size = "small",
|
|
srcs = ["model_analyzer_test.py"],
|
|
tags = [
|
|
"grappler",
|
|
"no_pip",
|
|
],
|
|
deps = [
|
|
":model_analyzer",
|
|
"//tensorflow/python/framework:constant_op",
|
|
"//tensorflow/python/framework:meta_graph",
|
|
"//tensorflow/python/framework:ops",
|
|
"//tensorflow/python/framework:test_lib",
|
|
"//tensorflow/python/ops:math_ops",
|
|
"//tensorflow/python/platform:client_testlib",
|
|
],
|
|
)
|
|
|
|
cuda_py_strict_test(
|
|
name = "auto_mixed_precision_test",
|
|
size = "medium",
|
|
srcs = [
|
|
"auto_mixed_precision_test.py",
|
|
],
|
|
tags = ["grappler"],
|
|
# This test analyzes the graph, but XLA changes the names of nodes.
|
|
xla_enable_strict_auto_jit = False,
|
|
deps = [
|
|
"//tensorflow/core:protos_all_py",
|
|
"//tensorflow/python:tf2",
|
|
"//tensorflow/python/client:session",
|
|
"//tensorflow/python/data/ops:dataset_ops",
|
|
"//tensorflow/python/framework:constant_op",
|
|
"//tensorflow/python/framework:dtypes",
|
|
"//tensorflow/python/framework:function",
|
|
"//tensorflow/python/framework:ops",
|
|
"//tensorflow/python/framework:random_seed",
|
|
"//tensorflow/python/framework:test_lib",
|
|
"//tensorflow/python/ops:array_ops",
|
|
"//tensorflow/python/ops:init_ops",
|
|
"//tensorflow/python/ops:math_ops",
|
|
"//tensorflow/python/ops:nn",
|
|
"//tensorflow/python/ops:nn_impl",
|
|
"//tensorflow/python/ops:random_ops",
|
|
"//tensorflow/python/ops:tensor_array_ops",
|
|
"//tensorflow/python/ops:variables",
|
|
"//tensorflow/python/ops:while_loop",
|
|
"//tensorflow/python/ops/losses",
|
|
"//tensorflow/python/platform:client_testlib",
|
|
"//tensorflow/python/platform:sysconfig",
|
|
"//tensorflow/python/training:adam",
|
|
"//tensorflow/python/training:gradient_descent",
|
|
"//tensorflow/python/util:_pywrap_utils",
|
|
"//third_party/py/numpy",
|
|
"@absl_py//absl/testing:parameterized",
|
|
],
|
|
)
|
|
|
|
cuda_py_strict_test(
|
|
name = "remapper_test",
|
|
size = "medium",
|
|
srcs = ["remapper_test.py"],
|
|
tags = ["grappler"],
|
|
# This test analyzes the graph, but XLA changes the names of nodes.
|
|
xla_enable_strict_auto_jit = False,
|
|
deps = [
|
|
":tf_optimizer",
|
|
"//tensorflow/core:protos_all_py",
|
|
"//tensorflow/python/client:session",
|
|
"//tensorflow/python/framework:constant_op",
|
|
"//tensorflow/python/framework:dtypes",
|
|
"//tensorflow/python/framework:meta_graph",
|
|
"//tensorflow/python/framework:ops",
|
|
"//tensorflow/python/framework:test_lib",
|
|
"//tensorflow/python/ops:array_ops",
|
|
"//tensorflow/python/ops:init_ops",
|
|
"//tensorflow/python/ops:math_ops",
|
|
"//tensorflow/python/ops:nn",
|
|
"//tensorflow/python/ops:nn_ops",
|
|
"//tensorflow/python/ops:random_ops",
|
|
"//tensorflow/python/ops:variables",
|
|
"//tensorflow/python/platform:client_testlib",
|
|
"//tensorflow/python/platform:sysconfig",
|
|
"//tensorflow/python/util:_pywrap_utils",
|
|
"@absl_py//absl/testing:parameterized",
|
|
],
|
|
)
|
|
|
|
tf_python_pybind_extension(
|
|
name = "_pywrap_graph_analyzer",
|
|
srcs = ["graph_analyzer_tool_wrapper.cc"],
|
|
enable_stub_generation = True,
|
|
pytype_srcs = [
|
|
"_pywrap_graph_analyzer.pyi",
|
|
],
|
|
deps = [
|
|
"//tensorflow/core/grappler/graph_analyzer:graph_analyzer_tool",
|
|
"@pybind11",
|
|
],
|
|
)
|
|
|
|
py_binary(
|
|
name = "graph_analyzer",
|
|
srcs = ["graph_analyzer.py"],
|
|
strict_deps = True,
|
|
deps = [
|
|
":_pywrap_graph_analyzer",
|
|
"@absl_py//absl:app",
|
|
],
|
|
)
|