diff options
| author |  Avijit <Avijit.Chakraborty@intel.com> | 2018-08-12 16:21:41 -0700 |
|---|---|---|
| committer | Avijit <Avijit.Chakraborty@intel.com> | 2018-08-12 16:21:41 -0700 |
| commit | 9523a98466d16cf01fc76a67b489f1124cf626ac (patch) | |
| tree | bd4c460b67fab60c2fb1a6c56bf22d1cbb5391e6 /tensorflow/core/kernels/data/map_and_batch_dataset_op.cc | |
| parent | 93e950c308071071f35d6dcb35b9f91b8a34876c (diff) | |
| parent | 1a22b0b982fa1a953651b98af8f3cd30542048fd (diff) | |
Merge remote-tracking branch 'upstream/master'
Diffstat (limited to 'tensorflow/core/kernels/data/map_and_batch_dataset_op.cc')
| -rw-r--r-- | tensorflow/core/kernels/data/map_and_batch_dataset_op.cc | 13 |
1 files changed, 7 insertions, 6 deletions
diff --git a/tensorflow/core/kernels/data/map_and_batch_dataset_op.cc b/tensorflow/core/kernels/data/map_and_batch_dataset_op.cc index 004f153af6..51a7fd23a8 100644 --- a/tensorflow/core/kernels/data/map_and_batch_dataset_op.cc +++ b/tensorflow/core/kernels/data/map_and_batch_dataset_op.cc @@ -144,11 +144,12 @@ class MapAndBatchDatasetOp : public UnaryDatasetOpKernel { } protected: - Status AsGraphDefInternal(OpKernelContext* ctx, DatasetGraphDefBuilder* b, + Status AsGraphDefInternal(SerializationContext* ctx, + DatasetGraphDefBuilder* b, Node** output) const override { - TF_RETURN_IF_ERROR(b->AddFunction(ctx, map_fn_.name())); + TF_RETURN_IF_ERROR(b->AddFunction(ctx->flib_def(), map_fn_.name())); Node* input_graph_node = nullptr; - TF_RETURN_IF_ERROR(b->AddParentDataset(ctx, input_, &input_graph_node)); + TF_RETURN_IF_ERROR(b->AddInputDataset(ctx, input_, &input_graph_node)); Node* batch_size_node; TF_RETURN_IF_ERROR(b->AddScalar(batch_size_, &batch_size_node)); Node* num_parallel_calls_node; @@ -232,7 +233,7 @@ class MapAndBatchDatasetOp : public UnaryDatasetOpKernel { cond_var_.wait(l); } CHECK_EQ(num_calls_, 0); - TF_RETURN_IF_ERROR(SaveParent(writer, input_impl_)); + TF_RETURN_IF_ERROR(SaveInput(writer, input_impl_)); TF_RETURN_IF_ERROR( writer->WriteScalar(full_name("call_counter"), call_counter_)); TF_RETURN_IF_ERROR(writer->WriteScalar(full_name("batch_results_size"), @@ -246,7 +247,7 @@ class MapAndBatchDatasetOp : public UnaryDatasetOpKernel { Status RestoreInternal(IteratorContext* ctx, IteratorStateReader* reader) override { mutex_lock l(mu_); - TF_RETURN_IF_ERROR(RestoreParent(ctx, reader, input_impl_)); + TF_RETURN_IF_ERROR(RestoreInput(ctx, reader, input_impl_)); TF_RETURN_IF_ERROR( reader->ReadScalar(full_name("call_counter"), &call_counter_)); int64 batch_results_size; @@ -383,7 +384,7 @@ class MapAndBatchDatasetOp : public UnaryDatasetOpKernel { #undef HANDLE_TYPE default: return errors::InvalidArgument("Unsupported data type: ", - value.dtype()); + DataTypeString(value.dtype())); } return Status::OK(); } |