Add the tests for BatchDatasetOp

tensorflow · Apr 16, 2019 · b22ebf9 · b22ebf9
1 parent dac2bf1
commit b22ebf9
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 2 deletions.
diff --git a/tensorflow/core/kernels/data/BUILD b/tensorflow/core/kernels/data/BUILD
@@ -178,6 +178,26 @@ tf_kernel_library(
     ],
 )
 
+tf_cc_test(
+    name = "batch_dataset_op_test",
+    size = "small",
+    srcs = ["batch_dataset_op_test.cc"],
+    deps = [
+        ":batch_dataset_op",
+        ":dataset_test_base",
+        ":dataset_utils",
+        ":iterator_ops",
+        ":range_dataset_op",
+        "//tensorflow/core:core_cpu_internal",
+        "//tensorflow/core:dataset_ops_op_lib",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:lib_internal",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+        "//tensorflow/core:testlib",
+    ],
+)
+
 tf_kernel_library(
     name = "shard_dataset_op",
     srcs = ["shard_dataset_op.cc"],

diff --git a/tensorflow/core/kernels/data/batch_dataset_op.cc b/tensorflow/core/kernels/data/batch_dataset_op.cc
@@ -114,8 +114,14 @@ class BatchDatasetOp : public UnaryDatasetOpKernel {
       TF_RETURN_IF_ERROR(b->AddScalar(batch_size_, &batch_size));
       Node* drop_remainder = nullptr;
       TF_RETURN_IF_ERROR(b->AddScalar(drop_remainder_, &drop_remainder));
-      TF_RETURN_IF_ERROR(b->AddDataset(
-          this, {input_graph_node, batch_size, drop_remainder}, output));
+      if (type_string() == "BatchDataset") {
+        TF_RETURN_IF_ERROR(
+            b->AddDataset(this, {input_graph_node, batch_size}, output));
+      } else {
+        TF_RETURN_IF_ERROR(b->AddDataset(
+            this, {input_graph_node, batch_size, drop_remainder}, output));
+      }
+
       return Status::OK();
     }