diff --git a/tensorflow/core/kernels/data/iterator_ops.cc b/tensorflow/core/kernels/data/iterator_ops.cc index f19de311703..68bf172268d 100644 --- a/tensorflow/core/kernels/data/iterator_ops.cc +++ b/tensorflow/core/kernels/data/iterator_ops.cc @@ -269,6 +269,7 @@ IteratorHandleOp::IteratorHandleOp(OpKernelConstruction* ctx) OP_REQUIRES_OK(ctx, ctx->GetAttr(kOutputTypes, &output_dtypes_)); OP_REQUIRES_OK(ctx, ctx->GetAttr(kOutputShapes, &output_shapes_)); OP_REQUIRES_OK(ctx, ctx->GetAttr("shared_name", &name_)); + OP_REQUIRES_OK(ctx, ctx->GetAttr("recoverable", &recoverable_)); } // The resource is deleted from the resource manager only when it is private @@ -308,7 +309,11 @@ void IteratorHandleOp::Compute(OpKernelContext* context) LOCKS_EXCLUDED(mu_) { } ResourceMgr* mgr = context->resource_manager(); - OP_REQUIRES_OK(context, cinfo_.Init(mgr, def(), true)); + if (recoverable_ == false) { + OP_REQUIRES_OK(context, cinfo_.Init(mgr, def(), false)); + } else { + OP_REQUIRES_OK(context, cinfo_.Init(mgr, def(), true)); + } IteratorResource* resource; OP_REQUIRES_OK( diff --git a/tensorflow/core/kernels/data/iterator_ops.h b/tensorflow/core/kernels/data/iterator_ops.h index 07b88d4ccc3..7277a8ac652 100644 --- a/tensorflow/core/kernels/data/iterator_ops.h +++ b/tensorflow/core/kernels/data/iterator_ops.h @@ -135,6 +135,7 @@ class IteratorHandleOp : public OpKernel { std::vector output_shapes_; const int graph_def_version_; string name_; + bool recoverable_; }; // Like IteratorHandleOp, but creates handles which are never shared, and does diff --git a/tensorflow/core/kernels/data/multi_device_iterator_ops.cc b/tensorflow/core/kernels/data/multi_device_iterator_ops.cc index c78f28de2ba..7a538d77d1b 100644 --- a/tensorflow/core/kernels/data/multi_device_iterator_ops.cc +++ b/tensorflow/core/kernels/data/multi_device_iterator_ops.cc @@ -442,7 +442,7 @@ class MultiDeviceIteratorHandleOp : public OpKernel { std::unique_ptr function_handle_cache = absl::make_unique(flr); ResourceMgr* mgr = context->resource_manager(); - OP_REQUIRES_OK(context, cinfo_.Init(mgr, def(), true)); + OP_REQUIRES_OK(context, cinfo_.Init(mgr, def())); MultiDeviceIterator* resource; diff --git a/tensorflow/core/ops/dataset_ops.cc b/tensorflow/core/ops/dataset_ops.cc index 3ed48c17224..6fe19a1471d 100644 --- a/tensorflow/core/ops/dataset_ops.cc +++ b/tensorflow/core/ops/dataset_ops.cc @@ -555,13 +555,24 @@ REGISTER_OP("Iterator") .Attr("output_shapes: list(shape) >= 1") .SetShapeFn(shape_inference::ScalarShape); +#ifndef TF_API_COMPATIBLE_1150 REGISTER_OP("IteratorV2") .Output("handle: resource") .Attr("shared_name: string") .Attr("container: string") + .Attr("recoverable: bool = false") .Attr("output_types: list(type) >= 1") .Attr("output_shapes: list(shape) >= 1") .SetShapeFn(shape_inference::ScalarShape); +#else +REGISTER_OP("IteratorV2") + .Output("handle: resource") + .Attr("shared_name: string") + .Attr("container: string") + .Attr("output_types: list(type) >= 1") + .Attr("output_shapes: list(shape) >= 1") + .SetShapeFn(shape_inference::ScalarShape); +#endif REGISTER_OP("AnonymousIterator") .Output("handle: resource")