Automated g4 rollback of changelist 189231636

PiperOrigin-RevId: 189258641

5 files changed, 46 insertions, 143 deletions

diff --git a/tensorflow/contrib/kafka/BUILD b/tensorflow/contrib/kafka/BUILD
index 1c3974871c..efb403462a 100644
--- a/tensorflow/contrib/kafka/BUILD
+++ b/tensorflow/contrib/kafka/BUILD
@@ -1,93 +1,66 @@
-package(default_visibility = ["//tensorflow:internal"])
+package(
+    default_visibility = ["//visibility:private"],
+)
 
 licenses(["notice"])  # Apache 2.0
 
 exports_files(["LICENSE"])
 
-load(
-    "//tensorflow:tensorflow.bzl",
-    "tf_gen_op_wrapper_py",
-    "tf_kernel_library",
-    "tf_custom_op_library",
-    "tf_custom_op_py_library",
-    "tf_gen_op_libs",
-    "tf_py_test",
-)
-
-py_library(
-    name = "kafka",
-    srcs = ["__init__.py"],
-    srcs_version = "PY2AND3",
-    deps = [
-        ":dataset_ops",
-    ],
-)
-
-tf_custom_op_library(
-    name = "_dataset_ops.so",
-    srcs = ["ops/dataset_ops.cc"],
-    deps = [":dataset_kernels"],
-)
-
-tf_gen_op_libs(
-    op_lib_names = ["dataset_ops"],
-)
+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", "tf_py_test")
 
