diff --git a/tensorflow/core/framework/resource_mgr.cc b/tensorflow/core/framework/resource_mgr.cc index ab7dd0c5475..55860d92271 100644 --- a/tensorflow/core/framework/resource_mgr.cc +++ b/tensorflow/core/framework/resource_mgr.cc @@ -24,6 +24,34 @@ limitations under the License. #include "tensorflow/core/platform/demangle.h" namespace tensorflow { +ResourceHandle MakeResourceHandle(OpKernelContext* ctx, const string& container, + const string& name, + const TypeIndex& type_index) { + ResourceHandle result; + result.set_device(ctx->device()->attributes().name()); + string actual_container; + if (!container.empty()) { + actual_container = container; + } else { + actual_container = ctx->resource_manager()->default_container(); + } + result.set_container(actual_container); + result.set_name(name); + result.set_hash_code(type_index.hash_code()); + result.set_maybe_type_name(type_index.name()); + return result; +} + +Status MakeResourceHandleToOutput(OpKernelContext* context, int output_index, + const string& container, const string& name, + const TypeIndex& type_index) { + Tensor* handle; + TF_RETURN_IF_ERROR( + context->allocate_output(output_index, TensorShape({}), &handle)); + handle->scalar()() = + MakeResourceHandle(context, container, name, type_index); + return Status::OK(); +} namespace internal { diff --git a/tensorflow/core/framework/resource_mgr.h b/tensorflow/core/framework/resource_mgr.h index 26a5766569f..0e1a5a82d3f 100644 --- a/tensorflow/core/framework/resource_mgr.h +++ b/tensorflow/core/framework/resource_mgr.h @@ -202,9 +202,20 @@ class ResourceMgr { // Makes a resource handle with the specified type for a given container / // name. +ResourceHandle MakeResourceHandle(OpKernelContext* ctx, const string& container, + const string& name, + const TypeIndex& type_index); + template ResourceHandle MakeResourceHandle(OpKernelContext* ctx, const string& container, - const string& name); + const string& name) { + return MakeResourceHandle(ctx, container, name, MakeTypeIndex()); +} + +Status MakeResourceHandleToOutput(OpKernelContext* context, int output_index, + const string& container, const string& name, + const TypeIndex& type_index); + template ResourceHandle MakePerStepResourceHandle(OpKernelContext* ctx, const string& name); @@ -423,25 +434,6 @@ Status GetResourceFromContext(OpKernelContext* ctx, const string& input_name, return ctx->resource_manager()->Lookup(container, shared_name, resource); } -template -ResourceHandle MakeResourceHandle(OpKernelContext* ctx, const string& container, - const string& name) { - ResourceHandle result; - result.set_device(ctx->device()->attributes().name()); - string actual_container; - if (!container.empty()) { - actual_container = container; - } else { - actual_container = ctx->resource_manager()->default_container(); - } - result.set_container(actual_container); - result.set_name(name); - auto type_index = MakeTypeIndex(); - result.set_hash_code(type_index.hash_code()); - result.set_maybe_type_name(type_index.name()); - return result; -} - template ResourceHandle MakePerStepResourceHandle(OpKernelContext* ctx, const string& name) { diff --git a/tensorflow/core/framework/resource_op_kernel.h b/tensorflow/core/framework/resource_op_kernel.h index de65657a9e5..813ec6eed58 100644 --- a/tensorflow/core/framework/resource_op_kernel.h +++ b/tensorflow/core/framework/resource_op_kernel.h @@ -95,11 +95,9 @@ class ResourceOpKernel : public OpKernel { resource_ = resource; } if (context->expected_output_dtype(0) == DT_RESOURCE) { - Tensor* handle; - OP_REQUIRES_OK(context, - context->allocate_output(0, TensorShape({}), &handle)); - handle->scalar()() = - MakeResourceHandle(context, cinfo_.container(), cinfo_.name()); + OP_REQUIRES_OK(context, MakeResourceHandleToOutput( + context, 0, cinfo_.container(), cinfo_.name(), + MakeTypeIndex())); } else { context->set_output_ref(0, &mu_, handle_.AccessTensor(context)); }