experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Reuse the rendezvous provided by the OpKernelContext for PartitionedCallOp.

Browse Source 
This will allow send/recv across different tf.functions.

PiperOrigin-RevId: 310567623
Change-Id: I64cc5c3c44fa6c64684570e6f20808c6b193172b
This commit is contained in:
Guangda Lai 2020-05-08 09:08:24 -07:00 committed by TensorFlower Gardener
parent b6c57a6e4b
commit 57d2ac3f07
1 changed files with 1 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
tensorflow/core/kernels/partitioned_function_ops.cc
View File

@ -245,6 +245,7 @@ void PartitionedCallOp::RunFunction(FunctionLibraryRuntime::Handle handle,
run_opts.source_device =
lib->device() == nullptr ? "" : lib->device()->name();
run_opts.allow_dead_tensors = true;
run_opts.rendezvous = ctx->rendezvous();
std::vector<Tensor>* rets = new std::vector<Tensor>;
const string& func_name = func_->name();