diff options
Diffstat (limited to 'tensorflow/contrib/tpu/BUILD')
| -rw-r--r-- | tensorflow/contrib/tpu/BUILD | 57 |
1 files changed, 55 insertions, 2 deletions
diff --git a/tensorflow/contrib/tpu/BUILD b/tensorflow/contrib/tpu/BUILD index 16696793bc..643a7cc13a 100644 --- a/tensorflow/contrib/tpu/BUILD +++ b/tensorflow/contrib/tpu/BUILD @@ -16,7 +16,6 @@ package( "//cloud/vmm/testing/tests/tpu:__subpackages__", "//learning/brain:__subpackages__", "//tensorflow:__subpackages__", - "//third_party/cloud_tpu:__subpackages__", ], ) @@ -37,6 +36,7 @@ cc_library( py_library( name = "tpu_estimator", srcs = [ + "python/tpu/error_handling.py", "python/tpu/tpu_config.py", "python/tpu/tpu_context.py", "python/tpu/tpu_estimator.py", @@ -161,12 +161,55 @@ py_library( ) py_library( + name = "keras_support", + srcs = [ + "python/tpu/keras_support.py", + ], + srcs_version = "PY2AND3", + visibility = [ + "//cloud/vmm/testing/tests/tpu:__subpackages__", + "//learning/brain:__subpackages__", + # TODO(b/111651964): Clean special visibility for keras_support. + # + # Note: If you are an end user, please do not add your project to this + # visibility. This feature is experimental, and will be made public + # when ready. + "//third_party/cloud_tpu/models/keras:__subpackages__", + "//tensorflow:__subpackages__", + ], + deps = [ + ":tpu_lib", + ":tpu_py", + "//tensorflow/contrib/cluster_resolver:tpu_cluster_resolver_py", + "//tensorflow/contrib/distribute/python:tpu_strategy", + "//tensorflow/contrib/framework:framework_py", + "//tensorflow/contrib/tpu/proto:compilation_result_proto_py", + "//tensorflow/core:protos_all_py", + "//tensorflow/python:array_ops", + "//tensorflow/python:dtypes", + "//tensorflow/python:framework_ops", + "//tensorflow/python:linalg_ops", + "//tensorflow/python:math_ops", + "//tensorflow/python:platform", + "//tensorflow/python:random_ops", + "//tensorflow/python:session", + "//tensorflow/python:tensor_spec", + "//tensorflow/python:variable_scope", + "//tensorflow/python/data/ops:dataset_ops", + "//tensorflow/python/estimator:model_fn", + "//tensorflow/python/keras:backend", + "//tensorflow/python/keras:engine", + "//tensorflow/python/keras:layers", + "//third_party/py/numpy", + ], +) + +py_library( name = "tpu_lib", srcs = [ "python/tpu/__init__.py", "python/tpu/bfloat16.py", "python/tpu/device_assignment.py", - "python/tpu/keras_support.py", "python/tpu/session_support.py", "python/tpu/topology.py", "python/tpu/tpu.py", @@ -307,3 +350,13 @@ tf_py_test( "//tensorflow/python:framework_test_lib", ], ) + +tf_py_test( + name = "topology_test", + size = "small", + srcs = ["python/tpu/topology_test.py"], + additional_deps = [ + ":tpu", + "//tensorflow/python:framework_test_lib", + ], +) |