From 1edb79fb352e103cbe9db101bf78308fa1d2b826 Mon Sep 17 00:00:00 2001 From: Ayush Dubey Date: Fri, 14 Feb 2020 14:28:02 -0800 Subject: [PATCH] Skip ScopedAllocatorOptimizer when the op inputs are `_Arg` nodes. The optimizer adds control inputs from the `_ScopedAllocator` node to the op input. But that would cause a dependency from the function body to function inputs. PiperOrigin-RevId: 295227438 Change-Id: I9955aef9b22d6b32f7cfb0fe4b97199b502bf131 --- .../grappler/optimizers/scoped_allocator_optimizer.cc | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/tensorflow/core/grappler/optimizers/scoped_allocator_optimizer.cc b/tensorflow/core/grappler/optimizers/scoped_allocator_optimizer.cc index 818bbc5e57a..ffb1083d64d 100644 --- a/tensorflow/core/grappler/optimizers/scoped_allocator_optimizer.cc +++ b/tensorflow/core/grappler/optimizers/scoped_allocator_optimizer.cc @@ -518,6 +518,11 @@ class UnaryElementwiseRewriter : public ScopedAllocatorOptimizer::Rewriter { // input nodes and mark them for allocation from backing tensor. for (int i = 0; i < inputs.size(); ++i) { auto& nd = inputs[i]; + if (IsArg(*nd.from_node_def)) { + return errors::Internal( + "ScopedAllocatorOptimizer does not work well when the op inputs " + "are _Arg ops; skipping this optimizer for this function"); + } VLOG(2) << "To input " << i << ": " << nd.from_node_def->name() << " add control input " << "^" << sa_name; @@ -530,8 +535,9 @@ class UnaryElementwiseRewriter : public ScopedAllocatorOptimizer::Rewriter { node_map->AddOutput(sa_name, nd.from_node_def->name()); } - // We add control edges in order to delay execution of the ScopedAllocatorOp - // until just before first use in order to conserve memory. + // We add control edges in order to (1) delay execution of the + // ScopedAllocatorOp until just before first use in order to conserve memory + // (2) ensure correctness in the presence of control flow related ops. bool added_delay_edge = false; for (auto& nd : inputs) { std::vector inputs_to_first;