message: [tf.data] Fix the bug when restoring the iterator, the buffered elements are not correctly recorded for ParallelInterleaveDataset.
diff: @@ -1278,6 +1278,7 @@ class ParallelInterleaveDatasetOp::Dataset : public DatasetBase {
  absl::StrCat(kResultsSuffix, "[", i, "][", j, "]"),
  &result->return_values.back()));
  }
  + RecordBufferEnqueue(ctx, result->return_values);
  element->results[i] = std::move(result);
  }
  if (!reader->Contains(iterator_name,
  @@ -1339,6 +1340,9 @@ class ParallelInterleaveDatasetOp::Dataset : public DatasetBase {
  TF_RETURN_IF_ERROR(
  ReadElementsParallel(ctx, reader, size, kCurrentElements, &elements));
  mutex_lock l(*mu_);
  + for (auto& element : current_elements_) {
  + DCHECK(element == nullptr);
  + }
  for (int idx = 0; idx < size; ++idx) {
  current_elements_[idx] = std::move(elements[idx]);
  }
  @@ -1361,6 +1365,9 @@ class ParallelInterleaveDatasetOp::Dataset : public DatasetBase {
  TF_RETURN_IF_ERROR(
  ReadElementsParallel(ctx, reader, size, kFutureElements, &elements));
  mutex_lock l(*mu_);
  + for (auto& element : future_elements_) {
  + DCHECK(element == nullptr);
  + }
  for (int idx = 0; idx < size; ++idx) {
  future_elements_[idx] = std::move(elements[idx]);
  }
  