From 19bbc31eee8b81bce6eb08b7ada539943fac6014 Mon Sep 17 00:00:00 2001 From: Derek Murray Date: Thu, 4 Jan 2018 14:48:09 -0800 Subject: Add `FunctionLibraryRuntime::InstantiateOptions` struct. This new struct allows optional arguments to be passed to the `FunctionLibraryRuntime::Instantiate()` API. The new struct is now used to configure the target device for a function instantiation (instead of an attr). PiperOrigin-RevId: 180848930 --- tensorflow/core/kernels/function_ops.cc | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) (limited to 'tensorflow/core/kernels/function_ops.cc') diff --git a/tensorflow/core/kernels/function_ops.cc b/tensorflow/core/kernels/function_ops.cc index f469f41e06..facac10f66 100644 --- a/tensorflow/core/kernels/function_ops.cc +++ b/tensorflow/core/kernels/function_ops.cc @@ -296,21 +296,19 @@ class RemoteCallOp : public AsyncOpKernel { void ComputeAsync(OpKernelContext* ctx, DoneCallback done) override { const Tensor* target; OP_REQUIRES_OK_ASYNC(ctx, ctx->input("target", &target), done); - AttrValueMap attr_values = func_.attr(); - AttrValue v; const string& target_device = DeviceNameUtils::CanonicalizeDeviceName(target->scalar()()); - v.set_s(target_device); - AddAttr("_target", v, &attr_values); FunctionLibraryRuntime* lib = ctx->function_library(); OP_REQUIRES_ASYNC(ctx, lib != nullptr, errors::Internal("No function library is provided."), done); + AttrValueMap attr_values = func_.attr(); FunctionLibraryRuntime::Handle handle; - OP_REQUIRES_OK_ASYNC( - ctx, lib->Instantiate(func_.name(), AttrSlice(&attr_values), &handle), - done); + OP_REQUIRES_OK_ASYNC(ctx, + lib->Instantiate(func_.name(), AttrSlice(&attr_values), + {target_device}, &handle), + done); OpInputList arguments; OP_REQUIRES_OK_ASYNC(ctx, ctx->input_list("args", &arguments), done); -- cgit v1.2.3