1load("//tensorflow:tensorflow.bzl", "pytype_strict_library") 2load("//tensorflow:tensorflow.bzl", "tf_py_test") 3load("//tensorflow:tensorflow.bzl", "cuda_py_test") 4 5package( 6 default_visibility = ["//tensorflow:internal"], 7 licenses = ["notice"], 8) 9 10cuda_py_test( 11 name = "cross_device_ops_test", 12 srcs = ["cross_device_ops_test.py"], 13 python_version = "PY3", 14 tags = [ 15 "multi_and_single_gpu", 16 "no_windows_gpu", # b/216367668 17 ], 18 deps = [ 19 "//tensorflow/python:array_ops", 20 "//tensorflow/python:collective_ops", 21 "//tensorflow/python:constant_op", 22 "//tensorflow/python:framework_ops", 23 "//tensorflow/python:kernels", 24 "//tensorflow/python:math_ops", 25 "//tensorflow/python:variables", 26 "//tensorflow/python/distribute:collective_all_reduce_strategy", 27 "//tensorflow/python/distribute:collective_util", 28 "//tensorflow/python/distribute:combinations", 29 "//tensorflow/python/distribute:cross_device_ops", 30 "//tensorflow/python/distribute:cross_device_utils", 31 "//tensorflow/python/distribute:device_util", 32 "//tensorflow/python/distribute:distribute_utils", 33 "//tensorflow/python/distribute:multi_worker_test_base", 34 "//tensorflow/python/distribute:multi_worker_util", 35 "//tensorflow/python/distribute:reduce_util", 36 "//tensorflow/python/distribute:strategy_combinations", 37 "//tensorflow/python/distribute:test_util", 38 "//tensorflow/python/distribute:values", 39 "//tensorflow/python/distribute/cluster_resolver:cluster_resolver_lib", 40 "//tensorflow/python/eager:context", 41 "//tensorflow/python/eager:test", 42 "@absl_py//absl/testing:parameterized", 43 ], 44) 45 46py_library( 47 name = "all_reduce", 48 srcs = [ 49 "all_reduce.py", 50 ], 51 srcs_version = "PY3", 52 deps = [ 53 "//tensorflow/python:array_ops", 54 "//tensorflow/python:framework_ops", 55 "//tensorflow/python:math_ops", 56 "//tensorflow/python:nccl_ops", 57 ], 58) 59 60tf_py_test( 61 name = "all_reduce_test", 62 srcs = ["all_reduce_test.py"], 63 deps = [ 64 ":all_reduce", 65 "//tensorflow/core:protos_all_py", 66 "//tensorflow/python:array_ops", 67 "//tensorflow/python:client", 68 "//tensorflow/python:client_testlib", 69 "//tensorflow/python:constant_op", 70 "//tensorflow/python:framework_ops", 71 "//tensorflow/python:framework_test_lib", 72 "//tensorflow/python:math_ops", 73 "//tensorflow/python:platform", 74 "//tensorflow/python:platform_test", 75 "//tensorflow/python:state_ops", 76 "//third_party/py/numpy", 77 ], 78) 79 80pytype_strict_library( 81 name = "input_lib", 82 srcs = ["input_lib.py"], 83 srcs_version = "PY3", 84 deps = [ 85 "//tensorflow/python:control_flow_ops", 86 "//tensorflow/python/data/experimental/ops:cardinality", 87 "//tensorflow/python/data/ops:dataset_ops", 88 "//tensorflow/python/data/ops:multi_device_iterator_ops", 89 "//tensorflow/python/data/ops:optional_ops", 90 "//tensorflow/python/distribute:input_lib", 91 "//tensorflow/python/eager:context", 92 "//tensorflow/python/framework:ops", 93 "//tensorflow/python/util", 94 ], 95) 96