"//tensorflow/core:lib_internal",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/grappler:grappler_item",
+ "//tensorflow/core/grappler:op_types",
"//tensorflow/core/grappler:utils",
],
)
"//tensorflow/cc:cc_ops",
"//tensorflow/core:all_kernels",
"//tensorflow/core:framework",
+ "//tensorflow/core:lib",
+ "//tensorflow/core:lib_internal",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/grappler/op_types.h"
#include "tensorflow/core/grappler/utils.h"
+#include "tensorflow/core/lib/gtl/map_util.h"
#include "tensorflow/core/lib/strings/scanner.h"
namespace tensorflow {
namespace grappler {
+namespace {
+
+Status OutputNameRange(const FunctionLibraryDefinition& flib,
+ const NodeDef& node,
+ tensorflow::NameRangeMap* outputs_range_map) {
+ const OpRegistrationData* registration;
+ TF_RETURN_IF_ERROR(flib.LookUp(node.op(), ®istration));
+ TF_RETURN_IF_ERROR(tensorflow::NameRangesForNode(node, registration->op_def,
+ nullptr, outputs_range_map));
+ return Status::OK();
+}
+
+Status RegisterFunctionBodyOutputs(const FunctionLibraryDefinition& flib,
+ const NodeDef& node,
+ GrapplerFunctionConnectivity* connectivity) {
+ tensorflow::NameRangeMap outputs_range_map;
+ TF_RETURN_IF_ERROR(OutputNameRange(flib, node, &outputs_range_map));
+ connectivity->RegisterFunctionBodyOutputs(node.name(), outputs_range_map);
+ return Status::OK();
+}
+
+// Replace the placeholder attribute values with the values specified in
+// instantiation attributes.
+Status ResolveFunctionBodyNodeAttrPlaceholders(
+ const AttrValueMap& func_instantiation_attr, NodeDef* node) {
+ for (auto& attr : *node->mutable_attr()) {
+ const string& placeholder = attr.second.placeholder();
+ if (placeholder.empty()) continue;
+
+ auto it = func_instantiation_attr.find(placeholder);
+ if (it != func_instantiation_attr.end()) {
+ attr.second = it->second;
+ } else {
+ return errors::InvalidArgument("Can't resolve placeholder: ",
+ placeholder);
+ }
+ }
+ return Status::OK();
+}
+
+} // namespace
+
void GrapplerFunctionConnectivity::RegisterInputArgExpansion(
const InputArgExpansion& input_arg_expansion) {
- input_arg_expansions_.insert(
- {input_arg_expansion.input_name, input_arg_expansion});
+ const auto& input_name = input_arg_expansion.input_name;
+ const auto& placeholders = input_arg_expansion.placeholders;
+ input_arg_expansions_.emplace(input_name, input_arg_expansion);
+ for (int i = 0; i < placeholders.size(); ++i) {
+ const string& placeholder = input_arg_expansion.placeholders[i];
+ input_arg_placeholders_.emplace(
+ placeholder, InputArgPlaceholder{input_name, /*position=*/i});
+ }
}
void GrapplerFunctionConnectivity::RegisterFunctionBodyOutputs(
const string& node_name, const tensorflow::NameRangeMap& outputs) {
- function_body_outputs_.insert({node_name, outputs});
+ function_body_outputs_[node_name] = outputs;
}
Status GrapplerFunctionConnectivity::ExpandFunctionDefInput(
const string& func_def_input, std::vector<string>* graph_def_inputs) const {
using ::tensorflow::strings::Scanner;
+ if (IsControlInput(func_def_input)) {
+ graph_def_inputs->push_back(func_def_input);
+ return Status::OK();
+ }
+
// Parse input format: "node_name[:node_output][:position]"
string node_name;
string node_output;
std::vector<string> expanded_inputs;
for (const string& function_def_input : function_body_node->input()) {
- if (!IsControlInput(function_def_input))
- TF_RETURN_IF_ERROR(
- ExpandFunctionDefInput(function_def_input, &expanded_inputs));
- else
- expanded_inputs.push_back(function_def_input);
+ TF_RETURN_IF_ERROR(
+ ExpandFunctionDefInput(function_def_input, &expanded_inputs));
}
function_body_node->clear_input();
return Status::OK();
}
-Status GrapplerFunctionItemBuilder::GetTypeAttr(const string& type_attr_name,
- DataType* data_type) const {
- auto it = func_attr_->find(type_attr_name);
- if (it == func_attr_->end()) {
+Status GrapplerFunctionConnectivity::AsFunctionDefInput(
+ const string& graph_def_input, string* func_def_input) const {
+ using gtl::FindOrNull;
+
+ if (IsControlInput(graph_def_input)) {
+ *func_def_input = graph_def_input;
+ return Status::OK();
+ }
+
+ int position;
+ string node_name = ParseNodeName(graph_def_input, &position);
+ CHECK_GE(position, 0);
+
+ // Check if it's an input arg placeholder
+ if (position == 0) {
+ const InputArgPlaceholder* placeholder =
+ FindOrNull(input_arg_placeholders_, node_name);
+ if (placeholder != nullptr) {
+ *func_def_input =
+ strings::StrCat(placeholder->input_name, ":", placeholder->position);
+ return Status::OK();
+ }
+ }
+
+ // It must be output from one of the function body nodes
+ const tensorflow::NameRangeMap* outputs_range_map =
+ FindOrNull(function_body_outputs_, node_name);
+ if (outputs_range_map != nullptr) {
+ for (const auto& el : *outputs_range_map) {
+ const auto& output_name = el.first;
+ const auto& output_range = el.second;
+ if (position >= output_range.first && position < output_range.second) {
+ int pos = position - output_range.first;
+ *func_def_input =
+ strings::StrCat(node_name, ":", output_name, ":", pos);
+ return Status::OK();
+ }
+ }
+ }
+
+ return errors::InvalidArgument("Unknown graph def input: ", graph_def_input);
+}
+
+Status GrapplerFunctionConnectivity::AsFunctionDefNode(
+ NodeDef* function_body_node) const {
+ string func_def_input;
+
+ for (int i = 0; i < function_body_node->input_size(); ++i) {
+ TF_RETURN_IF_ERROR(
+ AsFunctionDefInput(function_body_node->input(i), &func_def_input));
+ function_body_node->set_input(i, func_def_input);
+ }
+
+ return Status::OK();
+}
+
+Status GrapplerFunctionItemInstantiation::GetTypeAttr(
+ const string& type_attr_name, DataType* data_type) const {
+ auto it = func_instantiation_attr_->find(type_attr_name);
+ if (it == func_instantiation_attr_->end()) {
return errors::InvalidArgument("Type attribute ", type_attr_name,
" is not defined");
} else if (it->second.type() == DT_INVALID) {
return Status::OK();
}
-Status GrapplerFunctionItemBuilder::GetArgType(const OpDef::ArgDef& arg,
- DataType* data_type) const {
+Status GrapplerFunctionItemInstantiation::GetArgType(
+ const OpDef::ArgDef& arg, DataType* data_type) const {
if (arg.type() != DT_INVALID) {
*data_type = arg.type();
} else {
+ if (!arg.type_list_attr().empty() || !arg.number_attr().empty()) {
+ return errors::InvalidArgument(
+ "Arguments with sequence of tensors are not supported. Unsupported "
+ "argument name: ",
+ arg.name());
+ }
TF_RETURN_IF_ERROR(GetTypeAttr(arg.type_attr(), data_type));
}
return Status::OK();
}
GrapplerFunctionItem::GrapplerFunctionItem(
- const string& function_name,
+ const string& func_name, const AttrValueMap& func_attr,
const std::vector<InputArgExpansion>& input_arg_expansions,
const std::vector<OutputArgExpansion>& output_arg_expansions,
GraphDef&& function_body)
- : function_name_(function_name),
+ : func_attr_(func_attr),
input_arg_expansions_(input_arg_expansions),
output_arg_expansions_(output_arg_expansions) {
+ id = func_name;
+ // Fill the feed nodes with input placeholders
+ for (const InputArgExpansion& input_arg : input_arg_expansions_) {
+ for (const string& placeholder : input_arg.placeholders) {
+ feed.emplace_back(placeholder, Tensor());
+ input_arg_placeholders_.insert(placeholder);
+ }
+ }
+ // Fill the fetch nodes with outputs
+ for (const OutputArgExpansion& output_arg : output_arg_expansions_) {
+ for (const string& output_tensor : output_arg.output_tensors) {
+ fetch.push_back(output_tensor);
+ }
+ }
+ // Swap the graph body
graph.Swap(&function_body);
}
-const string& GrapplerFunctionItem::function_name() const {
- return function_name_;
-}
-
const std::vector<InputArgExpansion>& GrapplerFunctionItem::inputs() const {
return input_arg_expansions_;
}
return input_arg_expansions_.size();
}
+bool GrapplerFunctionItem::IsInputPlaceholder(const string& node_name) const {
+ return input_arg_placeholders_.find(node_name) !=
+ input_arg_placeholders_.end();
+}
+
const std::vector<OutputArgExpansion>& GrapplerFunctionItem::outputs() const {
return output_arg_expansions_;
}
return output_arg_expansions_.size();
}
+const AttrValueMap& GrapplerFunctionItem::func_attr() const {
+ return func_attr_;
+}
+
const GraphDef& GrapplerFunctionItem::function_body() const { return graph; }
GraphDef& GrapplerFunctionItem::mutable_function_body() { return graph; }
+GrapplerFunctionItem& GrapplerFunctionItem::SwapFunctionBody(GraphDef&& other) {
+ graph.Swap(&other);
+ return *this;
+}
+
std::vector<string> OutputTensors(const GrapplerFunctionItem& item) {
std::vector<string> output_tensors;
for (const OutputArgExpansion& output : item.outputs()) {
return output_tensors;
}
-Status MakeGrapplerFunctionItem(
- const FunctionDef& func,
- const std::unordered_map<string, AttrValue>& func_attr,
- const FunctionLibraryDefinition& func_library, GrapplerFunctionItem* item) {
+Status MakeGrapplerFunctionItem(const FunctionDef& func,
+ const AttrValueMap& func_instantiation_attr,
+ const FunctionLibraryDefinition& flib,
+ GrapplerFunctionItem* item) {
const OpDef& signature = func.signature();
if (signature.name().empty()) {
return errors::InvalidArgument("Function name must be specified");
}
- // Helper methods to lookup function attributes
- GrapplerFunctionItemBuilder builder(&func_attr);
+ // Function types will be resolved from function instantiation attributes. All
+ // other attributes will be lost during conversion to FunctionDef.
+ for (const OpDef::AttrDef& attr : signature.attr()) {
+ if (attr.type() != "type") {
+ return errors::InvalidArgument(
+ "Function signature must have only type attributes");
+ }
+ }
+
+ // Helper methods to lookup function instantiation attributes
+ GrapplerFunctionItemInstantiation instantiation(&func_instantiation_attr);
// Mapping from FunctionDef input format (name[:output][:position]) to
// GraphDef input format (name[:position])
std::vector<InputArgExpansion> inputs;
std::vector<OutputArgExpansion> outputs;
+
+ // Function body shares the library with the graph that instantiated it.
GraphDef function_body;
+ *function_body.mutable_library() = flib.ToProto();
// TODO(ezhulenev): support functions with tensor sequence inputs/outputs
}
DataType input_data_type;
- TF_RETURN_IF_ERROR(builder.GetArgType(input, &input_data_type));
+ TF_RETURN_IF_ERROR(instantiation.GetArgType(input, &input_data_type));
NodeDef* placeholder = function_body.add_node();
placeholder->set_name(input.name());
(*placeholder->mutable_attr())["T"].set_type(input_data_type);
InputArgExpansion input_expansion{/*input_name=*/input.name(),
+ /*data_type=*/input_data_type,
/*placeholders=*/{input.name()}};
connectivity.RegisterInputArgExpansion(input_expansion);
inputs.push_back(input_expansion);
NodeDef* new_node = function_body.add_node();
*new_node = func_def_node;
- // Replace the placeholder attribute values with the specified value
- for (auto& attr : *new_node->mutable_attr()) {
- const string& ph_name = attr.second.placeholder();
- auto it = func_attr.find(ph_name);
- if (it != func_attr.end()) {
- attr.second = it->second;
- }
- }
-
- // Functions use a custom format to encode connectivity. Map these custom
- // strings to regular ones.
- tensorflow::NameRangeMap outputs_range_map;
- const OpRegistrationData* registration;
- TF_RETURN_IF_ERROR(func_library.LookUp(func_def_node.op(), ®istration));
- TF_RETURN_IF_ERROR(tensorflow::NameRangesForNode(
- func_def_node, registration->op_def, nullptr, &outputs_range_map));
- connectivity.RegisterFunctionBodyOutputs(func_def_node.name(),
- outputs_range_map);
+ // Resolve all placeholder values using function instantiation attributes.
+ TF_RETURN_IF_ERROR(ResolveFunctionBodyNodeAttrPlaceholders(
+ func_instantiation_attr, new_node));
+ // Register node output range in a function connectivity.
+ TF_RETURN_IF_ERROR(
+ RegisterFunctionBodyOutputs(flib, func_def_node, &connectivity));
}
// Rewrite inputs to use GraphDef format
for (const OpDef::ArgDef& out : signature.output_arg()) {
std::vector<string> output_tensors;
auto ret = func.ret().find(out.name());
- if (ret != func.ret().end()) {
- // Expand outputs using provided output mapping
- TF_RETURN_IF_ERROR(
- connectivity.ExpandFunctionDefInput(ret->second, &output_tensors));
- } else {
- // Otherwise output must be one of the function inputs
- TF_RETURN_IF_ERROR(
- connectivity.ExpandFunctionDefInput(out.name(), &output_tensors));
+ TF_RETURN_IF_ERROR(
+ ret != func.ret().end()
+ // Expand outputs using provided output mapping
+ ? connectivity.ExpandFunctionDefInput(ret->second, &output_tensors)
+ // Otherwise output must be one of the function inputs
+ : connectivity.ExpandFunctionDefInput(out.name(), &output_tensors));
+
+ DataType output_data_type;
+ TF_RETURN_IF_ERROR(instantiation.GetArgType(out, &output_data_type));
+
+ OutputArgExpansion output{/*output_name=*/out.name(),
+ /*data_type=*/output_data_type,
+ /*output_tensors=*/output_tensors};
+ outputs.push_back(output);
+ }
+
+ *item = GrapplerFunctionItem(
+ /*func_name=*/signature.name(),
+ /*func_attr=*/AttrValueMap(func.attr().begin(), func.attr().end()),
+ inputs, outputs, std::move(function_body));
+ return Status::OK();
+}
+
+// Register GrapplerFunctionItem input arg expansion and function body outputs
+// in the GrapplerFunctionConnectivity
+Status RegisterGrapplerFunctionConnectivity(
+ const GrapplerFunctionItem& item, const FunctionLibraryDefinition& flib,
+ GrapplerFunctionConnectivity* connectivity) {
+ for (const InputArgExpansion& input : item.inputs()) {
+ connectivity->RegisterInputArgExpansion(input);
+ }
+ for (const NodeDef& func_body_node : item.function_body().node()) {
+ TF_RETURN_IF_ERROR(
+ RegisterFunctionBodyOutputs(flib, func_body_node, connectivity));
+ }
+ return Status::OK();
+}
+
+Status MakeSpecializedFunctionDef(const GrapplerFunctionItem& item,
+ const FunctionLibraryDefinition& flib,
+ FunctionDef* func) {
+ func->mutable_signature()->set_name(item.id);
+
+ // Build a GrapplerFunctionConnectivity from inputs and new function body.
+ GrapplerFunctionConnectivity connectivity;
+ TF_RETURN_IF_ERROR(
+ RegisterGrapplerFunctionConnectivity(item, flib, &connectivity));
+
+ // Add function input arguments.
+ for (const InputArgExpansion& input_arg : item.inputs()) {
+ OpDef::ArgDef arg_def;
+ arg_def.set_name(input_arg.input_name);
+ arg_def.set_type(input_arg.data_type);
+ *func->mutable_signature()->add_input_arg() = arg_def;
+ }
+
+ // Add function output arguments.
+ for (const OutputArgExpansion& output_arg : item.outputs()) {
+ OpDef::ArgDef arg_def;
+ arg_def.set_name(output_arg.output_name);
+ arg_def.set_type(output_arg.data_type);
+ *func->mutable_signature()->add_output_arg() = arg_def;
+
+ CHECK(output_arg.output_tensors.size() == 1) // do some sanity checking
+ << "Outputs of tensor sequences are not supported";
+
+ string ret;
+ for (const string& output_tensor : output_arg.output_tensors) {
+ TF_RETURN_IF_ERROR(connectivity.AsFunctionDefInput(output_tensor, &ret));
+ (*func->mutable_ret())[output_arg.output_name] = ret;
}
- outputs.push_back({out.name(), output_tensors});
}
- *item = GrapplerFunctionItem(signature.name(), inputs, outputs,
- std::move(function_body));
+ // Copy function definition specific attributes.
+ for (const auto& attr : item.func_attr()) {
+ const auto& attr_name = attr.first;
+ const auto& attr_value = attr.second;
+ (*func->mutable_attr())[attr_name] = attr_value;
+ }
+
+ // Copy function body nodes to the FunctionDef and update input format
+ for (const NodeDef& func_body_node : item.function_body().node()) {
+ // Do not copy input placeholders
+ if (item.IsInputPlaceholder(func_body_node.name())) continue;
+
+ NodeDef* func_def_node = func->add_node_def();
+ *func_def_node = func_body_node;
+ TF_RETURN_IF_ERROR(connectivity.AsFunctionDefNode(func_def_node));
+ }
+
return Status::OK();
}
namespace tensorflow {
namespace grappler {
+using AttrValueMap = std::unordered_map<string, AttrValue>;
+
// Depending on the function instantiation attributes, input argument to the
// function might be a single tensor, list of tensors of the same type, or a
// list of tensors of different types.
//
// InputArgExpansion keeps track of the placeholders that were added to the
-// function body in place of function inputs.
+// function body in place of function inputs and a resolved input data type.
struct InputArgExpansion {
+ // TODO(ezhulenev): Add support for functions with tensor sequence inputs of
+ // different data types
string input_name; // name of the function input argument
+ DataType data_type; // input data type
std::vector<string> placeholders; // names of placeholder nodes in the
// function body
};
// to one or more outputs of one of the function body nodes.
//
// OutputArgExpansion keeps mapping from a function output arg to the output
-// tensors of a function body nodes, that compute function outputs.
+// tensors of a function body nodes and a resolved output data type
struct OutputArgExpansion {
+ // TODO(ezhulenev): Add support for functions with tensor sequence outputs of
+ // different data types
string output_name; // name of the function output argument
- std::vector<string> output_tensors; // names of output tensors from the
- // function body graph nodes
+ DataType data_type; // output data type
+ std::vector<string> output_tensors; // names of output tensor from the
+ // function body nodes
};
// FunctionDef uses different connectivity encoding for the function body nodes,
Status ExpandFunctionDefInput(const string& func_def_input,
std::vector<string>* graph_def_inputs) const;
- // Update Node inputs from FunctionDef to GraphDef format
+ // Update Node inputs from FunctionDef to GraphDef format.
Status ExpandNodeInputs(NodeDef* function_body_node) const;
- // TODO(ezhulenev): fold GraphDef inputs back to FunctionDef format
- // Status FoldGraphDefInputs(const std::vector<sting> graph_def_inputs,
- // std::vector<string>* function_def_inputs) const;
+ // When expanding inputs in function def format, single input might be
+ // expanded into multiple tensors. When converting back to the function def
+ // format from graph def format, it's always a 1-to-1 relationship.
+ // FunctionDef built from GrapplerFunctionItem is always specialized to it's
+ // instantiation attributes and length of input args (and node def outputs) is
+ // known.
+
+ // Map from GraphDef input format to FunctionDef input format using registered
+ // input arg expansion and function body outputs.
+ Status AsFunctionDefInput(const string& graph_def_input,
+ string* func_def_input) const;
+
+ // Update Node inputs from GraphDef to FunctionDef format.
+ Status AsFunctionDefNode(NodeDef* function_body_node) const;
private:
+ // Mapping from input name to input arg expansion.
std::unordered_map<string, InputArgExpansion> input_arg_expansions_;
+ // Mapping from function body node name to output names range map.
std::unordered_map<string, tensorflow::NameRangeMap> function_body_outputs_;
+
+ struct InputArgPlaceholder {
+ string input_name;
+ int position;
+ };
+
+ // Mapping from input arg placeholder to the function input tensor.
+ std::unordered_map<string, InputArgPlaceholder> input_arg_placeholders_;
};
-// Helper methods to build GrapplerFunctionItem from a function def and function
-// attributes.
-class GrapplerFunctionItemBuilder {
+// Get Function type attributes using attributes of a node that instantiated
+// a function.
+class GrapplerFunctionItemInstantiation {
public:
- using FunctionAttr = std::unordered_map<string, AttrValue>;
-
- explicit GrapplerFunctionItemBuilder(const FunctionAttr* func_attr)
- : func_attr_(func_attr) {}
+ explicit GrapplerFunctionItemInstantiation(
+ const AttrValueMap* func_instantiation_attr)
+ : func_instantiation_attr_(func_instantiation_attr) {}
// Get DataType from attributes by name. Return error if attribute is missing,
// or it doesn't define a valid data type.
Status GetArgType(const OpDef::ArgDef& arg, DataType* data_type) const;
private:
- const FunctionAttr* func_attr_; // do not own
+ const AttrValueMap* func_instantiation_attr_; // do not own
};
// A special case of GrapplerItem, constructed from a TensorFlow Function.
class GrapplerFunctionItem : public GrapplerItem {
public:
- GrapplerFunctionItem() {}
+ GrapplerFunctionItem() = default;
GrapplerFunctionItem(
- const string& function_name,
+ const string& func_name, const AttrValueMap& func_attr,
const std::vector<InputArgExpansion>& input_arg_expansions,
const std::vector<OutputArgExpansion>& output_arg_expansions,
GraphDef&& function_body);
- const string& function_name() const;
+ bool IsInputPlaceholder(const string& node_name) const;
const std::vector<InputArgExpansion>& inputs() const;
const InputArgExpansion& input(int i) const;
const OutputArgExpansion& output(int i) const;
const std::size_t output_size() const;
+ const AttrValueMap& func_attr() const;
const GraphDef& function_body() const;
GraphDef& mutable_function_body();
+ GrapplerFunctionItem& SwapFunctionBody(GraphDef&& other);
+
private:
- string function_name_;
+ AttrValueMap func_attr_; // Attributes specific to function definition that
+ // produced this item (FuncDef.attr field).
+
std::vector<InputArgExpansion> input_arg_expansions_;
std::vector<OutputArgExpansion> output_arg_expansions_;
+
+ std::set<string> input_arg_placeholders_;
};
// Return all output tensors referenced by item output args.
// Return error if the given function def cannot be converted.
Status MakeGrapplerFunctionItem(
const FunctionDef& func,
- const std::unordered_map<string, AttrValue>& func_attr,
- const FunctionLibraryDefinition& func_library, GrapplerFunctionItem* item);
+ const std::unordered_map<string, AttrValue>& func_instantiation_attr,
+ const FunctionLibraryDefinition& flib, GrapplerFunctionItem* item);
+
+// Register GrapplerFunctionItem input arg expansion and function body outputs
+// in the GrapplerFunctionConnectivity. Use function library definition to
+// lookup function body nodes output names and ranges.
+Status RegisterGrapplerFunctionConnectivity(
+ const GrapplerFunctionItem& item, const FunctionLibraryDefinition& flib,
+ GrapplerFunctionConnectivity* connectivity);
+
+// Make a specialized FunctionDef from the GrapplerFunctionItem. Use function
+// library definition to lookup function body nodes output names and ranges.
+Status MakeSpecializedFunctionDef(const GrapplerFunctionItem& item,
+ const FunctionLibraryDefinition& flib,
+ FunctionDef* func);
} // end namespace grappler
} // end namespace tensorflow
#include "tensorflow/core/framework/node_def_util.h"
#include "tensorflow/core/framework/tensor_testutil.h"
#include "tensorflow/core/lib/core/status_test_util.h"
+#include "tensorflow/core/lib/gtl/map_util.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/protobuf/meta_graph.pb.h"
TEST_F(FunctionsTest, GrapplerFunctionConnectivity_ExpandFunctionDefInput) {
GrapplerFunctionConnectivity connectivity;
- connectivity.RegisterInputArgExpansion({"inputA", {"inputA"}});
- connectivity.RegisterInputArgExpansion({"inputB", {"inputB_0", "inputB_1"}});
+ connectivity.RegisterInputArgExpansion({"inputA", DT_FLOAT, {"inputA"}});
+ connectivity.RegisterInputArgExpansion(
+ {"inputB", DT_FLOAT, {"inputB_0", "inputB_1"}});
connectivity.RegisterFunctionBodyOutputs("Add", {{"z", {0, 1}}});
connectivity.RegisterFunctionBodyOutputs("Func",
EXPECT_EQ("Func:3", inputs[0]);
}
+TEST_F(FunctionsTest, GrapplerFunctionConnectivity_AsFunctionDefInput) {
+ GrapplerFunctionConnectivity connectivity;
+
+ connectivity.RegisterInputArgExpansion({"inputA", DT_FLOAT, {"inputA"}});
+ connectivity.RegisterInputArgExpansion(
+ {"inputB", DT_FLOAT, {"inputB_0", "inputB_1"}});
+
+ connectivity.RegisterFunctionBodyOutputs("Add", {{"z", {0, 1}}});
+ connectivity.RegisterFunctionBodyOutputs("Func",
+ {{"o1", {0, 2}}, {"o2", {2, 4}}});
+
+ string input;
+
+ TF_EXPECT_OK(connectivity.AsFunctionDefInput("inputA", &input));
+ EXPECT_EQ("inputA:0", input);
+
+ TF_EXPECT_OK(connectivity.AsFunctionDefInput("inputB_0", &input));
+ EXPECT_EQ("inputB:0", input);
+
+ TF_EXPECT_OK(connectivity.AsFunctionDefInput("inputB_1", &input));
+ EXPECT_EQ("inputB:1", input);
+
+ TF_EXPECT_OK(connectivity.AsFunctionDefInput("Add", &input));
+ EXPECT_EQ("Add:z:0", input);
+
+ TF_EXPECT_OK(connectivity.AsFunctionDefInput("Func", &input));
+ EXPECT_EQ("Func:o1:0", input);
+
+ TF_EXPECT_OK(connectivity.AsFunctionDefInput("Func:1", &input));
+ EXPECT_EQ("Func:o1:1", input);
+
+ TF_EXPECT_OK(connectivity.AsFunctionDefInput("Func:2", &input));
+ EXPECT_EQ("Func:o2:0", input);
+
+ TF_EXPECT_OK(connectivity.AsFunctionDefInput("Func:3", &input));
+ EXPECT_EQ("Func:o2:1", input);
+}
+
TEST_F(FunctionsTest, GrapplerFunctionConnectivity_ExpandNodeInputs) {
GrapplerFunctionConnectivity connectivity;
- connectivity.RegisterInputArgExpansion({"inputA", {"inputA"}});
- connectivity.RegisterInputArgExpansion({"inputB", {"inputB_0", "inputB_1"}});
+ connectivity.RegisterInputArgExpansion({"inputA", DT_FLOAT, {"inputA"}});
+ connectivity.RegisterInputArgExpansion(
+ {"inputB", DT_FLOAT, {"inputB_0", "inputB_1"}});
NodeDef node;
node.add_input("inputA:0");
std::unordered_map<string, AttrValue> func_attr;
func_attr["T"].set_type(DT_FLOAT);
- FunctionLibraryDefinition library(OpRegistry::Global(), FunctionDefLibrary());
+ FunctionLibraryDefinition flib(OpRegistry::Global(), FunctionDefLibrary());
GrapplerFunctionItem item;
- TF_EXPECT_OK(MakeGrapplerFunctionItem(func, func_attr, library, &item));
+ TF_EXPECT_OK(MakeGrapplerFunctionItem(func, func_attr, flib, &item));
- EXPECT_EQ("XTimesTwo", item.function_name());
+ EXPECT_EQ("XTimesTwo", item.id);
EXPECT_EQ(4, item.function_body().node_size());
EXPECT_EQ(1, item.input_size());
std::unordered_map<string, AttrValue> func_attr;
func_attr["T"].set_type(DT_FLOAT);
- FunctionLibraryDefinition library(OpRegistry::Global(), FunctionDefLibrary());
+ FunctionLibraryDefinition flib(OpRegistry::Global(), FunctionDefLibrary());
GrapplerFunctionItem item;
- TF_EXPECT_OK(MakeGrapplerFunctionItem(func, func_attr, library, &item));
+ TF_EXPECT_OK(MakeGrapplerFunctionItem(func, func_attr, flib, &item));
- EXPECT_EQ("SubGrad", item.function_name());
+ EXPECT_EQ("SubGrad", item.id);
EXPECT_EQ(12, item.function_body().node_size());
ASSERT_EQ(3, item.input_size());
}
TEST_F(FunctionsTest, FromFunctionDefWithNestedFuncs) {
- FunctionLibraryDefinition library(OpRegistry::Global(), FunctionDefLibrary());
- TF_ASSERT_OK(library.AddFunctionDef(FunctionDefHelper::Define(
+ FunctionLibraryDefinition flib(OpRegistry::Global(), FunctionDefLibrary());
+ TF_ASSERT_OK(flib.AddFunctionDef(FunctionDefHelper::Define(
// Name
"Swap",
// Args
func_attr["T"].set_type(DT_FLOAT);
GrapplerFunctionItem item;
- TF_EXPECT_OK(MakeGrapplerFunctionItem(func, func_attr, library, &item));
+ TF_EXPECT_OK(MakeGrapplerFunctionItem(func, func_attr, flib, &item));
int count = 0;
for (const NodeDef &node : item.function_body().node()) {
{{"out", "Exp:y:0"}});
std::unordered_map<string, AttrValue> func_attr;
- FunctionLibraryDefinition library(OpRegistry::Global(), FunctionDefLibrary());
+ FunctionLibraryDefinition flib(OpRegistry::Global(), FunctionDefLibrary());
GrapplerFunctionItem item;
- TF_EXPECT_OK(MakeGrapplerFunctionItem(func, func_attr, library, &item));
+ TF_EXPECT_OK(MakeGrapplerFunctionItem(func, func_attr, flib, &item));
EXPECT_EQ(1, item.output_size());
EXPECT_EQ("Exp", item.output(0).output_tensors[0]);
{{"out0", "in0"}});
std::unordered_map<string, AttrValue> func_attr;
- FunctionLibraryDefinition library(OpRegistry::Global(), FunctionDefLibrary());
+ FunctionLibraryDefinition flib(OpRegistry::Global(), FunctionDefLibrary());
GrapplerFunctionItem item;
- TF_EXPECT_OK(MakeGrapplerFunctionItem(func, func_attr, library, &item));
+ TF_EXPECT_OK(MakeGrapplerFunctionItem(func, func_attr, flib, &item));
- EXPECT_EQ("ForwardInputs", item.function_name());
+ EXPECT_EQ("ForwardInputs", item.id);
EXPECT_EQ(5, item.function_body().node_size());
EXPECT_EQ(3, item.output_size());
std::unordered_map<string, AttrValue> func_attr;
func_attr["T"].set_type(DT_FLOAT);
- FunctionLibraryDefinition library(OpRegistry::Global(), FunctionDefLibrary());
+ FunctionLibraryDefinition flib(OpRegistry::Global(), FunctionDefLibrary());
GrapplerFunctionItem item;
- TF_EXPECT_OK(MakeGrapplerFunctionItem(func, func_attr, library, &item));
+ TF_EXPECT_OK(MakeGrapplerFunctionItem(func, func_attr, flib, &item));
EXPECT_EQ(0, item.input_size());
EXPECT_EQ(1, item.output_size());
EXPECT_EQ("two", cast.input(0));
}
+TEST_F(FunctionsTest, MakeSpecializedFunctionDef) {
+ const Tensor kTwo = test::AsScalar<int64>(2);
+ FunctionDef func = FunctionDefHelper::Define(
+ // Name
+ "XTimesTwo",
+ // Args
+ {"x: T"},
+ // Return values
+ {"y: T"},
+ // Attr def
+ {"T: {float, double, int32, int64}"},
+ // Nodes
+ {
+ {{"two"}, "Const", {}, {{"value", kTwo}, {"dtype", DT_INT64}}},
+ {{"scale"}, "Cast", {"two"}, {{"SrcT", DT_INT64}, {"DstT", "$T"}}},
+ {{"y"}, "Mul", {"x", "scale"}, {{"T", "$T"}}},
+ });
+
+ std::unordered_map<string, AttrValue> func_attr;
+ func_attr["T"].set_type(DT_FLOAT);
+ FunctionLibraryDefinition flib(OpRegistry::Global(), FunctionDefLibrary());
+
+ GrapplerFunctionItem item;
+ TF_EXPECT_OK(MakeGrapplerFunctionItem(func, func_attr, flib, &item));
+
+ FunctionDef specialized;
+ TF_EXPECT_OK(MakeSpecializedFunctionDef(item, flib, &specialized));
+
+ // Input and output types are resolved based on instantiation attributes.
+ EXPECT_EQ("x", specialized.signature().input_arg(0).name());
+ EXPECT_EQ(DT_FLOAT, specialized.signature().input_arg(0).type());
+ EXPECT_EQ("y", specialized.signature().output_arg(0).name());
+ EXPECT_EQ(DT_FLOAT, specialized.signature().output_arg(0).type());
+
+ // Function body specialized for instantiation types
+ int count = 0;
+ for (const NodeDef &node : specialized.node_def()) {
+ if (node.name() == "scale" && count++) {
+ EXPECT_EQ(DT_FLOAT, node.attr().at("DstT").type());
+ } else if (node.name() == "y" && count++) {
+ EXPECT_EQ("Mul", node.op());
+ EXPECT_EQ("x:0", node.input(0));
+ EXPECT_EQ("scale:y:0", node.input(1));
+ EXPECT_EQ(DT_FLOAT, node.attr().at("T").type());
+ }
+ }
+ EXPECT_EQ(2, count);
+}
+
+TEST_F(FunctionsTest, SwapFunctionBodyAndMakeSpecializedFunctionDef) {
+ using test::function::NDef;
+
+ FunctionDef mul_func = FunctionDefHelper::Create(
+ "MyMul", {"x:T", "y:T"}, {"z:T"}, {"T: {float, double}"},
+ {{{"output"}, "Mul", {"x", "y"}, {{"T", "$T"}}}},
+ /* Mapping between function returns and function node outputs. */
+ {{"z", "output:z:0"}});
+
+ FunctionDef func = FunctionDefHelper::Create(
+ "MySquare", {"x:T"}, {"z:T"}, {"T: {float, double}"},
+ {{{"output"}, "MyMul", {"x", "x"}, {{"T", "$T"}}}},
+ /* Mapping between function returns and function node outputs. */
+ {{"z", "output:z:0"}});
+
+ GraphDef id_func_body = test::function::GDef(
+ {/* pass input to output through identity */
+ NDef("output", "Identity", {"x"}, {{"T", "float"}})});
+
+ std::unordered_map<string, AttrValue> func_attr;
+ func_attr["T"].set_type(DT_FLOAT);
+
+ FunctionDefLibrary lib_def;
+ *lib_def.add_function() = func;
+ *lib_def.add_function() = mul_func;
+ FunctionLibraryDefinition flib(OpRegistry::Global(), lib_def);
+
+ GrapplerFunctionItem item;
+ TF_EXPECT_OK(MakeGrapplerFunctionItem(func, func_attr, flib, &item));
+
+ // Replace function body with identity function
+ item.SwapFunctionBody(std::move(id_func_body));
+ FunctionDef specialized;
+ TF_EXPECT_OK(MakeSpecializedFunctionDef(item, flib, &specialized));
+
+ // Check that graph body was updated.
+ int count = 0;
+ for (const NodeDef &node : specialized.node_def()) {
+ if (node.name() == "output" && count++) {
+ EXPECT_EQ("Identity", node.op());
+ EXPECT_EQ("x:0", node.input(0));
+ }
+ }
+ EXPECT_EQ(1, count);
+
+ // And return tensor mapping was updated with a new output name (z->output).
+ EXPECT_EQ("output:output:0", (*specialized.mutable_ret())["z"]);
+}
+
} // namespace
} // namespace grappler
} // namespace tensorflow
projects / platform / upstream / tensorflow.git / commitdiff