• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Tensorflow util package
2
3load("//tensorflow:tensorflow.bzl", "py_strict_library")
4
5# buildifier: disable=same-origin-load
6load("//tensorflow:tensorflow.bzl", "get_compatible_with_portable")
7
8# buildifier: disable=same-origin-load
9load("//tensorflow:tensorflow.bzl", "tf_python_pybind_extension")
10
11# buildifier: disable=same-origin-load
12load("//tensorflow:tensorflow.bzl", "tf_py_test")
13
14# buildifier: disable=same-origin-load
15load("//tensorflow/core/platform:build_config.bzl", "tf_proto_library")  # @unused
16load("//tensorflow/core/platform:build_config_root.bzl", "if_static")
17
18visibility = [
19    "//engedu/ml/tf_from_scratch:__pkg__",
20    "//third_party/cloud_tpu/convergence_tools:__subpackages__",
21    "//third_party/mlperf:__subpackages__",
22    "//tensorflow:internal",
23    "//tensorflow/lite/toco/python:__pkg__",
24    "//tensorflow_models:__subpackages__",
25    "//tensorflow_model_optimization:__subpackages__",
26    "//third_party/py/cleverhans:__subpackages__",
27    "//third_party/py/launchpad:__subpackages__",
28    "//third_party/py/reverb:__subpackages__",
29    "//third_party/py/neural_structured_learning:__subpackages__",
30    "//third_party/py/tensorflow_examples:__subpackages__",
31    "//third_party/py/tf_agents:__subpackages__",  # For benchmarks.
32    "//third_party/py/tf_slim:__subpackages__",
33    "//third_party/py/tensorflow_docs:__subpackages__",
34    "//third_party/py/keras:__subpackages__",
35]
36
37package(
38    default_visibility = visibility,
39    licenses = ["notice"],  # Apache 2.0
40)
41
42py_strict_library(
43    name = "core",
44    deps = [
45        ":tf_decorator",
46        ":tf_export",
47        ":tf_stack",
48    ],
49)
50
51# TODO(mdan): Move this utility outside of TF.
52cc_library(
53    name = "kernel_registry",
54    srcs = ["kernel_registry.cc"],
55    hdrs = ["kernel_registry.h"],
56    deps = [
57        "//tensorflow/core:framework",
58        "//tensorflow/core:lib",
59        "//tensorflow/core:protos_all_cc",
60    ],
61    alwayslink = 1,
62)
63
64tf_python_pybind_extension(
65    name = "_pywrap_tfprof",
66    srcs = ["tfprof_wrapper.cc"],
67    module_name = "_pywrap_tfprof",
68    deps = [
69        "//tensorflow/core:framework_headers_lib",
70        "//tensorflow/core:lib_headers_for_pybind",
71        "//tensorflow/core/profiler/internal:print_model_analysis_hdr",
72        "//third_party/eigen3",
73        "//third_party/python_runtime:headers",
74        "@com_google_absl//absl/strings",
75        "@pybind11",
76    ],
77)
78
79tf_python_pybind_extension(
80    name = "_pywrap_utils",
81    srcs = ["util_wrapper.cc"],
82    hdrs = ["util.h"],
83    module_name = "_pywrap_utils",
84    deps = [
85        "//tensorflow/core/platform:platform_port",
86        "//tensorflow/python:pybind11_lib",
87        "//third_party/python_runtime:headers",
88        "@pybind11",
89    ],
90)
91
92tf_python_pybind_extension(
93    name = "_pywrap_nest",
94    srcs = ["nest_wrapper.cc"],
95    hdrs = ["nest.h"],
96    module_name = "_pywrap_nest",
97    deps = [
98        "//tensorflow/python:pybind11_lib",
99        "//third_party/python_runtime:headers",
100        "@pybind11",
101    ],
102)
103
104cc_library(
105    name = "cpp_nest",
106    srcs = ["nest.cc"],
107    hdrs = ["nest.h"],
108    deps = [
109        ":cpp_python_util",
110        "//tensorflow/core:lib",
111        "//tensorflow/core:lib_internal",
112        "//tensorflow/core/platform:logging",
113        "//tensorflow/core/platform:stringpiece",
114        "//tensorflow/python/lib/core:safe_pyobject_ptr",
115        "//third_party/python_runtime:headers",
116    ],
117    alwayslink = 1,
118)
119
120tf_python_pybind_extension(
121    name = "_pywrap_kernel_registry",
122    srcs = ["kernel_registry_wrapper.cc"],
123    hdrs = ["kernel_registry.h"],
124    module_name = "_pywrap_kernel_registry",
125    deps = [
126        "//tensorflow/core:framework_headers_lib",
127        "//tensorflow/core:lib_headers_for_pybind",
128        "//tensorflow/core:protos_all_cc",
129        "//tensorflow/python:pybind11_lib",
130        "//third_party/python_runtime:headers",
131        "@pybind11",
132    ],
133)
134
135tf_python_pybind_extension(
136    name = "_pywrap_stat_summarizer",
137    srcs = ["stat_summarizer_wrapper.cc"],
138    module_name = "_pywrap_stat_summarizer",
139    deps = [
140        "//tensorflow/core:framework_headers_lib",
141        "//tensorflow/core:protos_all_cc",
142        "//third_party/eigen3",
143        "//third_party/python_runtime:headers",
144        "@com_google_absl//absl/memory",
145        "@pybind11",
146    ],
147)
148
149tf_python_pybind_extension(
150    name = "_pywrap_tensor_float_32_execution",
151    srcs = ["tensor_float_32.cc"],
152    hdrs = ["//tensorflow/core/platform:tensor_float_32_hdr"],
153    compatible_with = get_compatible_with_portable(),
154    module_name = "_pywrap_tensor_float_32_execution",
155    deps = [
156        "@pybind11",
157    ],
158)
159
160tf_python_pybind_extension(
161    name = "_pywrap_util_port",
162    srcs = ["port_wrapper.cc"],
163    hdrs = ["//tensorflow/core/util:port_hdrs"],
164    module_name = "_pywrap_util_port",
165    deps = [
166        "//tensorflow/core/util:port",
167        "//third_party/python_runtime:headers",
168        "@pybind11",
169    ],
170)
171
172tf_python_pybind_extension(
173    name = "_pywrap_transform_graph",
174    srcs = ["transform_graph_wrapper.cc"],
175    hdrs = ["//tensorflow/tools/graph_transforms:transform_graph_hdrs"],
176    module_name = "_pywrap_transform_graph",
177    deps = [
178        "//tensorflow/core:framework_headers_lib",
179        "//tensorflow/core:lib_headers_for_pybind",
180        "//tensorflow/core:protos_all_cc",
181        "//tensorflow/python/lib/core:pybind11_status",
182        "//third_party/python_runtime:headers",
183        "@pybind11",
184    ],
185)
186
187tf_python_pybind_extension(
188    name = "_pywrap_checkpoint_reader",
189    srcs = ["py_checkpoint_reader_wrapper.cc"],
190    hdrs = [
191        "//tensorflow/c:checkpoint_reader_hdrs",
192        "//tensorflow/c:headers",
193        "//tensorflow/c/eager:headers",
194        "//tensorflow/python/lib/core:ndarray_tensor_hdr",
195        "//tensorflow/python/lib/core:py_exception_registry_hdr",
196        "//tensorflow/python/lib/core:safe_ptr_hdr",
197    ],
198    module_name = "_pywrap_checkpoint_reader",
199    deps = [
200        "//tensorflow/core:lib_headers_for_pybind",
201        "//tensorflow/core:op_gen_lib",
202        "//tensorflow/core:protos_all_cc",
203        "//tensorflow/core/util/tensor_bundle:tensor_bundle_headers_lib",
204        "//tensorflow/python:pybind11_lib",
205        "//tensorflow/python:pybind11_status",
206        "//tensorflow/python/lib/core:safe_pyobject_ptr",
207        "//third_party/py/numpy:headers",
208        "//third_party/python_runtime:headers",
209        "@com_google_absl//absl/strings",
210        "@pybind11",
211    ],
212)
213
214cc_library(
215    name = "cpp_python_util",
216    srcs = ["util.cc"],
217    hdrs = ["util.h"],
218    deps = [
219        "//tensorflow/core:lib",
220        "//tensorflow/core:lib_internal",
221        "//tensorflow/python/lib/core:safe_ptr",
222        "//tensorflow/python/lib/core:safe_pyobject_ptr",
223        "//third_party/python_runtime:headers",
224        "@com_google_absl//absl/memory",
225    ],
226)
227
228tf_py_test(
229    name = "decorator_utils_test",
230    srcs = ["decorator_utils_test.py"],
231    python_version = "PY3",
232    deps = [
233        ":util",
234        "//tensorflow/python:client_testlib",
235        "//tensorflow/python:platform",
236    ],
237)
238
239tf_py_test(
240    name = "deprecation_test",
241    srcs = ["deprecation_test.py"],
242    python_version = "PY3",
243    deps = [
244        ":util",
245        "//tensorflow/python:client_testlib",
246        "//tensorflow/python:platform",
247    ],
248)
249
250tf_py_test(
251    name = "dispatch_test",
252    srcs = ["dispatch_test.py"],
253    python_version = "PY3",
254    deps = [
255        ":util",
256        "//tensorflow/python:client_testlib",
257        "//tensorflow/python:platform",
258    ],
259)
260
261tf_py_test(
262    name = "keyword_args_test",
263    srcs = ["keyword_args_test.py"],
264    python_version = "PY3",
265    deps = [
266        ":util",
267        "//tensorflow/python:client_testlib",
268    ],
269)
270
271py_strict_library(
272    name = "tf_export",
273    srcs = ["tf_export.py"],
274    compatible_with = get_compatible_with_portable(),
275    srcs_version = "PY3",
276    visibility = ["//tensorflow:__subpackages__"],
277    deps = [
278        ":tf_decorator",
279    ],
280)
281
282tf_py_test(
283    name = "tf_export_test",
284    srcs = ["tf_export_test.py"],
285    python_version = "PY3",
286    deps = [
287        ":util",
288        "//tensorflow/python:client_testlib",
289        "//tensorflow/python:platform",
290    ],
291)
292
293# Leaf library: may not depend on anything else inside TensorFlow.
294# TODO(mdan): Move this utility outside of TF.
295py_strict_library(
296    name = "tf_decorator",
297    srcs = [
298        "tf_contextlib.py",
299        "tf_decorator.py",
300        "tf_inspect.py",
301    ],
302    compatible_with = get_compatible_with_portable(),
303    srcs_version = "PY3",
304    visibility = [
305        "//tensorflow:__subpackages__",
306        # TODO(mdan): Remove these dependencies.
307        "//third_party/py/tf_slim:__subpackages__",
308        "//learning/deepmind/research/language/translation/lm:__subpackages__",
309    ],
310    deps = [
311        "@six_archive//:six",
312    ],
313)
314
315# Note: this is a heavyweight library specialized for TensorFlow graphs. Do not use for
316# other purposes.
317py_strict_library(
318    name = "tf_stack",
319    srcs = ["tf_stack.py"],
320    srcs_version = "PY3",
321    # TODO(mdan): Remove public visibility.
322    visibility = ["//visibility:public"],
323    deps = [
324        ":_tf_stack",
325        "@six_archive//:six",
326    ],
327)
328
329tf_python_pybind_extension(
330    name = "_tf_stack",
331    srcs = ["tf_stack.cc"],
332    hdrs = [
333        "//tensorflow/c:headers",
334        "//tensorflow/c/eager:headers",
335        # Using header directly is required to avoid ODR violations.
336        "stack_trace.h",
337    ],
338    # TODO(b/138203821): change to "util._tf_stack" once the bug is fixed.
339    module_name = "_tf_stack",
340    deps = [
341        "@com_google_absl//absl/algorithm:container",
342        "@com_google_absl//absl/strings",
343        "@com_google_absl//absl/strings:str_format",
344        "@com_google_absl//absl/types:span",
345        "@pybind11",
346        "//third_party/python_runtime:headers",  # buildcleaner: keep
347        "//tensorflow/c:pywrap_required_hdrs",
348        "//tensorflow/core/common_runtime:core_cpu_headers_lib",
349        "//tensorflow/core/framework:pywrap_required_hdrs",
350        "//tensorflow/core/platform:path",
351    ] + if_static([
352        ":stack_trace",
353    ]),
354)
355
356tf_py_test(
357    name = "tf_stack_test",
358    srcs = ["tf_stack_test.py"],
359    python_version = "PY3",
360    deps = [
361        ":tf_export",
362        ":tf_stack",
363        "//tensorflow/python:client_testlib",
364    ],
365)
366
367cc_library(
368    name = "stack_trace",
369    srcs = ["stack_trace.cc"],
370    hdrs = ["stack_trace.h"],
371    deps = [
372        "//tensorflow/core/platform:status",
373        "//tensorflow/core/platform:str_util",
374        "//tensorflow/core/platform:stringpiece",
375        "//tensorflow/core/util:managed_stack_trace",
376        "//third_party/python_runtime:headers",  # buildcleaner: keep
377        "@com_google_absl//absl/base:core_headers",
378        "@com_google_absl//absl/container:flat_hash_map",
379        "@com_google_absl//absl/container:flat_hash_set",
380        "@com_google_absl//absl/container:inlined_vector",
381        "@com_google_absl//absl/types:optional",
382    ],
383)
384
385cc_library(
386    name = "function_parameter_canonicalizer",
387    srcs = ["function_parameter_canonicalizer.cc"],
388    hdrs = ["function_parameter_canonicalizer.h"],
389    deps = [
390        "//tensorflow/core/platform:logging",
391        "//tensorflow/core/platform:macros",
392        "//tensorflow/python/lib/core:py_util",
393        "//tensorflow/python/lib/core:safe_pyobject_ptr",
394        "//third_party/python_runtime:headers",  # buildcleaner: keep
395        "@com_google_absl//absl/container:flat_hash_set",
396        "@com_google_absl//absl/types:span",
397    ],
398)
399
400tf_python_pybind_extension(
401    name = "_function_parameter_canonicalizer_binding_for_test",
402    testonly = True,
403    srcs = ["function_parameter_canonicalizer_binding_for_test.cc"],
404    hdrs = [
405        "function_parameter_canonicalizer.h",
406    ],
407    module_name = "_function_parameter_canonicalizer_binding_for_test",
408    deps = [
409        "//tensorflow/core:lib",
410        "//tensorflow/python/lib/core:safe_pyobject_ptr_required_hdrs",
411        "//third_party/python_runtime:headers",  # buildcleaner: keep
412        "@com_google_absl//absl/types:span",
413        "@pybind11",
414    ],
415)
416
417tf_py_test(
418    name = "function_parameter_canonicalizer_test",
419    srcs = ["function_parameter_canonicalizer_test.py"],
420    python_version = "PY3",
421    tags = [
422        "no_pip",  # b/168621686
423        "no_windows",  # b/169275019
424    ],
425    deps = [
426        ":_function_parameter_canonicalizer_binding_for_test",
427        "//tensorflow/python:client_testlib",
428    ],
429)
430
431py_library(
432    name = "util",
433    srcs = glob(
434        ["**/*.py"],
435        exclude = [
436            "example_parser*",
437            "tf_contextlib.py",
438            "tf_should_use.py",
439            "tf_export.py",
440            "tf_stack.py",
441            "tf_decorator.py",
442            "**/*_test.py",
443        ],
444    ),
445    compatible_with = get_compatible_with_portable(),
446    srcs_version = "PY3",
447    visibility = visibility + [
448        "//tensorflow:__pkg__",
449        "//third_party/py/tensorflow_core:__subpackages__",
450        "//third_party/py/tf_agents:__subpackages__",
451        "//third_party/py/tfx:__subpackages__",
452    ],
453    deps = [
454        ":_pywrap_tensor_float_32_execution",
455        # global_test_configuration is added here because all major tests depend on this
456        # library. It isn't possible to add these test dependencies via tensorflow.bzl's
457        # py_test because not all tensorflow tests use tensorflow.bzl's py_test.
458        "//tensorflow/python:global_test_configuration",
459        ":tf_decorator",
460        ":tf_export",
461        "@org_python_pypi_backports_weakref",
462        "@com_google_protobuf//:protobuf_python",
463        "//third_party/py/numpy",
464        "@six_archive//:six",
465        "@wrapt",
466        "//tensorflow/tools/docs:doc_controls",
467        "//tensorflow/tools/compatibility:all_renames_v2",
468    ],
469)
470
471tf_py_test(
472    name = "object_identity_test",
473    size = "small",
474    srcs = ["object_identity_test.py"],
475    python_version = "PY3",
476)
477
478# Placeholder for intenal nest_test comments.
479tf_py_test(
480    name = "nest_test",
481    size = "small",
482    srcs = ["nest_test.py"],
483    main = "nest_test.py",
484    python_version = "PY3",
485    deps = [":nest_test_main_lib"],
486)
487
488py_library(
489    name = "nest_test_main_lib",
490    testonly = True,
491    srcs = ["nest_test.py"],
492    srcs_version = "PY3",
493    deps = [
494        ":util",
495        "//tensorflow/python:array_ops",
496        "//tensorflow/python:client_testlib",
497        "//tensorflow/python:framework",
498        "//tensorflow/python:framework_for_generated_wrappers",
499        "//tensorflow/python:math_ops",
500        "//third_party/py/numpy",
501        "@absl_py//absl/testing:parameterized",
502    ],
503)
504
505tf_py_test(
506    name = "serialization_test",
507    size = "small",
508    srcs = ["serialization_test.py"],
509    main = "serialization_test.py",
510    python_version = "PY3",
511    deps = [
512        ":util",
513        "//tensorflow/python:client_testlib",
514    ],
515)
516
517tf_py_test(
518    name = "function_utils_test",
519    srcs = ["function_utils_test.py"],
520    python_version = "PY3",
521    deps = [
522        ":util",
523        "//tensorflow/python:client_testlib",
524    ],
525)
526
527tf_py_test(
528    name = "tf_contextlib_test",
529    size = "small",
530    srcs = ["tf_contextlib_test.py"],
531    python_version = "PY3",
532    deps = [
533        ":util",
534        "//tensorflow/python:client_testlib",
535    ],
536)
537
538tf_py_test(
539    name = "tf_decorator_test",
540    size = "small",
541    srcs = ["tf_decorator_test.py"],
542    python_version = "PY3",
543    deps = [
544        ":util",
545        "//tensorflow/python:client_testlib",
546    ],
547)
548
549py_library(
550    name = "tf_should_use",
551    srcs = ["tf_should_use.py"],
552    srcs_version = "PY3",
553    deps = [
554        ":util",
555        "//tensorflow/python:framework_ops",
556        "//tensorflow/python/eager:context",
557        "@six_archive//:six",
558    ],
559)
560
561tf_py_test(
562    name = "tf_should_use_test",
563    size = "small",
564    srcs = ["tf_should_use_test.py"],
565    python_version = "PY3",
566    deps = [
567        ":tf_should_use",
568        "//tensorflow/python:client_testlib",
569    ],
570)
571
572tf_py_test(
573    name = "tf_inspect_test",
574    size = "small",
575    srcs = ["tf_inspect_test.py"],
576    python_version = "PY3",
577    deps = [
578        ":util",
579        "//tensorflow/python:client_testlib",
580    ],
581)
582
583py_library(
584    name = "example_parser_configuration",
585    srcs = ["example_parser_configuration.py"],
586    srcs_version = "PY3",
587    visibility = ["//visibility:public"],
588    deps = [
589        "//tensorflow/core:protos_all_py",
590        "//tensorflow/python:framework",
591        "//tensorflow/python:framework_for_generated_wrappers",
592    ],
593)
594
595tf_py_test(
596    name = "lock_util_test",
597    size = "small",
598    srcs = ["lock_util_test.py"],
599    main = "lock_util_test.py",
600    python_version = "PY3",
601    deps = [
602        ":util",
603        "//tensorflow/python:client_testlib",
604        "@absl_py//absl/testing:parameterized",
605    ],
606)
607
608tf_python_pybind_extension(
609    name = "fast_module_type",
610    srcs = ["fast_module_type.cc"],
611    module_name = "fast_module_type",
612    deps = [
613        "//tensorflow/core/platform:logging",
614        "//third_party/python_runtime:headers",
615        "@com_google_absl//absl/container:flat_hash_map",
616        "@pybind11",
617    ],
618)
619
620tf_py_test(
621    name = "fast_module_type_test",
622    srcs = ["fast_module_type_test.py"],
623    python_version = "PY3",
624    deps = [
625        ":fast_module_type",
626        "//tensorflow/python:platform",
627    ],
628)
629
630tf_py_test(
631    name = "module_wrapper_test",
632    size = "small",
633    srcs = ["module_wrapper_test.py"],
634    python_version = "PY3",
635    deps = [
636        ":fast_module_type",
637        ":util",
638        "//tensorflow/python:client_testlib",
639        "//tensorflow/tools/compatibility:all_renames_v2",
640        "@six_archive//:six",
641    ],
642)
643
644tf_proto_library(
645    name = "compare_test_proto",
646    testonly = 1,
647    srcs = ["protobuf/compare_test.proto"],
648    cc_api_version = 2,
649)
650
651tf_py_test(
652    name = "protobuf_compare_test",
653    size = "small",
654    srcs = ["protobuf/compare_test.py"],
655    main = "protobuf/compare_test.py",
656    python_version = "PY3",
657    tags = ["no_pip"],  # compare_test_pb2 proto is not available in pip.
658    deps = [
659        ":compare_test_proto_py",
660        ":util",
661        "//tensorflow/python:platform_test",
662        "@six_archive//:six",
663    ],
664)
665
666tf_py_test(
667    name = "example_parser_configuration_test",
668    size = "small",
669    srcs = ["example_parser_configuration_test.py"],
670    main = "example_parser_configuration_test.py",
671    python_version = "PY3",
672    deps = [
673        ":example_parser_configuration",
674        "//tensorflow/python:array_ops",
675        "//tensorflow/python:client",
676        "//tensorflow/python:client_testlib",
677        "//tensorflow/python:framework_for_generated_wrappers",
678        "//tensorflow/python:parsing_ops",
679    ],
680)
681
682filegroup(
683    name = "util_hdr",
684    srcs = ["util.h"],
685)
686
687filegroup(
688    name = "compare_test_proto_src",
689    srcs = ["protobuf/compare_test.proto"],
690)
691
692# copybara:uncomment_begin(google-only)
693# py_proto_library(
694#     name = "compare_test_py_pb2",
695#     testonly = 1,
696#     has_services = 0,
697#     api_version = 2,
698#     deps = [":compare_test_proto"],
699# )
700# copybara:uncomment_end
701