From 5c1a75db4528c14e50d754550e2cf72c390ddc54 Mon Sep 17 00:00:00 2001
From: Zachary Garrett <zachgarrett@google.com>
Date: Mon, 21 Sep 2020 15:09:41 -0700
Subject: [PATCH] [tf.data] Add a pass-thru implementation of `ModelDataset`
 for tf-mobile.

Mobile already has a pass-thru for `OptimizeDataset` and hence doesn't benefit from the modeling.

PiperOrigin-RevId: 332944390
Change-Id: If5b085a7ec77ff003f514e685fc08269185e4f07
---
 tensorflow/core/kernels/data/BUILD            |  1 +
 .../kernels/data/model_dataset_op_mobile.cc   | 37 +++++++++++++++++++
 2 files changed, 38 insertions(+)
 create mode 100644 tensorflow/core/kernels/data/model_dataset_op_mobile.cc

diff --git a/tensorflow/core/kernels/data/BUILD b/tensorflow/core/kernels/data/BUILD
index e9c5430d0fb..438f3a6e1bf 100644
--- a/tensorflow/core/kernels/data/BUILD
+++ b/tensorflow/core/kernels/data/BUILD
@@ -1449,6 +1449,7 @@ filegroup(
         exclude = [
             "dataset_ops*",  # includes grappler dependency, which isn't supported on mobile.
             "optimize_dataset_op.*",  # includes grappler dependency, which isn't supported on mobile.
+            "model_dataset_op.*",  # not supported on mobile.
             "rewrite_utils*",  # includes grappler dependency, which isn't supported on mobile.
             "*test.cc",
             "*test.h",
diff --git a/tensorflow/core/kernels/data/model_dataset_op_mobile.cc b/tensorflow/core/kernels/data/model_dataset_op_mobile.cc
new file mode 100644
index 00000000000..4de0e84130c
--- /dev/null
+++ b/tensorflow/core/kernels/data/model_dataset_op_mobile.cc
@@ -0,0 +1,37 @@
+/* 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.
+==============================================================================*/
+#include "tensorflow/core/framework/dataset.h"
+
+namespace tensorflow {
+namespace data {
+
+class ModelDatasetOp : public UnaryDatasetOpKernel {
+ public:
+  explicit ModelDatasetOp(OpKernelConstruction* ctx)
+      : UnaryDatasetOpKernel(ctx) {}
+
+  void MakeDataset(OpKernelContext* ctx, DatasetBase* input,
+                   DatasetBase** output) {
+    input->Ref();
+    *output = input;
+  }
+};
+
+namespace {
+REGISTER_KERNEL_BUILDER(Name("ModelDataset").Device(DEVICE_CPU),
+                        ModelDatasetOp);
+}  // namespace
+}  // namespace data
+}  // namespace tensorflow