diff options
| author |  Rohan Jain <rohanj@google.com> | 2017-03-23 16:36:51 -0800 |
|---|---|---|
| committer |  TensorFlower Gardener <gardener@tensorflow.org> | 2017-03-23 17:52:19 -0700 |
| commit | b2ac11c0797447eb36d427a0907fb2f05af57fcc (patch) | |
| tree | eafeef30eed4451bc55c421d966b01d2fda96327 /tensorflow/contrib/memory_stats | |
| parent | 494fe43926f8928bb3b10c1729abea9be5113f21 (diff) | |
Internal changes
Change: 151077399
Diffstat (limited to 'tensorflow/contrib/memory_stats')
| -rw-r--r-- | tensorflow/contrib/memory_stats/BUILD | 21 |
1 files changed, 19 insertions, 2 deletions
diff --git a/tensorflow/contrib/memory_stats/BUILD b/tensorflow/contrib/memory_stats/BUILD index bbb78383fa..8b9d30dcfd 100644 --- a/tensorflow/contrib/memory_stats/BUILD +++ b/tensorflow/contrib/memory_stats/BUILD @@ -10,7 +10,9 @@ package(default_visibility = ["//tensorflow:__subpackages__"]) load("//tensorflow:tensorflow.bzl", "tf_custom_op_library") load("//tensorflow:tensorflow.bzl", "tf_gen_op_libs") load("//tensorflow:tensorflow.bzl", "tf_gen_op_wrapper_py") +load("//tensorflow:tensorflow.bzl", "tf_kernel_library") load("//tensorflow:tensorflow.bzl", "cuda_py_test") +load("//tensorflow:tensorflow.bzl", "tf_custom_op_py_library") tf_custom_op_library( name = "python/ops/_memory_stats_ops.so", @@ -20,6 +22,17 @@ tf_custom_op_library( ], ) +tf_kernel_library( + name = "memory_stats_kernels", + srcs = [ + "kernels/memory_stats_ops.cc", + ], + deps = [ + "//tensorflow/core:framework", + ], + alwayslink = 1, +) + tf_gen_op_libs( op_lib_names = ["memory_stats_ops"], deps = [ @@ -32,15 +45,19 @@ tf_gen_op_wrapper_py( deps = [":memory_stats_ops_op_lib"], ) -py_library( +tf_custom_op_py_library( name = "memory_stats_py", srcs = [ "__init__.py", "python/ops/memory_stats_ops.py", ], - data = [ + dso = [ ":python/ops/_memory_stats_ops.so", ], + kernels = [ + ":memory_stats_kernels", + ":memory_stats_ops_op_lib", + ], srcs_version = "PY2AND3", visibility = ["//visibility:public"], deps = [ |