Function GraphToFunctionDef

tensorflow/core/framework/graph_to_functiondef.cc:113–292

// code in tensorflow/python/framework/function.py.

Status GraphToFunctionDef(const Graph& graph, const string& name,
                          const ControlRetMapping& control_ret,
                          FunctionDef* fdef) {
  fdef->mutable_signature()->set_name(name);

  std::unordered_map<string, string> tensor_renaming;
  std::unordered_map<string, string> return_values;
  NodeNameMapping node_names;

  for (Node const* node : graph.op_nodes()) {
    if (node->IsArg()) {
      int index;
      DataType type;
      TF_RETURN_IF_ERROR(GetNodeAttr(node->attrs(), "T", &type));
      TF_RETURN_IF_ERROR(GetNodeAttr(node->attrs(), "index", &index));
      while (fdef->signature().input_arg_size() <= index) {
        fdef->mutable_signature()->add_input_arg();
      }
      OpDef::ArgDef* argdef =
          fdef->mutable_signature()->mutable_input_arg(index);
      argdef->set_type(type);
      const string normalized = node_names.Normalize(node->name());
      argdef->set_name(normalized);
      tensor_renaming[strings::StrCat(node->name(), ":0")] = normalized;
      continue;
    }

    if (node->IsRetval()) {
      int index;
      DataType type;
      TF_RETURN_IF_ERROR(GetNodeAttr(node->attrs(), "T", &type));
      TF_RETURN_IF_ERROR(GetNodeAttr(node->attrs(), "index", &index));
      while (fdef->signature().output_arg_size() <= index) {
        fdef->mutable_signature()->add_output_arg();
      }
      OpDef::ArgDef* argdef =
          fdef->mutable_signature()->mutable_output_arg(index);
      argdef->set_type(type);
      const string normalized = node_names.Normalize(node->name());
      argdef->set_name(normalized);
      Edge const* edge;
      TF_RETURN_IF_ERROR(node->input_edge(0, &edge));
      return_values[normalized] =
          strings::StrCat(edge->src()->name(), ":", edge->src_output());
      continue;
    }

    NodeDef* node_def = fdef->add_node_def();
    *node_def = node->def();
    if (!node->assigned_device_name().empty()) {
      node_def->set_device(node->assigned_device_name());
    }
    node_def->set_name(node_names.Uniquify(node->name()));
    MergeDebugInfo(NodeDebugInfo(node->def()), node_def);

    // Check if a node must be a part of control return set.
    absl::optional<string> maybe_control_ret =
        control_ret ? control_ret(node) : absl::nullopt;