diff options
author | Derek Murray <mrry@google.com> | 2018-01-05 18:45:12 -0800 |
---|---|---|
committer | TensorFlower Gardener <gardener@tensorflow.org> | 2018-01-05 18:51:25 -0800 |
commit | 32d138db751c541e951d1958cac4918214e9644e (patch) | |
tree | 78a2940696bd3162199216bd37a692a979d4062a /tensorflow/core/common_runtime/function_test.cc | |
parent | cfbcbc66a8756d087fd895ff6ec3f5cea32e5157 (diff) |
Add `FunctionLibraryRuntime::InstantiateOptions` struct (take 2).
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).
This version fixes an issue in the previous attempt whereby identical
functions instantiated on different devices in the same task would be
canonicalized to the same key and hence receive the same handle in the
ProcFLR, but they would only be instantiated locally in one per-device
FLR. The fix ensures that the device name is part of the canonicalized
form. A test is added to function_test.py that failed with the old
version and passed in the new version.
END_PUBLIC
BEGIN_PUBLIC
Automated g4 rollback of changelist 180979141
PiperOrigin-RevId: 181003413
Diffstat (limited to 'tensorflow/core/common_runtime/function_test.cc')
-rw-r--r-- | tensorflow/core/common_runtime/function_test.cc | 16 |
1 files changed, 9 insertions, 7 deletions
diff --git a/tensorflow/core/common_runtime/function_test.cc b/tensorflow/core/common_runtime/function_test.cc index d4181ff48c..2dacacea7b 100644 --- a/tensorflow/core/common_runtime/function_test.cc +++ b/tensorflow/core/common_runtime/function_test.cc @@ -191,11 +191,14 @@ class FunctionLibraryRuntimeTest : public ::testing::Test { Status Instantiate(FunctionLibraryRuntime* flr, const string& name, test::function::Attrs attrs, FunctionLibraryRuntime::Handle* handle) { - Status status = flr->Instantiate(name, attrs, handle); - if (!status.ok()) { - return status; - } - return Status::OK(); + return flr->Instantiate(name, attrs, handle); + } + + Status Instantiate(FunctionLibraryRuntime* flr, const string& name, + test::function::Attrs attrs, + const FunctionLibraryRuntime::InstantiateOptions& options, + FunctionLibraryRuntime::Handle* handle) { + return flr->Instantiate(name, attrs, options, handle); } Status InstantiateAndRun(FunctionLibraryRuntime* flr, const string& name, @@ -1088,8 +1091,7 @@ TEST_F(FunctionLibraryRuntimeTest, Gradient_AddSum) { TEST_F(FunctionLibraryRuntimeTest, CrossDevice) { Init({test::function::FindDevice()}); FunctionLibraryRuntime::Handle handle; - TF_CHECK_OK(Instantiate(flr0_, "FindDevice", {{"_target", "/device:CPU:1"}}, - &handle)); + TF_CHECK_OK(Instantiate(flr0_, "FindDevice", {}, {"/device:CPU:1"}, &handle)); Tensor y; FunctionLibraryRuntime::Options opts; |