#tf-data Remove invocation_results_size from the checkpoint when us…

…ing symbolic checkpointing in ParallelMap

PiperOrigin-RevId: 543570204

tensorflow/core/kernels/data/parallel_map_dataset_op.cc

@@ -33,6 +33,7 @@ limitations under the License.
 #include "tensorflow/core/framework/tensor.h"
 #include "tensorflow/core/lib/core/errors.h"
 #include "tensorflow/core/lib/random/random.h"
+#include "tensorflow/core/platform/status.h"
 #include "tensorflow/core/platform/stringprintf.h"
 #include "tensorflow/core/profiler/lib/traceme.h"
 #include "tensorflow/core/profiler/lib/traceme_encode.h"
@@ -317,8 +318,7 @@ class ParallelMapDatasetOp::Dataset : public DatasetBase {
       TF_RETURN_IF_ERROR(ctx->HandleCheckExternalStateStatus(
           dataset()->captured_func_->CheckExternalState()));
       if (ctx->symbolic_checkpoint()) {
-        return writer->WriteScalar(
-            prefix(), absl::StrCat(kInvocationResults, "_", kSize), 0);
+        return OkStatus();
       }
       mutex_lock l(*mu_);
       // Wait for all in-flight calls to complete.
@@ -357,11 +357,14 @@ class ParallelMapDatasetOp::Dataset : public DatasetBase {
                            IteratorStateReader* reader) override {
       mutex_lock l(*mu_);
       TF_RETURN_IF_ERROR(RestoreInput(ctx, reader, input_impl_));
+      DCHECK(invocation_results_.empty());
+      if (ctx->symbolic_checkpoint()) {
+        return OkStatus();
+      }
       int64_t invocation_results_size;
       TF_RETURN_IF_ERROR(reader->ReadScalar(
           prefix(), absl::StrCat(kInvocationResults, "_", kSize),
           &invocation_results_size));
-      DCHECK(invocation_results_.empty());
       for (size_t i = 0; i < invocation_results_size; i++) {
         invocation_results_.push_back(std::make_shared<InvocationResult>(ctx));
         auto& result = *invocation_results_.back();