wehub-resource-sync
1042 lines
34 KiB
Python
1042 lines
34 KiB
Python
load("//tensorflow:pytype.default.bzl", "pytype_strict_library")
|
|
load("//tensorflow:tensorflow.bzl", "if_google", "py_test")
|
|
load(
|
|
"//tensorflow/dtensor:build_defs.bzl",
|
|
"ALL_BACKENDS",
|
|
"GPU_2DEVS_BACKEND",
|
|
"TPU_V3_DONUT_BACKEND",
|
|
"TPU_V4_DONUT_BACKEND",
|
|
"dtensor_test",
|
|
)
|
|
|
|
# copybara:uncomment package(default_applicable_licenses = ["//tensorflow:license"])
|
|
|
|
# File used by internal tests.
|
|
exports_files([
|
|
"spmd_test.py",
|
|
"collective_test.py",
|
|
])
|
|
|
|
pytype_strict_library(
|
|
name = "test_util",
|
|
testonly = if_google(
|
|
True,
|
|
oss_value = False, # build_pip_package depends on this target.
|
|
),
|
|
srcs = [
|
|
"test_backend_name.py",
|
|
"test_backend_util.py",
|
|
"test_util.py",
|
|
"test_util_ops.py",
|
|
],
|
|
visibility = [
|
|
"//tensorflow/dtensor:dtensor-internal",
|
|
"//tensorflow/dtensor:dtensor-users",
|
|
"//tensorflow/tools/pip_package:__subpackages__",
|
|
],
|
|
deps = [
|
|
"//tensorflow/dtensor/python:api",
|
|
"//tensorflow/dtensor/python:config",
|
|
"//tensorflow/dtensor/python:layout",
|
|
"//tensorflow/dtensor/python:numpy_util",
|
|
"//tensorflow/dtensor/python:tpu_util",
|
|
"//tensorflow/python/compat:v2_compat",
|
|
"//tensorflow/python/eager:context",
|
|
"//tensorflow/python/framework:config",
|
|
"//tensorflow/python/framework:constant_op",
|
|
"//tensorflow/python/framework:device",
|
|
"//tensorflow/python/framework:dtypes",
|
|
"//tensorflow/python/ops:array_ops",
|
|
"//tensorflow/python/ops:array_ops_gen",
|
|
"//tensorflow/python/ops:bitwise_ops_gen",
|
|
"//tensorflow/python/ops:clip_ops",
|
|
"//tensorflow/python/ops:math_ops",
|
|
"//tensorflow/python/ops:math_ops_gen",
|
|
"//tensorflow/python/ops:nn_ops",
|
|
"//tensorflow/python/ops:nn_ops_gen",
|
|
"//tensorflow/python/ops:resource_variable_ops",
|
|
"//tensorflow/python/ops:special_math_ops",
|
|
"//tensorflow/python/ops:spectral_ops_gen",
|
|
"//tensorflow/python/ops:stateless_random_ops_gen",
|
|
"//tensorflow/python/ops:stateless_random_ops_v2_gen",
|
|
"//tensorflow/python/platform:client_testlib",
|
|
"//tensorflow/python/util:numpy_compat",
|
|
"//third_party/py/numpy",
|
|
"@absl_py//absl/flags",
|
|
"@absl_py//absl/logging",
|
|
"@absl_py//absl/testing:parameterized",
|
|
],
|
|
)
|
|
|
|
py_test(
|
|
name = "api_test",
|
|
srcs = [
|
|
"api_test.py",
|
|
],
|
|
strict_deps = True,
|
|
deps = [
|
|
":test_util",
|
|
"//tensorflow/dtensor/python:api",
|
|
"//tensorflow/dtensor/python:d_random",
|
|
"//tensorflow/dtensor/python:layout",
|
|
"//tensorflow/dtensor/python:numpy_util",
|
|
"//tensorflow/python/eager/polymorphic_function",
|
|
"//tensorflow/python/framework:constant_op",
|
|
"//tensorflow/python/framework:dtypes",
|
|
"//tensorflow/python/framework:errors",
|
|
"//tensorflow/python/framework:ops",
|
|
"//tensorflow/python/ops:array_ops",
|
|
"//tensorflow/python/ops:nn_ops",
|
|
"//tensorflow/python/ops:stateless_random_ops",
|
|
"//tensorflow/python/platform:client_testlib",
|
|
"//third_party/py/numpy",
|
|
"@absl_py//absl/testing:parameterized",
|
|
],
|
|
)
|
|
|
|
# TODO(b/301286466): Investigate why python annotation type mismatch is not catptured by the type
|
|
# strict BUILD rules.
|
|
|
|
dtensor_test(
|
|
name = "array_ops_test",
|
|
srcs = ["array_ops_test.py"],
|
|
additional_backends = [],
|
|
deps = [
|
|
":test_util",
|
|
"//tensorflow/dtensor/python:api",
|
|
"//tensorflow/dtensor/python:layout",
|
|
"//tensorflow/python/eager/polymorphic_function",
|
|
"//tensorflow/python/framework:combinations",
|
|
"//tensorflow/python/ops:array_ops",
|
|
"//tensorflow/python/platform:client_testlib",
|
|
"//third_party/py/numpy",
|
|
],
|
|
)
|
|
|
|
dtensor_test(
|
|
name = "batchparallel_spmd_test",
|
|
srcs = ["batchparallel_spmd_test.py"],
|
|
additional_backends = [TPU_V4_DONUT_BACKEND],
|
|
main = "batchparallel_spmd_test.py",
|
|
shard_count = {
|
|
"cpu": 4,
|
|
"gpu": 4,
|
|
"tpu": 4,
|
|
TPU_V4_DONUT_BACKEND: 8,
|
|
},
|
|
deps = [
|
|
":test_util",
|
|
"//tensorflow/dtensor/python:api",
|
|
"//tensorflow/dtensor/python:layout",
|
|
"//tensorflow/dtensor/python:numpy_util",
|
|
"//tensorflow/python/eager:backprop",
|
|
"//tensorflow/python/eager:context",
|
|
"//tensorflow/python/eager/polymorphic_function",
|
|
"//tensorflow/python/framework:constant_op",
|
|
"//tensorflow/python/framework:dtypes",
|
|
"//tensorflow/python/framework:ops",
|
|
"//tensorflow/python/ops:array_ops",
|
|
"//tensorflow/python/ops:array_ops_gen",
|
|
"//tensorflow/python/ops:image_ops_gen",
|
|
"//tensorflow/python/ops:linalg_ops_gen",
|
|
"//tensorflow/python/ops:nn_impl",
|
|
"//tensorflow/python/ops:nn_ops",
|
|
"//tensorflow/python/platform:client_testlib",
|
|
"//third_party/py/numpy",
|
|
"@absl_py//absl/testing:parameterized",
|
|
],
|
|
)
|
|
|
|
dtensor_test(
|
|
name = "cache_test",
|
|
srcs = ["cache_test.py"],
|
|
main = "cache_test.py",
|
|
tags = [
|
|
"nomultivm",
|
|
],
|
|
deps = [
|
|
":test_util",
|
|
"//tensorflow/dtensor/python:api",
|
|
"//tensorflow/dtensor/python:d_variable",
|
|
"//tensorflow/dtensor/python:layout",
|
|
"//tensorflow/python/eager/polymorphic_function",
|
|
"//tensorflow/python/framework:combinations",
|
|
"//tensorflow/python/framework:constant_op",
|
|
"//tensorflow/python/framework:dtypes",
|
|
"//tensorflow/python/ops:array_ops",
|
|
"//tensorflow/python/ops:math_ops",
|
|
"//tensorflow/python/ops:stateless_random_ops_gen",
|
|
"//tensorflow/python/platform:client_testlib",
|
|
"//third_party/py/numpy",
|
|
],
|
|
)
|
|
|
|
dtensor_test(
|
|
name = "config_test",
|
|
srcs = ["config_test.py"],
|
|
main = "config_test.py",
|
|
deps = [
|
|
":test_util",
|
|
"//tensorflow/dtensor/python:config",
|
|
"//tensorflow/python/eager:context",
|
|
"//tensorflow/python/framework:device",
|
|
"//tensorflow/python/platform:client_testlib",
|
|
],
|
|
)
|
|
|
|
dtensor_test(
|
|
name = "collective_test",
|
|
srcs = ["collective_test.py"],
|
|
additional_backends = [
|
|
TPU_V3_DONUT_BACKEND,
|
|
GPU_2DEVS_BACKEND,
|
|
],
|
|
deps = [
|
|
":test_util",
|
|
"//tensorflow/dtensor/python:api",
|
|
"//tensorflow/dtensor/python:d_variable",
|
|
"//tensorflow/dtensor/python:dtensor_device",
|
|
"//tensorflow/dtensor/python:layout",
|
|
"//tensorflow/python/eager/polymorphic_function",
|
|
"//tensorflow/python/framework:constant_op",
|
|
"//tensorflow/python/framework:dtypes",
|
|
"//tensorflow/python/ops:array_ops",
|
|
"//tensorflow/python/ops:array_ops_gen",
|
|
"//tensorflow/python/ops:math_ops",
|
|
"//tensorflow/python/ops:math_ops_gen",
|
|
"//tensorflow/python/platform:client_testlib",
|
|
"//third_party/py/numpy",
|
|
"@absl_py//absl/testing:parameterized",
|
|
],
|
|
)
|
|
|
|
dtensor_test(
|
|
name = "collective_combine_all_reduce_test",
|
|
srcs = [":collective_test.py"],
|
|
args = if_google(
|
|
[
|
|
"--vmodule=dtensor_graph_to_mlir_pass=4",
|
|
],
|
|
[],
|
|
),
|
|
env = {
|
|
"DTENSOR_ENABLE_COMBINE_ALL_REDUCES_OPTIMIZATION": "1",
|
|
},
|
|
deps = [
|
|
":test_util",
|
|
"//tensorflow/dtensor/python:api",
|
|
"//tensorflow/dtensor/python:d_variable",
|
|
"//tensorflow/dtensor/python:dtensor_device",
|
|
"//tensorflow/dtensor/python:layout",
|
|
"//tensorflow/python/eager/polymorphic_function",
|
|
"//tensorflow/python/framework:constant_op",
|
|
"//tensorflow/python/framework:dtypes",
|
|
"//tensorflow/python/ops:array_ops",
|
|
"//tensorflow/python/ops:array_ops_gen",
|
|
"//tensorflow/python/ops:math_ops",
|
|
"//tensorflow/python/ops:math_ops_gen",
|
|
"//tensorflow/python/platform:client_testlib",
|
|
"//third_party/py/numpy",
|
|
"@absl_py//absl/testing:parameterized",
|
|
],
|
|
)
|
|
|
|
dtensor_test(
|
|
name = "conv_test",
|
|
srcs = [
|
|
"conv_test.py",
|
|
],
|
|
additional_backends = [TPU_V3_DONUT_BACKEND],
|
|
# All tests require 8 TPUs.
|
|
disable = ["tpu"],
|
|
shard_count = {
|
|
"cpu": 4,
|
|
"gpu": 4,
|
|
TPU_V3_DONUT_BACKEND: 4,
|
|
},
|
|
deps = [
|
|
":test_util",
|
|
"//tensorflow/dtensor/python:layout",
|
|
"//tensorflow/dtensor/python:numpy_util",
|
|
"//tensorflow/python/eager:backprop",
|
|
"//tensorflow/python/eager/polymorphic_function",
|
|
"//tensorflow/python/framework:constant_op",
|
|
"//tensorflow/python/framework:dtypes",
|
|
"//tensorflow/python/ops:nn_ops",
|
|
"//tensorflow/python/ops:special_math_ops",
|
|
"//tensorflow/python/platform:client_testlib",
|
|
"//third_party/py/numpy",
|
|
"@absl_py//absl/testing:parameterized",
|
|
],
|
|
)
|
|
|
|
dtensor_test(
|
|
name = "device_test",
|
|
srcs = ["device_test.py"],
|
|
additional_backends = [TPU_V3_DONUT_BACKEND],
|
|
main = "device_test.py",
|
|
shard_count = {
|
|
TPU_V3_DONUT_BACKEND: 32,
|
|
},
|
|
tags = [
|
|
"nomultivm",
|
|
],
|
|
deps = [
|
|
":test_util",
|
|
"//tensorflow/dtensor/python:api",
|
|
"//tensorflow/dtensor/python:d_variable",
|
|
"//tensorflow/dtensor/python:dtensor_device",
|
|
"//tensorflow/dtensor/python:layout",
|
|
"//tensorflow/python/eager/polymorphic_function",
|
|
"//tensorflow/python/framework:constant_op",
|
|
"//tensorflow/python/framework:dtypes",
|
|
"//tensorflow/python/framework:errors",
|
|
"//tensorflow/python/framework:ops",
|
|
"//tensorflow/python/ops:array_ops",
|
|
"//tensorflow/python/ops:collective_ops_gen",
|
|
"//tensorflow/python/ops:math_ops",
|
|
"//tensorflow/python/ops:math_ops_gen",
|
|
"//tensorflow/python/ops:sparse_ops",
|
|
"//tensorflow/python/ops:stateless_random_ops",
|
|
"//tensorflow/python/ops:variables",
|
|
"//tensorflow/python/platform:client_testlib",
|
|
"//third_party/py/numpy",
|
|
"@absl_py//absl/testing:parameterized",
|
|
],
|
|
)
|
|
|
|
py_test(
|
|
name = "input_util_test",
|
|
size = "medium",
|
|
srcs = ["input_util_test.py"],
|
|
shard_count = 8,
|
|
strict_deps = True,
|
|
deps = [
|
|
":test_util",
|
|
"//tensorflow/dtensor/python:api",
|
|
"//tensorflow/dtensor/python:input_util",
|
|
"//tensorflow/dtensor/python:layout",
|
|
"//tensorflow/dtensor/python:mesh_util",
|
|
"//tensorflow/python/data/ops:dataset_ops",
|
|
"//tensorflow/python/eager/polymorphic_function",
|
|
"//tensorflow/python/framework:constant_op",
|
|
"//tensorflow/python/framework:dtypes",
|
|
"//tensorflow/python/framework:errors",
|
|
"//tensorflow/python/framework:tensor_shape",
|
|
"//tensorflow/python/framework:tensor_spec",
|
|
"//tensorflow/python/ops:math_ops",
|
|
"//tensorflow/python/ops:stateless_random_ops",
|
|
"//tensorflow/python/ops:variables",
|
|
"//tensorflow/python/platform:client_testlib",
|
|
"//third_party/py/numpy",
|
|
"@absl_py//absl/testing:parameterized",
|
|
],
|
|
)
|
|
|
|
py_test(
|
|
name = "multi_client_input_util_test",
|
|
timeout = "long",
|
|
srcs = ["multi_client_input_util_test.py"],
|
|
env = {
|
|
"TF2_BEHAVIOR": "1",
|
|
},
|
|
shard_count = 8,
|
|
strict_deps = True,
|
|
tags = [
|
|
# ThreadSanitizer does not support starting new threads after multi-threaded fork.
|
|
"notsan",
|
|
"no_oss", # Fails on OSS.
|
|
"nosan", # b/195537906
|
|
],
|
|
deps = [
|
|
":multi_client_test_util",
|
|
":test_util",
|
|
"//tensorflow/core:protos_all_py",
|
|
"//tensorflow/dtensor/python:accelerator_util",
|
|
"//tensorflow/dtensor/python:api",
|
|
"//tensorflow/dtensor/python:config",
|
|
"//tensorflow/dtensor/python:input_util",
|
|
"//tensorflow/dtensor/python:layout",
|
|
"//tensorflow/dtensor/python:mesh_util",
|
|
"//tensorflow/python/data/experimental/service:server_lib",
|
|
"//tensorflow/python/data/ops:dataset_ops",
|
|
"//tensorflow/python/data/ops:readers",
|
|
"//tensorflow/python/eager:context",
|
|
"//tensorflow/python/framework:config",
|
|
"//tensorflow/python/framework:device_spec",
|
|
"//tensorflow/python/framework:dtypes",
|
|
"//tensorflow/python/lib/io:tf_record",
|
|
"//tensorflow/python/ops:array_ops_stack",
|
|
"//tensorflow/python/ops:check_ops",
|
|
"//tensorflow/python/ops:io_ops",
|
|
"//tensorflow/python/ops:parsing_config",
|
|
"//tensorflow/python/ops:parsing_ops",
|
|
"//tensorflow/python/ops:parsing_ops_gen",
|
|
"//tensorflow/python/ops:stateless_random_ops",
|
|
"//tensorflow/python/platform:client_testlib",
|
|
"//third_party/py/numpy",
|
|
"@absl_py//absl/logging",
|
|
"@absl_py//absl/testing:parameterized",
|
|
],
|
|
)
|
|
|
|
dtensor_test(
|
|
name = "layout_test",
|
|
srcs = ["layout_test.py"],
|
|
disable = [
|
|
"gpu",
|
|
"tpu",
|
|
],
|
|
main = "layout_test.py",
|
|
tags = [
|
|
"no_windows",
|
|
],
|
|
deps = [
|
|
":test_util",
|
|
"//tensorflow/dtensor/python:api",
|
|
"//tensorflow/dtensor/python:layout",
|
|
"//tensorflow/python/eager:backprop",
|
|
"//tensorflow/python/eager/polymorphic_function",
|
|
"//tensorflow/python/framework:combinations",
|
|
"//tensorflow/python/framework:dtypes",
|
|
"//tensorflow/python/framework:errors",
|
|
"//tensorflow/python/framework:ops",
|
|
"//tensorflow/python/framework:tensor_shape",
|
|
"//tensorflow/python/ops:array_ops",
|
|
"//tensorflow/python/ops:math_ops",
|
|
"//tensorflow/python/ops:stateless_random_ops",
|
|
"//tensorflow/python/platform:client_testlib",
|
|
"//third_party/py/numpy",
|
|
"@absl_py//absl/testing:parameterized",
|
|
],
|
|
)
|
|
|
|
dtensor_test(
|
|
name = "layout_propagation_test",
|
|
srcs = ["layout_propagation_test.py"],
|
|
args = if_google(
|
|
[
|
|
"--vmodule=dtensor_mlir_passes=4",
|
|
],
|
|
[],
|
|
),
|
|
disable = [
|
|
"gpu",
|
|
"tpu",
|
|
],
|
|
main = "layout_propagation_test.py",
|
|
shard_count = {
|
|
"cpu": 5,
|
|
},
|
|
deps = [
|
|
":test_util",
|
|
"//tensorflow/dtensor/python:api",
|
|
"//tensorflow/dtensor/python:d_variable",
|
|
"//tensorflow/dtensor/python:layout",
|
|
"//tensorflow/dtensor/python:numpy_util",
|
|
"//tensorflow/python/eager:backprop",
|
|
"//tensorflow/python/eager/polymorphic_function",
|
|
"//tensorflow/python/framework:constant_op",
|
|
"//tensorflow/python/framework:dtypes",
|
|
"//tensorflow/python/ops:array_ops",
|
|
"//tensorflow/python/ops:math_ops",
|
|
"//tensorflow/python/ops:nn_ops",
|
|
"//tensorflow/python/ops:nn_ops_gen",
|
|
"//tensorflow/python/ops:stateless_random_ops",
|
|
"//tensorflow/python/ops:variables",
|
|
"//tensorflow/python/platform:client_testlib",
|
|
"//third_party/py/numpy",
|
|
"@absl_py//absl/testing:parameterized",
|
|
],
|
|
)
|
|
|
|
dtensor_test(
|
|
name = "mesh_util_test",
|
|
srcs = ["mesh_util_test.py"],
|
|
additional_backends = [
|
|
TPU_V3_DONUT_BACKEND,
|
|
GPU_2DEVS_BACKEND,
|
|
],
|
|
env = {
|
|
"TF_FORCE_GPU_ALLOW_GROWTH": "true",
|
|
},
|
|
deps = [
|
|
":test_util",
|
|
"//tensorflow/dtensor/python:config",
|
|
"//tensorflow/dtensor/python:mesh_util",
|
|
"//tensorflow/dtensor/python:tpu_util",
|
|
"//tensorflow/python/eager:context",
|
|
"//tensorflow/python/framework:config",
|
|
"//tensorflow/python/framework:device",
|
|
"//tensorflow/python/platform:client_testlib",
|
|
"@absl_py//absl/testing:parameterized",
|
|
],
|
|
)
|
|
|
|
pytype_strict_library(
|
|
name = "multi_client_test_util",
|
|
testonly = if_google(
|
|
True,
|
|
oss_value = False, # build_pip_package depends on this target.
|
|
),
|
|
srcs = ["multi_client_test_util.py"],
|
|
visibility = [
|
|
"//tensorflow/dtensor:dtensor-internal",
|
|
"//tensorflow/tools/pip_package:__subpackages__",
|
|
],
|
|
deps = [
|
|
":test_util",
|
|
"//tensorflow/python/platform:client_testlib",
|
|
"@absl_py//absl/flags",
|
|
"@pypi//portpicker",
|
|
],
|
|
)
|
|
|
|
dtensor_test(
|
|
name = "multi_client_test",
|
|
srcs = ["multi_client_test.py"],
|
|
additional_backends = [
|
|
GPU_2DEVS_BACKEND,
|
|
TPU_V3_DONUT_BACKEND,
|
|
],
|
|
disable = [
|
|
"gpu", # multi-client gpu is tested via GPU_2DEVS_BACKEND.
|
|
"tpu", # multi-client tpu is tested via TPU_V3_DONUT_BACKEND.
|
|
],
|
|
disable_tfrt = [
|
|
"cpu", # TODO(b/217969210): Re-enable in TFRT CPU.
|
|
GPU_2DEVS_BACKEND, # TODO(b/230679405): Re-enable in TFRT GPU.
|
|
],
|
|
tags = [
|
|
"no_windows",
|
|
"nosan",
|
|
], # b/195537906
|
|
deps = [
|
|
":multi_client_test_util",
|
|
":test_util",
|
|
"//tensorflow/dtensor/python:accelerator_util",
|
|
"//tensorflow/dtensor/python:api",
|
|
"//tensorflow/dtensor/python:config",
|
|
"//tensorflow/dtensor/python:d_variable",
|
|
"//tensorflow/dtensor/python:layout",
|
|
"//tensorflow/dtensor/python:mesh_util",
|
|
"//tensorflow/python/eager:backprop",
|
|
"//tensorflow/python/eager/polymorphic_function",
|
|
"//tensorflow/python/framework:config",
|
|
"//tensorflow/python/framework:dtypes",
|
|
"//tensorflow/python/framework:ops",
|
|
"//tensorflow/python/ops:array_ops",
|
|
"//tensorflow/python/ops:math_ops",
|
|
"//tensorflow/python/ops:nn_ops",
|
|
"//tensorflow/python/ops:stateless_random_ops",
|
|
"//tensorflow/python/ops:variables",
|
|
"//tensorflow/python/platform:client_testlib",
|
|
"//third_party/py/numpy",
|
|
"@absl_py//absl/flags",
|
|
],
|
|
)
|
|
|
|
dtensor_test(
|
|
name = "multi_client_test_nccl_local", # Use a suffix for coverage, b/23027507#comment47
|
|
srcs = ["multi_client_test.py"],
|
|
additional_backends = [
|
|
GPU_2DEVS_BACKEND,
|
|
],
|
|
args = [
|
|
"--num_clients=0",
|
|
"--num_local_devices=2",
|
|
"--model_dim_size=2",
|
|
],
|
|
disable = ALL_BACKENDS,
|
|
env = {
|
|
"DTENSOR_GPU_USE_NCCL_COMMUNICATION": "1",
|
|
"NCCL_P2P_DISABLE": "1", # FIXME(b/251183104): p2p detection in cuda 10.1+ is broken.
|
|
},
|
|
tags = [
|
|
"no_windows",
|
|
"nosan", # b/195537906
|
|
],
|
|
deps = [
|
|
":multi_client_test_util",
|
|
":test_util",
|
|
"//tensorflow/dtensor/python:accelerator_util",
|
|
"//tensorflow/dtensor/python:api",
|
|
"//tensorflow/dtensor/python:config",
|
|
"//tensorflow/dtensor/python:d_variable",
|
|
"//tensorflow/dtensor/python:layout",
|
|
"//tensorflow/dtensor/python:mesh_util",
|
|
"//tensorflow/python/eager:backprop",
|
|
"//tensorflow/python/eager/polymorphic_function",
|
|
"//tensorflow/python/framework:config",
|
|
"//tensorflow/python/framework:dtypes",
|
|
"//tensorflow/python/framework:ops",
|
|
"//tensorflow/python/ops:array_ops",
|
|
"//tensorflow/python/ops:math_ops",
|
|
"//tensorflow/python/ops:nn_ops",
|
|
"//tensorflow/python/ops:stateless_random_ops",
|
|
"//tensorflow/python/ops:variables",
|
|
"//tensorflow/python/platform:client_testlib",
|
|
"//third_party/py/numpy",
|
|
"@absl_py//absl/flags",
|
|
],
|
|
)
|
|
|
|
dtensor_test(
|
|
name = "multi_client_test_nccl", # Use a suffix for coverage, b/23027507#comment47
|
|
srcs = ["multi_client_test.py"],
|
|
additional_backends = [
|
|
GPU_2DEVS_BACKEND,
|
|
],
|
|
args = [
|
|
"--num_clients=2",
|
|
"--num_local_devices=1",
|
|
"--model_dim_size=2",
|
|
],
|
|
disable = ALL_BACKENDS,
|
|
env = {
|
|
"DTENSOR_GPU_USE_NCCL_COMMUNICATION": "1",
|
|
"NCCL_P2P_DISABLE": "1", # FIXME(b/251183104): p2p detection in cuda 10.1+ is broken.
|
|
"NCCL_CUMEM_HOST_ENABLE": "0", # TODO(b/414413290): workaround for cuda 12.8.1 failure.
|
|
},
|
|
tags = [
|
|
"no_windows",
|
|
"nosan", # b/195537906
|
|
],
|
|
deps = [
|
|
":multi_client_test_util",
|
|
":test_util",
|
|
"//tensorflow/dtensor/python:accelerator_util",
|
|
"//tensorflow/dtensor/python:api",
|
|
"//tensorflow/dtensor/python:config",
|
|
"//tensorflow/dtensor/python:d_variable",
|
|
"//tensorflow/dtensor/python:layout",
|
|
"//tensorflow/dtensor/python:mesh_util",
|
|
"//tensorflow/python/eager:backprop",
|
|
"//tensorflow/python/eager/polymorphic_function",
|
|
"//tensorflow/python/framework:config",
|
|
"//tensorflow/python/framework:dtypes",
|
|
"//tensorflow/python/framework:ops",
|
|
"//tensorflow/python/ops:array_ops",
|
|
"//tensorflow/python/ops:math_ops",
|
|
"//tensorflow/python/ops:nn_ops",
|
|
"//tensorflow/python/ops:stateless_random_ops",
|
|
"//tensorflow/python/ops:variables",
|
|
"//tensorflow/python/platform:client_testlib",
|
|
"//third_party/py/numpy",
|
|
"@absl_py//absl/flags",
|
|
],
|
|
)
|
|
|
|
dtensor_test(
|
|
name = "multi_mesh_test",
|
|
srcs = ["multi_mesh_test.py"],
|
|
disable_tfrt = [
|
|
"gpu",
|
|
"tpu",
|
|
], # TODO(b/192095157)
|
|
shard_count = {
|
|
"cpu": 5,
|
|
"gpu": 5,
|
|
"tpu": 10,
|
|
TPU_V3_DONUT_BACKEND: 10,
|
|
},
|
|
deps = [
|
|
":test_util",
|
|
"//tensorflow/dtensor/python:accelerator_util",
|
|
"//tensorflow/dtensor/python:api",
|
|
"//tensorflow/dtensor/python:config",
|
|
"//tensorflow/dtensor/python:d_variable",
|
|
"//tensorflow/dtensor/python:layout",
|
|
"//tensorflow/python/eager:backprop",
|
|
"//tensorflow/python/eager:context",
|
|
"//tensorflow/python/eager/polymorphic_function",
|
|
"//tensorflow/python/framework:constant_op",
|
|
"//tensorflow/python/framework:dtypes",
|
|
"//tensorflow/python/framework:errors",
|
|
"//tensorflow/python/framework:ops",
|
|
"//tensorflow/python/ops:array_ops",
|
|
"//tensorflow/python/ops:math_ops",
|
|
"//tensorflow/python/ops:math_ops_gen",
|
|
"//tensorflow/python/platform:client_testlib",
|
|
"//third_party/py/numpy",
|
|
"@absl_py//absl/testing:parameterized",
|
|
],
|
|
)
|
|
|
|
dtensor_test(
|
|
name = "numpy_util_test",
|
|
srcs = ["numpy_util_test.py"],
|
|
main = "numpy_util_test.py",
|
|
deps = [
|
|
":test_util",
|
|
"//tensorflow/dtensor/python:layout",
|
|
"//tensorflow/dtensor/python:numpy_util",
|
|
"//tensorflow/python/platform:client_testlib",
|
|
"//third_party/py/numpy",
|
|
],
|
|
)
|
|
|
|
dtensor_test(
|
|
name = "xla_spmd_test",
|
|
srcs = ["spmd_test.py"],
|
|
additional_backends = [
|
|
TPU_V3_DONUT_BACKEND,
|
|
],
|
|
disable = ALL_BACKENDS,
|
|
env = {
|
|
"DTENSOR_TEST_USE_XLA_SPMD": "1",
|
|
},
|
|
main = "spmd_test.py",
|
|
shard_count = {
|
|
TPU_V3_DONUT_BACKEND: 32,
|
|
},
|
|
deps = [
|
|
":test_util",
|
|
"//tensorflow/dtensor/python:api",
|
|
"//tensorflow/dtensor/python:d_variable",
|
|
"//tensorflow/dtensor/python:layout",
|
|
"//tensorflow/dtensor/python:numpy_util",
|
|
"//tensorflow/python/eager:backprop",
|
|
"//tensorflow/python/eager:context",
|
|
"//tensorflow/python/eager/polymorphic_function",
|
|
"//tensorflow/python/framework:config",
|
|
"//tensorflow/python/framework:constant_op",
|
|
"//tensorflow/python/framework:dtypes",
|
|
"//tensorflow/python/framework:errors",
|
|
"//tensorflow/python/framework:ops",
|
|
"//tensorflow/python/ops:array_ops",
|
|
"//tensorflow/python/ops:array_ops_gen",
|
|
"//tensorflow/python/ops:array_ops_stack",
|
|
"//tensorflow/python/ops:bitwise_ops_gen",
|
|
"//tensorflow/python/ops:io_ops_gen",
|
|
"//tensorflow/python/ops:linalg_ops_gen",
|
|
"//tensorflow/python/ops:list_ops_gen",
|
|
"//tensorflow/python/ops:math_ops",
|
|
"//tensorflow/python/ops:math_ops_gen",
|
|
"//tensorflow/python/ops:nn_ops",
|
|
"//tensorflow/python/ops:nn_ops_gen",
|
|
"//tensorflow/python/ops:random_ops",
|
|
"//tensorflow/python/ops:resource_variable_ops_gen",
|
|
"//tensorflow/python/ops:special_math_ops",
|
|
"//tensorflow/python/ops:spectral_ops_gen",
|
|
"//tensorflow/python/ops:stateless_random_ops",
|
|
"//tensorflow/python/ops:stateless_random_ops_gen",
|
|
"//tensorflow/python/ops:string_ops_gen",
|
|
"//tensorflow/python/ops:variables",
|
|
"//tensorflow/python/platform:client_testlib",
|
|
"//tensorflow/python/util:nest",
|
|
"//third_party/py/numpy",
|
|
"@absl_py//absl/logging",
|
|
"@absl_py//absl/testing:parameterized",
|
|
],
|
|
)
|
|
|
|
dtensor_test(
|
|
name = "multi_device_spmd_test",
|
|
srcs = ["spmd_test.py"],
|
|
additional_backends = [
|
|
TPU_V3_DONUT_BACKEND,
|
|
],
|
|
env = {
|
|
"DTENSOR_ENABLE_MULTI_DEVICE_EXPANSION": "1",
|
|
},
|
|
main = "spmd_test.py",
|
|
shard_count = {
|
|
"cpu": 25,
|
|
"gpu": 10,
|
|
"tpu": 10,
|
|
TPU_V3_DONUT_BACKEND: 32,
|
|
},
|
|
deps = [
|
|
":test_util",
|
|
"//tensorflow/dtensor/python:api",
|
|
"//tensorflow/dtensor/python:d_variable",
|
|
"//tensorflow/dtensor/python:layout",
|
|
"//tensorflow/dtensor/python:numpy_util",
|
|
"//tensorflow/python/eager:backprop",
|
|
"//tensorflow/python/eager:context",
|
|
"//tensorflow/python/eager/polymorphic_function",
|
|
"//tensorflow/python/framework:config",
|
|
"//tensorflow/python/framework:constant_op",
|
|
"//tensorflow/python/framework:dtypes",
|
|
"//tensorflow/python/framework:errors",
|
|
"//tensorflow/python/framework:ops",
|
|
"//tensorflow/python/ops:array_ops",
|
|
"//tensorflow/python/ops:array_ops_gen",
|
|
"//tensorflow/python/ops:array_ops_stack",
|
|
"//tensorflow/python/ops:bitwise_ops_gen",
|
|
"//tensorflow/python/ops:io_ops_gen",
|
|
"//tensorflow/python/ops:linalg_ops_gen",
|
|
"//tensorflow/python/ops:list_ops_gen",
|
|
"//tensorflow/python/ops:math_ops",
|
|
"//tensorflow/python/ops:math_ops_gen",
|
|
"//tensorflow/python/ops:nn_ops",
|
|
"//tensorflow/python/ops:nn_ops_gen",
|
|
"//tensorflow/python/ops:random_ops",
|
|
"//tensorflow/python/ops:resource_variable_ops_gen",
|
|
"//tensorflow/python/ops:special_math_ops",
|
|
"//tensorflow/python/ops:spectral_ops_gen",
|
|
"//tensorflow/python/ops:stateless_random_ops",
|
|
"//tensorflow/python/ops:stateless_random_ops_gen",
|
|
"//tensorflow/python/ops:string_ops_gen",
|
|
"//tensorflow/python/ops:variables",
|
|
"//tensorflow/python/platform:client_testlib",
|
|
"//tensorflow/python/util:nest",
|
|
"//third_party/py/numpy",
|
|
"@absl_py//absl/testing:parameterized",
|
|
],
|
|
)
|
|
|
|
dtensor_test(
|
|
name = "spmd_test",
|
|
srcs = ["spmd_test.py"],
|
|
additional_backends = [TPU_V3_DONUT_BACKEND],
|
|
main = "spmd_test.py",
|
|
shard_count = {
|
|
"cpu": 25,
|
|
"gpu": 10,
|
|
"tpu": 10,
|
|
TPU_V3_DONUT_BACKEND: 32,
|
|
},
|
|
tags = [
|
|
"no_oss_py38", # TODO(b/267017937)
|
|
],
|
|
deps = [
|
|
":test_util",
|
|
"//tensorflow/dtensor/python:api",
|
|
"//tensorflow/dtensor/python:d_variable",
|
|
"//tensorflow/dtensor/python:layout",
|
|
"//tensorflow/dtensor/python:numpy_util",
|
|
"//tensorflow/python/eager:backprop",
|
|
"//tensorflow/python/eager:context",
|
|
"//tensorflow/python/eager/polymorphic_function",
|
|
"//tensorflow/python/framework:config",
|
|
"//tensorflow/python/framework:constant_op",
|
|
"//tensorflow/python/framework:dtypes",
|
|
"//tensorflow/python/framework:errors",
|
|
"//tensorflow/python/framework:ops",
|
|
"//tensorflow/python/ops:array_ops",
|
|
"//tensorflow/python/ops:array_ops_gen",
|
|
"//tensorflow/python/ops:array_ops_stack",
|
|
"//tensorflow/python/ops:bitwise_ops_gen",
|
|
"//tensorflow/python/ops:io_ops_gen",
|
|
"//tensorflow/python/ops:linalg_ops_gen",
|
|
"//tensorflow/python/ops:list_ops_gen",
|
|
"//tensorflow/python/ops:math_ops",
|
|
"//tensorflow/python/ops:math_ops_gen",
|
|
"//tensorflow/python/ops:nn_ops",
|
|
"//tensorflow/python/ops:nn_ops_gen",
|
|
"//tensorflow/python/ops:random_ops",
|
|
"//tensorflow/python/ops:resource_variable_ops_gen",
|
|
"//tensorflow/python/ops:special_math_ops",
|
|
"//tensorflow/python/ops:spectral_ops_gen",
|
|
"//tensorflow/python/ops:stateless_random_ops",
|
|
"//tensorflow/python/ops:stateless_random_ops_gen",
|
|
"//tensorflow/python/ops:string_ops_gen",
|
|
"//tensorflow/python/ops:variables",
|
|
"//tensorflow/python/platform:client_testlib",
|
|
"//tensorflow/python/util:nest",
|
|
"//third_party/py/numpy",
|
|
"@absl_py//absl/testing:parameterized",
|
|
],
|
|
)
|
|
|
|
dtensor_test(
|
|
name = "rng_test",
|
|
size = "medium",
|
|
srcs = ["rng_test.py"],
|
|
additional_backends = [TPU_V3_DONUT_BACKEND],
|
|
# Requires at least 8 TPUs to run the tests.
|
|
disable = ["tpu"],
|
|
disable_tfrt = [
|
|
"gpu",
|
|
TPU_V3_DONUT_BACKEND,
|
|
],
|
|
main = "rng_test.py",
|
|
shard_count = {
|
|
"cpu": 20,
|
|
"tpu": 10,
|
|
"gpu": 30,
|
|
TPU_V3_DONUT_BACKEND: 20,
|
|
},
|
|
deps = [
|
|
":test_util",
|
|
"//tensorflow/dtensor/python:api",
|
|
"//tensorflow/dtensor/python:d_variable",
|
|
"//tensorflow/dtensor/python:layout",
|
|
"//tensorflow/dtensor/python:numpy_util",
|
|
"//tensorflow/python/distribute:tpu_strategy",
|
|
"//tensorflow/python/distribute/cluster_resolver/tpu:tpu_cluster_resolver_py",
|
|
"//tensorflow/python/eager:remote",
|
|
"//tensorflow/python/eager/polymorphic_function",
|
|
"//tensorflow/python/framework:constant_op",
|
|
"//tensorflow/python/framework:dtypes",
|
|
"//tensorflow/python/framework:ops",
|
|
"//tensorflow/python/ops:bitwise_ops_gen",
|
|
"//tensorflow/python/ops:math_ops",
|
|
"//tensorflow/python/ops:stateful_random_ops_gen",
|
|
"//tensorflow/python/ops:stateless_random_ops_v2_gen",
|
|
"//tensorflow/python/ops:variables",
|
|
"//tensorflow/python/platform:client_testlib",
|
|
"//tensorflow/python/tpu:device_assignment",
|
|
"@absl_py//absl/testing:parameterized",
|
|
],
|
|
)
|
|
|
|
dtensor_test(
|
|
name = "save_restore_v2_test",
|
|
srcs = ["save_restore_v2_test.py"],
|
|
additional_backends = [
|
|
TPU_V3_DONUT_BACKEND,
|
|
TPU_V4_DONUT_BACKEND,
|
|
],
|
|
main = "save_restore_v2_test.py",
|
|
shard_count = {
|
|
"cpu": 8,
|
|
"gpu": 8,
|
|
TPU_V3_DONUT_BACKEND: 8,
|
|
},
|
|
deps = [
|
|
":test_util",
|
|
"//tensorflow/dtensor/python:api",
|
|
"//tensorflow/dtensor/python:d_variable",
|
|
"//tensorflow/dtensor/python:layout",
|
|
"//tensorflow/dtensor/python:numpy_util",
|
|
"//tensorflow/python/checkpoint",
|
|
"//tensorflow/python/checkpoint:checkpoint_management",
|
|
"//tensorflow/python/framework:constant_op",
|
|
"//tensorflow/python/framework:dtypes",
|
|
"//tensorflow/python/module",
|
|
"//tensorflow/python/ops:array_ops",
|
|
"//tensorflow/python/ops:math_ops",
|
|
"//tensorflow/python/ops:stateless_random_ops",
|
|
"//tensorflow/python/platform:client_testlib",
|
|
"//third_party/py/numpy",
|
|
"@absl_py//absl/testing:parameterized",
|
|
],
|
|
)
|
|
|
|
dtensor_test(
|
|
name = "variable_test",
|
|
srcs = ["variable_test.py"],
|
|
disable_tfrt = [
|
|
"tpu",
|
|
"gpu",
|
|
], # b/198521331 timeout on TFRT TPU.
|
|
main = "variable_test.py",
|
|
tags = [
|
|
"nomultivm",
|
|
],
|
|
deps = [
|
|
":test_util",
|
|
"//tensorflow/dtensor/python:api",
|
|
"//tensorflow/dtensor/python:d_variable",
|
|
"//tensorflow/dtensor/python:layout",
|
|
"//tensorflow/python/eager/polymorphic_function",
|
|
"//tensorflow/python/framework:constant_op",
|
|
"//tensorflow/python/framework:dtypes",
|
|
"//tensorflow/python/framework:errors",
|
|
"//tensorflow/python/framework:ops",
|
|
"//tensorflow/python/ops:array_ops",
|
|
"//tensorflow/python/ops:array_ops_stack",
|
|
"//tensorflow/python/ops:math_ops",
|
|
"//tensorflow/python/ops:variables",
|
|
"//tensorflow/python/platform:client_testlib",
|
|
"//third_party/py/numpy",
|
|
],
|
|
)
|
|
|
|
dtensor_test(
|
|
name = "mnist_test",
|
|
size = "large",
|
|
srcs = ["mnist_test.py"],
|
|
shard_count = {
|
|
"tpu": 2,
|
|
},
|
|
tags = ["nosan"], # Non-opt builds has slow XLA compilation.
|
|
deps = [
|
|
":test_util",
|
|
"//tensorflow/dtensor/python:api",
|
|
"//tensorflow/dtensor/python:d_variable",
|
|
"//tensorflow/dtensor/python:input_util",
|
|
"//tensorflow/dtensor/python:layout",
|
|
"//tensorflow/python/data/ops:dataset_ops",
|
|
"//tensorflow/python/eager:backprop",
|
|
"//tensorflow/python/eager/polymorphic_function",
|
|
"//tensorflow/python/framework:dtypes",
|
|
"//tensorflow/python/framework:ops",
|
|
"//tensorflow/python/ops:array_ops",
|
|
"//tensorflow/python/ops:array_ops_stack",
|
|
"//tensorflow/python/ops:math_ops",
|
|
"//tensorflow/python/ops:nn_ops",
|
|
"//tensorflow/python/ops:stateless_random_ops",
|
|
"//tensorflow/python/ops:variables",
|
|
"//tensorflow/python/platform:client_testlib",
|
|
"//third_party/py/numpy",
|
|
"@absl_py//absl/testing:parameterized",
|
|
],
|
|
)
|
|
|
|
dtensor_test(
|
|
name = "numerics_test",
|
|
srcs = ["numerics_test.py"],
|
|
additional_backends = [TPU_V3_DONUT_BACKEND],
|
|
disable = ALL_BACKENDS,
|
|
enable = [
|
|
"tpu",
|
|
],
|
|
deps = [
|
|
":test_util",
|
|
"//tensorflow/dtensor/python:accelerator_util",
|
|
"//tensorflow/dtensor/python:layout",
|
|
"//tensorflow/dtensor/python:numpy_util",
|
|
"//tensorflow/python/eager/polymorphic_function",
|
|
"//tensorflow/python/framework:dtypes",
|
|
"//tensorflow/python/ops:math_ops",
|
|
"//tensorflow/python/ops:stateless_random_ops",
|
|
"//tensorflow/python/platform:client_testlib",
|
|
"//third_party/py/numpy",
|
|
"@absl_py//absl/testing:parameterized",
|
|
],
|
|
)
|
|
|
|
dtensor_test(
|
|
name = "sparse_test",
|
|
srcs = ["sparse_test.py"],
|
|
main = "sparse_test.py",
|
|
shard_count = {
|
|
"cpu": 4,
|
|
},
|
|
deps = [
|
|
":test_util",
|
|
"//tensorflow/dtensor/python:layout",
|
|
"//tensorflow/dtensor/python:numpy_util",
|
|
"//tensorflow/python/eager/polymorphic_function",
|
|
"//tensorflow/python/framework:dtypes",
|
|
"//tensorflow/python/ops:array_ops",
|
|
"//tensorflow/python/ops:math_ops",
|
|
"//tensorflow/python/platform:client_testlib",
|
|
"//third_party/py/numpy",
|
|
"@absl_py//absl/testing:parameterized",
|
|
],
|
|
)
|
|
|
|
dtensor_test(
|
|
name = "tpu_device_assignment_test",
|
|
srcs = ["tpu_device_assignment_test.py"],
|
|
disable = ALL_BACKENDS,
|
|
enable = [
|
|
"tpu",
|
|
],
|
|
deps = [
|
|
":test_util",
|
|
"//tensorflow/dtensor/python:layout",
|
|
"//tensorflow/dtensor/python:numpy_util",
|
|
"//tensorflow/dtensor/python:tpu_util",
|
|
"//tensorflow/python/framework:constant_op",
|
|
"//tensorflow/python/framework:dtypes",
|
|
"//tensorflow/python/ops:math_ops",
|
|
"//tensorflow/python/platform:client_testlib",
|
|
],
|
|
)
|