Skip to content
Permalink
Browse files Browse the repository at this point in the history
Prevent stack overflow when FunctionLib in GraphDef has a self-recurs…
…ive function.

It is likely that no recursivity is supported, but we should handle this separately.

PiperOrigin-RevId: 414860329
Change-Id: I02a2270e86282b37362ddd485eeef16fb986a9e0
  • Loading branch information
mihaimaruseac authored and tensorflower-gardener committed Dec 8, 2021
1 parent 7b1eba4 commit 448a161
Showing 1 changed file with 18 additions and 0 deletions.
18 changes: 18 additions & 0 deletions tensorflow/cc/saved_model/loader.cc
Expand Up @@ -25,6 +25,7 @@ limitations under the License.
#include "tensorflow/core/framework/attr_value.pb.h"
#include "tensorflow/core/framework/function.pb.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/framework/op_def.pb.h"
#include "tensorflow/core/framework/tensor.pb.h"
#include "tensorflow/core/lib/io/path.h"
#include "tensorflow/core/lib/monitoring/counter.h"
Expand Down Expand Up @@ -99,6 +100,19 @@ static Status ValidateNode(const NodeDef& node) {
return Status::OK();
}

static Status ValidateFunctionNotRecursive(const FunctionDef& function) {
const auto& function_name = function.signature().name();
for (const auto& node : function.node_def()) {
if (node.op() == function_name) {
return errors::FailedPrecondition(
"Function ", function_name,
" is self recursive and TensorFlow does not support this scenario.");
}
}

return Status::OK();
}

static Status ValidateSavedTensors(const GraphDef& graph_def) {
for (const auto& node : graph_def.node()) {
TF_RETURN_IF_ERROR(ValidateNode(node));
Expand All @@ -110,6 +124,10 @@ static Status ValidateSavedTensors(const GraphDef& graph_def) {
for (const auto& node : function.node_def()) {
TF_RETURN_IF_ERROR(ValidateNode(node));
}

// Also check that there is no recursivity in the library
// TODO(mihaimaruseac): Do more than self-recursivity
TF_RETURN_IF_ERROR(ValidateFunctionNotRecursive(function));
}
}

Expand Down

0 comments on commit 448a161

Please sign in to comment.