1 files changed, 26 insertions, 14 deletions
diff --git a/tensorflow/contrib/data/python/ops/BUILD b/tensorflow/contrib/data/python/ops/BUILD
index 1b81cf5be9..727c5d1c38 100644
--- a/tensorflow/contrib/data/python/ops/BUILD
+++ b/tensorflow/contrib/data/python/ops/BUILD
@@ -12,20 +12,6 @@ load(
 load("//tensorflow:tensorflow.bzl", "tf_custom_op_py_library")
 
 py_library(
-    name = "dataset_ops",
-    srcs = [
-        "dataset_ops.py",
-    ],
-    srcs_version = "PY2AND3",
-    deps = [
-        ":transformation_ops",
-        "//tensorflow/python:util",
-        "//tensorflow/python/data/ops:dataset_ops",
-        "//tensorflow/python/data/util:nest",
-    ],
-)
-
-py_library(
     name = "iterator_ops",
     srcs = [
         "iterator_ops.py",
@@ -73,6 +59,7 @@ py_library(
     ],
     srcs_version = "PY2AND3",
     deps = [
+        ":gen_dataset_ops",
         "//tensorflow/python:array_ops",
         "//tensorflow/python:control_flow_ops",
         "//tensorflow/python:dataset_ops_gen",
@@ -128,6 +115,31 @@ tf_custom_op_py_library(
     ],
 )
 
+tf_gen_op_wrapper_py(
+    name = "gen_dataset_ops",
+    out = "gen_dataset_ops.py",
+    deps = ["//tensorflow/contrib/data:dataset_ops_op_lib"],
+)
+
+tf_custom_op_py_library(
+    name = "dataset_ops",
+    srcs = ["dataset_ops.py"],
+    dso = ["//tensorflow/contrib/data:_dataset_ops.so"],
+    kernels = [
+        "//tensorflow/contrib/data:dataset_ops_op_lib",
+    ],
+    srcs_version = "PY2AND3",
+    deps = [
+        ":gen_dataset_ops",
+        ":transformation_ops",
+        "//tensorflow/contrib/util:util_py",
+        "//tensorflow/python:platform",
+        "//tensorflow/python:util",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/util:nest",
+    ],
+)
+
 filegroup(
     name = "all_files",
     srcs = glob(