-cc_library(
-    name = "dataset_kernels",
+tf_kernel_library(
+    name = "kafka_kernels",
     srcs = ["kernels/kafka_dataset_ops.cc"],
+    visibility = ["//visibility:public"],
     deps = [
-        "//tensorflow/core:framework_headers_lib",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:lib_internal",
+        "//tensorflow/core/kernels:bounds_check_lib",
+        "//tensorflow/core/kernels:dataset",
         "//third_party/eigen3",
         "@kafka",
-        "@protobuf_archive//:protobuf_headers",
     ],
-    alwayslink = 1,
 )
 
-py_library(
-    name = "dataset_ops",
-    srcs = [
-        "python/ops/kafka_dataset_ops.py",
-    ],
-    srcs_version = "PY2AND3",
+tf_gen_op_libs(
+    op_lib_names = ["kafka_ops"],
     deps = [
-        ":kafka_op_loader",
-        "//tensorflow/python:dataset_ops_gen",
-        "//tensorflow/python:util",
-        "//tensorflow/python/data/ops:dataset_ops",
-        "//tensorflow/python/data/util:nest",
+        "//tensorflow/core:lib",
     ],
 )
 
 tf_gen_op_wrapper_py(
-    name = "gen_dataset_ops",
-    out = "python/ops/gen_dataset_ops.py",
-    deps = ["//tensorflow/contrib/kafka:dataset_ops_op_lib"],
-)
-
-tf_kernel_library(
-    name = "dataset_ops_kernels",
-    deps = [
-        ":dataset_kernels",
-        "//tensorflow/core:framework",
-    ],
-    alwayslink = 1,
+    name = "gen_kafka_ops",
+    out = "python/ops/gen_kafka_ops.py",
+    require_shape_functions = True,
+    deps = [":kafka_ops_op_lib"],
 )
 
-tf_custom_op_py_library(
-    name = "kafka_op_loader",
-    srcs = ["python/ops/kafka_op_loader.py"],
-    dso = ["//tensorflow/contrib/kafka:_dataset_ops.so"],
-    kernels = [
-        ":dataset_ops_kernels",
-        "//tensorflow/contrib/kafka:dataset_ops_op_lib",
+py_library(
+    name = "kafka",
+    srcs = [
+        "__init__.py",
+        "python/ops/kafka_dataset_ops.py",
     ],
     srcs_version = "PY2AND3",
+    visibility = ["//visibility:public"],
     deps = [
-        ":gen_dataset_ops",
+        ":gen_kafka_ops",
         "//tensorflow/contrib/util:util_py",
+        "//tensorflow/python:array_ops",
+        "//tensorflow/python:control_flow_ops",
+        "//tensorflow/python:framework",
+        "//tensorflow/python:framework_for_generated_wrappers",
         "//tensorflow/python:platform",
+        "//tensorflow/python:state_ops",
+        "//tensorflow/python:training",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/ops:iterator_ops",
+        "//tensorflow/python/data/ops:readers",
     ],
 )
 
@@ -115,7 +88,6 @@ tf_py_test(
     ],
     tags = [
         "manual",
-        "no_windows",
         "notap",
     ],
 )
@@ -123,9 +95,7 @@ tf_py_test(
 filegroup(
     name = "all_files",
     srcs = glob(
-        include = [
-            "**/*",
-        ],
+        ["**/*"],
         exclude = [
             "**/METADATA",
             "**/OWNERS",
diff --git a/tensorflow/contrib/kafka/kernels/kafka_dataset_ops.cc b/tensorflow/contrib/kafka/kernels/kafka_dataset_ops.cc
index a4cd4a2cc4..88ef5f3571 100644
--- a/tensorflow/contrib/kafka/kernels/kafka_dataset_ops.cc
+++ b/tensorflow/contrib/kafka/kernels/kafka_dataset_ops.cc
@@ -13,7 +13,9 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#include "tensorflow/core/framework/dataset.h"
+#include "tensorflow/core/kernels/dataset.h"
+
+#include "tensorflow/core/framework/tensor.h"
 
 #include "src-cpp/rdkafkacpp.h"
 
diff --git a/tensorflow/contrib/kafka/ops/dataset_ops.cc b/tensorflow/contrib/kafka/ops/dataset_ops.cc
deleted file mode 100644
index 8cdf16103b..0000000000
--- a/tensorflow/contrib/kafka/ops/dataset_ops.cc
+++ /dev/null
@@ -1,44 +0,0 @@
-/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
-
-Licensed under the Apache License, Version 2.0 (the "License");
-you may not use this file except in compliance with the License.
-You may obtain a copy of the License at
-
-    http://www.apache.org/licenses/LICENSE-2.0
-
-Unless required by applicable law or agreed to in writing, software
-distributed under the License is distributed on an "AS IS" BASIS,
-WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-See the License for the specific language governing permissions and
-limitations under the License.
-==============================================================================*/
-
-#include "tensorflow/core/framework/common_shape_fns.h"
-#include "tensorflow/core/framework/op.h"
-#include "tensorflow/core/framework/shape_inference.h"
-
-namespace tensorflow {
-
-REGISTER_OP("KafkaDataset")
-    .Input("topics: string")
-    .Input("servers: string")
-    .Input("group: string")
-    .Input("eof: bool")
-    .Input("timeout: int64")
-    .Output("handle: variant")
-    .SetIsStateful()
-    .SetShapeFn(shape_inference::ScalarShape)
-    .Doc(R"doc(
-Creates a dataset that emits the messages of one or more Kafka topics.
-
-topics: A `tf.string` tensor containing one or more subscriptions,
-  in the format of [topic:partition:offset:length],
-  by default length is -1 for unlimited.
-servers: A list of bootstrap servers.
-group: The consumer group id.
-eof: If True, the kafka reader will stop on EOF.
-timeout: The timeout value for the Kafka Consumer to wait
-  (in millisecond).
-)doc");
-
-}  // namespace tensorflow
diff --git a/tensorflow/contrib/kafka/python/ops/kafka_dataset_ops.py b/tensorflow/contrib/kafka/python/ops/kafka_dataset_ops.py
index a1624614d1..8e51d27a34 100644
--- a/tensorflow/contrib/kafka/python/ops/kafka_dataset_ops.py
+++ b/tensorflow/contrib/kafka/python/ops/kafka_dataset_ops.py
@@ -17,9 +17,8 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-from tensorflow.contrib.kafka.python.ops import kafka_op_loader  # pylint: disable=unused-import
-from tensorflow.contrib.kafka.python.ops import gen_dataset_ops
-from tensorflow.python.data.ops.dataset_ops import Dataset
+from tensorflow.contrib.kafka.python.ops import gen_kafka_ops
+from tensorflow.python.data.ops.readers import Dataset
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import tensor_shape
@@ -59,8 +58,8 @@ class KafkaDataset(Dataset):
         timeout, dtype=dtypes.int64, name="timeout")
 
   def _as_variant_tensor(self):
-    return gen_dataset_ops.kafka_dataset(self._topics, self._servers,
-                                         self._group, self._eof, self._timeout)
+    return gen_kafka_ops.kafka_dataset(self._topics, self._servers, self._group,
+                                       self._eof, self._timeout)
 
   @property
   def output_classes(self):
diff --git a/tensorflow/contrib/kafka/python/ops/kafka_op_loader.py b/tensorflow/contrib/kafka/python/ops/kafka_op_loader.py
deleted file mode 100644
index ec2fdea962..0000000000
--- a/tensorflow/contrib/kafka/python/ops/kafka_op_loader.py
+++ /dev/null
@@ -1,24 +0,0 @@
-# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-"""Python helper for loading kafka ops and kernels."""
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from tensorflow.contrib.util import loader
-from tensorflow.python.platform import resource_loader
-
-_dataset_ops = loader.load_op_library(
-    resource_loader.get_path_to_datafile("../../_dataset_ops.so"))