commit e4281619f160dae0e04fc71c634da609c11ec52e
author Yong Tang <yong.tang.github@outlook.com> Mon May 04 02:57:56 2020 +0000
committer Yong Tang <yong.tang.github@outlook.com> Mon May 04 13:59:39 2020 +0000
tree 31cbdca3388502199979267142ac547f6125e24f
parent ae76544efc2731acb3f1ab401ffa1a35c02f57d6

Add cardinality calculation for Dataset.unbatch() when possible

This PR tries to address the issue raised in 39136 where cardinality
of Dataset.unbatch() was always UNKNOWN, even if it might be known
in certain situations.

This PR add the cardinality calculation in case the input cardinality
is known and the leading dim of the output shape is known.

This PR fixes 39136.

Signed-off-by: Yong Tang <yong.tang.github@outlook.com>

tensorflow/core/kernels/data/experimental/unbatch_dataset_op.cc

diff --git a/tensorflow/core/kernels/data/experimental/unbatch_dataset_op.cc b/tensorflow/core/kernels/data/experimental/unbatch_dataset_op.cc
index 111afa2..9b2434d 100644
--- a/tensorflow/core/kernels/data/experimental/unbatch_dataset_op.cc
+++ b/tensorflow/core/kernels/data/experimental/unbatch_dataset_op.cc
@@ -38,8 +38,14 @@
     explicit Dataset(OpKernelContext* ctx, DatasetBase* input)
         : DatasetBase(DatasetContext(ctx)), input_(input) {
       input_->Ref();
+      known_batch_size_ = -1;
       for (const PartialTensorShape& shape : input->output_shapes()) {
         if (!shape.unknown_rank()) {
+          if (known_batch_size_ < 0) {
+            if (shape.dim_size(0) >= 0) {
+              known_batch_size_ = shape.dim_size(0);
+            }
+          }
           gtl::InlinedVector<int64, 4> partial_dim_sizes;
           for (int i = 1; i < shape.dims(); ++i) {
             partial_dim_sizes.push_back(shape.dim_size(i));
@@ -69,6 +75,17 @@
 
     string DebugString() const override { return "UnbatchDatasetOp::Dataset"; }
 
+    int64 Cardinality() const override {
+      int64 n = input_->Cardinality();
+      if (n == kInfiniteCardinality || n == kUnknownCardinality) {
+        return n;
+      }
+      if (known_batch_size_ > 0) {
+        return n * known_batch_size_;
+      }
+      return kUnknownCardinality;
+    }
+
     Status CheckExternalState() const override {
       return input_->CheckExternalState();
     }
@@ -222,6 +239,7 @@
 
     const DatasetBase* const input_;
     std::vector<PartialTensorShape> shapes_;
+    int64 known_batch_size_;
   };
 };