chore: import upstream snapshot with attribution

tensorflow/cc/tools/BUILD

@@ -0,0 +1,55 @@
+#include "third_party/absl/strings/str_cat.h"
+#Description:
+# TensorFlow cc tools.
+
+load(
+    "//tensorflow:tensorflow.bzl",
+    "tf_cc_test",
+)
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = ["//visibility:public"],
+    licenses = ["notice"],
+)
+
+cc_library(
+    name = "freeze_saved_model",
+    srcs = ["freeze_saved_model.cc"],
+    hdrs = ["freeze_saved_model.h"],
+    deps = [
+        "//tensorflow/cc/saved_model:loader",
+        "//tensorflow/core:core_cpu",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:protos_all_cc",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/status",
+    ],
+)
+
+tf_cc_test(
+    name = "freeze_saved_model_test",
+    srcs = ["freeze_saved_model_test.cc"],
+    deps = [
+        ":freeze_saved_model",
+        "//tensorflow/cc:cc_ops",
+        "//tensorflow/cc:ops",
+        "//tensorflow/cc:resource_variable_ops",
+        "//tensorflow/cc:scope",
+        "//tensorflow/cc/saved_model:loader",
+        "//tensorflow/core:core_cpu",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:framework_internal",
+        "//tensorflow/core:portable_gif_internal",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core:tensorflow",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+        "//tensorflow/core:testlib",
+        "//tensorflow/core/platform:status",
+        "@com_google_absl//absl/status",
+        "@xla//xla/tsl/platform:errors",
+    ],
+)

tensorflow/cc/tools/freeze_saved_model.cc

@@ -0,0 +1,304 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/cc/tools/freeze_saved_model.h"
+
+#include <cstddef>
+#include <queue>
+#include <string>
+#include <unordered_map>
+#include <unordered_set>
+#include <vector>
+
+#include "absl/log/log.h"
+#include "absl/status/status.h"
+#include "tensorflow/cc/saved_model/loader.h"
+#include "xla/tsl/platform/errors.h"
+#include "tensorflow/core/framework/attr_value.pb.h"
+#include "tensorflow/core/framework/function.pb.h"
+#include "tensorflow/core/framework/graph.pb.h"
+#include "tensorflow/core/framework/node_def.pb.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/versions.pb.h"
+#include "tensorflow/core/lib/strings/str_util.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/platform/statusor.h"
+#include "tensorflow/core/platform/types.h"
+#include "tensorflow/core/protobuf/meta_graph.pb.h"
+#include "tensorflow/core/public/session.h"
+
+namespace tensorflow {
+
+namespace {
+
+// Gets tensor names from tensor_info and inserts them into the set of tensor
+// names.
+void GetTensorNamesFromTensorInfo(
+    const TensorInfo& tensor_info,
+    std::unordered_set<std::string>* tensor_names) {
+  if (tensor_info.has_coo_sparse()) {
+    // If the tensor is sparse we have to add all three tensors of the sparse
+    // representations.
+    const TensorInfo_CooSparse& coo_sparse = tensor_info.coo_sparse();
+    tensor_names->insert(coo_sparse.values_tensor_name());
+    tensor_names->insert(coo_sparse.indices_tensor_name());
+    tensor_names->insert(coo_sparse.dense_shape_tensor_name());
+  } else if (tensor_info.has_composite_tensor()) {
+    for (const auto& component : tensor_info.composite_tensor().components()) {
+      tensor_names->insert(component.name());
+    }
+  } else {
+    tensor_names->insert(tensor_info.name());
+  }
+}
+
+// Gets the union of all inputs and outputs of all SignatureDefs in the bundle
+void GetSignatureDefsInputsAndOutputs(
+    const SavedModelBundle& saved_model_bundle,
+    std::unordered_set<std::string>* inputs,
+    std::unordered_set<std::string>* outputs) {
+  for (auto& sigdef_elem : saved_model_bundle.meta_graph_def.signature_def()) {
+    const SignatureDef& signature_def = sigdef_elem.second;
+    for (auto& input_elem : signature_def.inputs()) {
+      GetTensorNamesFromTensorInfo(input_elem.second, inputs);
+    }
+    for (auto& output_elem : signature_def.outputs()) {
+      GetTensorNamesFromTensorInfo(output_elem.second, outputs);
+    }
+  }
+}
+
+// Gets a map from string node name to NodeDef.
+void GetNodeNameToNodeDefMap(
+    GraphDef* graph_def,
+    std::unordered_map<std::string, NodeDef*>* name_to_node_map) {
+  for (size_t i = 0; i < graph_def->node_size(); i++) {
+    NodeDef* node = graph_def->mutable_node(i);
+    (*name_to_node_map)[node->name()] = node;
+  }
+}
+
+// Strips off the tensor part of the tensor_name to get the node_name.
+const std::string GetNodeNameFromTensorName(std::string tensor_name) {
+  if (tensor_name[0] == '^') {
+    tensor_name.erase(0, 1);
+  }
+  std::vector<std::string> tensor_name_parts =
+      str_util::Split(tensor_name, ':');
+  return tensor_name_parts[0];
+}
+
+// Gets the set of node names needed by `outputs` and the corresponding set of
+// variable nodes to convert.
+void GetReachableNodesAndVariables(
+    GraphDef* graph_def, const std::unordered_set<std::string>& outputs,
+    const std::unordered_map<std::string, NodeDef*>& name_to_node_map,
+    std::unordered_set<std::string>* reachable_node_names,
+    std::unordered_set<std::string>* variable_node_names) {
+  // TODO(suharshs): Add support for ResourceVariables.
+  static const std::unordered_set<std::string>* kVariableTypes =
+      new std::unordered_set<std::string>(
+          {"Variable", "VariableV2", "VarHandleOp"});
+
+  std::queue<std::string> nodes_to_visit;
+  for (const std::string& output_tensor_name : outputs) {
+    nodes_to_visit.push(GetNodeNameFromTensorName(output_tensor_name));
+  }
+  // We do a traversal backwards from the outputs specified in the MetaGraphDef.
+  while (!nodes_to_visit.empty()) {
+    const std::string node_name = nodes_to_visit.front();
+    nodes_to_visit.pop();
+    if (reachable_node_names->find(node_name) != reachable_node_names->end()) {
+      continue;
+    }
+    reachable_node_names->insert(node_name);
+    NodeDef* node = name_to_node_map.at(node_name);
+    if (kVariableTypes->find(node->op()) != kVariableTypes->end()) {
+      variable_node_names->insert(node->name());
+    }
+    for (const std::string& input_tensor_name : node->input()) {
+      nodes_to_visit.push(GetNodeNameFromTensorName(input_tensor_name));
+    }
+  }
+}
+
+// Gets a map from variable name to variable value.
+absl::Status GetVariableNameToTensorMap(
+    Session* session,
+    const std::unordered_map<std::string, NodeDef*>& name_to_node_map,
+    std::unordered_set<std::string> variable_names_set,
+    std::unordered_map<std::string, Tensor>* variable_name_to_value_map) {
+  if (variable_names_set.empty()) {
+    return absl::OkStatus();
+  }
+  std::vector<std::string> variable_names;
+  variable_names.reserve(variable_names_set.size());
+  std::vector<std::string> tensor_names;
+  tensor_names.reserve(variable_names_set.size());
+  for (const std::string& node_name : variable_names_set) {
+    variable_names.push_back(node_name);
+    NodeDef* node_def = name_to_node_map.at(node_name);
+    if (node_def->op() == "VarHandleOp") {
+      // If this is a resource variable, we have to run the corresponding
+      // ReadVariableOp.
+      tensor_names.push_back(node_name + "/Read/ReadVariableOp:0");
+    } else {
+      tensor_names.push_back(node_name + ":0");
+    }
+  }
+  std::vector<Tensor> outputs;
+  TF_RETURN_IF_ERROR(
+      session->Run(/* inputs */ {}, tensor_names, /* targets */ {}, &outputs));
+  for (size_t i = 0; i < variable_names.size(); i++) {
+    (*variable_name_to_value_map)[variable_names[i]] = outputs[i];
+  }
+  return absl::OkStatus();
+}
+
+// Converts a Variable NodeDef into a Constant NodeDef.
+void ConvertVariableToConstant(const NodeDef& variable_node,
+                               const Tensor& variable_value,
+                               NodeDef* const_node) {
+  const_node->set_name(variable_node.name());
+  const_node->set_op("Const");
+  (*const_node->mutable_attr())["dtype"] = variable_node.attr().at("dtype");
+  variable_value.AsProtoTensorContent(
+      (*const_node->mutable_attr())["value"].mutable_tensor());
+}
+
+// Converts a ReadVariableOp NodeDef to an Identity NodeDef.
+void ConvertReadVariableOpToIdentity(const NodeDef& node,
+                                     NodeDef* identity_node) {
+  identity_node->set_name(node.name());
+  identity_node->set_op("Identity");
+  (*identity_node->mutable_attr())["T"] = node.attr().at("dtype");
+  identity_node->add_input(node.input(0));
+}
+
+// Returns the name of the VarHandleOp that provides input (possibly indirectly)
+// to node with node_name. A typical indirect chain of nodes (that can occur due
+// to graph inlining) is the following: VarHandleOp -> Identity -> Identity ->
+// ReadVariableOp. Calling the function on any of these nodes would return the
+// name of the VarHandleOp.
+absl::StatusOr<std::string> GetVarHandleName(
+    const std::unordered_map<std::string, NodeDef*>& name_to_node_map,
+    std::string node_name) {
+  const NodeDef* node = name_to_node_map.at(node_name);
+  while (node->input_size() > 0) {
+    auto parent = name_to_node_map.find(node->input(0));
+    if (parent == name_to_node_map.end()) break;
+    node = parent->second;
+    if (node->op() != "Identity") {
+      VLOG(2) << "Stopping at non-identity node " << node->op();
+      break;
+    }
+  }
+  if (node->op() == "VarHandleOp") {
+    return node->name();
+  }
+  return absl::NotFoundError("No VarHandleOp ancestor found");
+}
+
+// Looks up the variable handle that provides input to node with node_name,
+// and returns the handle name if the handle corresponds to a variable that we
+// want to freeze (i.e. its name is contained in variable_node_names). If there
+// is no such handle in the graph (or we do not want to save that variable)
+// then NotFound error is returned.
+absl::StatusOr<std::string> GetHandleNameIfNeedsToFreeze(
+    const std::unordered_map<std::string, NodeDef*>& name_to_node_map,
+    std::string node_name,
+    const std::unordered_set<std::string>& variable_node_names) {
+  absl::StatusOr<std::string> var_handle_name =
+      GetVarHandleName(name_to_node_map, node_name);
+  if (var_handle_name.ok() && variable_node_names.count(*var_handle_name)) {
+    return var_handle_name;
+  }
+  return absl::NotFoundError("No VarHandleOp ancestor found");
+}
+
+// Freezes the subgraph of all nodes needed by `outputs`.
+absl::Status FreezeGraphDef(const SavedModelBundle& saved_model_bundle,
+                            const std::unordered_set<std::string>& outputs,
+                            GraphDef* frozen_graph_def) {
+  GraphDef graph_def = saved_model_bundle.meta_graph_def.graph_def();
+  // Copy versions and library as-is from original graph.
+  *frozen_graph_def->mutable_versions() = graph_def.versions();
+  *frozen_graph_def->mutable_library() = graph_def.library();
+  // If the graph is empty there is nothing left to do.
+  if (graph_def.node_size() == 0) {
+    return absl::OkStatus();
+  }
+  // name_to_node_map is needed to get the inputs from the NodeDef corresponding
+  // the a string node name. These inputs are used when doing our backwards
+  // traversal.
+  std::unordered_map<std::string, NodeDef*> name_to_node_map;
+  GetNodeNameToNodeDefMap(&graph_def, &name_to_node_map);
+  std::unordered_set<std::string> reachable_node_names;
+  std::unordered_set<std::string> variable_node_names;
+  GetReachableNodesAndVariables(&graph_def, outputs, name_to_node_map,
+                                &reachable_node_names, &variable_node_names);
+  std::unordered_map<std::string, Tensor> variable_to_value_map;
+  TF_RETURN_IF_ERROR(GetVariableNameToTensorMap(
+      saved_model_bundle.session.get(), name_to_node_map, variable_node_names,
+      &variable_to_value_map));
+  // We copy the nodes in the same order they were in the original graph_def.
+  for (const NodeDef& node : graph_def.node()) {
+    if (reachable_node_names.find(node.name()) == reachable_node_names.end()) {
+      continue;
+    }
+    if (variable_node_names.find(node.name()) != variable_node_names.end()) {
+      ConvertVariableToConstant(node, variable_to_value_map[node.name()],
+                                frozen_graph_def->add_node());
+      continue;
+    } else if (node.op() == "ReadVariableOp" &&
+               GetHandleNameIfNeedsToFreeze(name_to_node_map, node.name(),
+                                            variable_node_names)
+                   .ok()) {
+      // If the node is a ReadVariableOp, its input VarHandleOp will be
+      // converted to a Constant, so we will need to convert it to an Identity.
+      ConvertReadVariableOpToIdentity(node, frozen_graph_def->add_node());
+      continue;
+    } else if (node.op() == "Identity") {
+      absl::StatusOr<std::string> handle_name = GetHandleNameIfNeedsToFreeze(
+          name_to_node_map, node.name(), variable_node_names);
+      if (handle_name.ok()) {
+        // Identity node that is forwarding the value of a frozen
+        // VarhandleOp. We ensure that the dtype matches of the variable dtype.
+        NodeDef* new_node = frozen_graph_def->add_node();
+        *new_node = node;
+        (*new_node->mutable_attr())["T"] =
+            name_to_node_map.at(*handle_name)->attr().at("dtype");
+        continue;
+      }
+    }
+    // If the node isn't a variable, just copy the node as-is.
+    *frozen_graph_def->add_node() = node;
+  }
+  return absl::OkStatus();
+}
+
+}  // namespace
+
+absl::Status FreezeSavedModel(const SavedModelBundle& saved_model_bundle,
+                              GraphDef* frozen_graph_def,
+                              std::unordered_set<std::string>* inputs,
+                              std::unordered_set<std::string>* outputs) {
+  GetSignatureDefsInputsAndOutputs(saved_model_bundle, inputs, outputs);
+  TF_RETURN_IF_ERROR(
+      FreezeGraphDef(saved_model_bundle, *outputs, frozen_graph_def));
+  return absl::OkStatus();
+}
+
+}  // namespace tensorflow

tensorflow/cc/tools/freeze_saved_model.h

@@ -0,0 +1,45 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_CC_TOOLS_FREEZE_SAVED_MODEL_H_
+#define TENSORFLOW_CC_TOOLS_FREEZE_SAVED_MODEL_H_
+
+#include <unordered_set>
+
+#include "absl/status/status.h"
+#include "tensorflow/cc/saved_model/loader.h"
+#include "tensorflow/core/framework/graph.pb.h"
+#include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/platform/types.h"
+
+namespace tensorflow {
+
+// Returns a frozen GraphDef, input tensors, and output tensors from the loaded
+// SavedModelBundle.
+// `inputs` and `outputs` consist of the union of all inputs and outputs in the
+// SignatureDefs in the SavedModelBundle.
+// FreezeSavedModel sets `frozen_graph_def` to a GraphDef of all nodes needed by
+// `outputs`. All variables in the supplied SavedModelBundle are converted to
+// constants, set to the value of the variables, by running the restored Session
+// in the SavedModelBundle.
+// WARNING: Only the variable checkpoints will be reflected in the frozen
+// graph_def. All saved_model assets will be ignored.
+absl::Status FreezeSavedModel(const SavedModelBundle& saved_model_bundle,
+                              GraphDef* frozen_graph_def,
+                              std::unordered_set<std::string>* inputs,
+                              std::unordered_set<std::string>* outputs);
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CC_TOOLS_FREEZE_SAVED_MODEL_H_

tensorflow/cc/tools/freeze_saved_model_test.cc

@@ -0,0 +1,521 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/cc/tools/freeze_saved_model.h"
+
+#include <cstddef>
+#include <memory>
+#include <string>
+#include <unordered_set>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "tensorflow/cc/framework/ops.h"
+#include "tensorflow/cc/framework/scope.h"
+#include "tensorflow/cc/ops/array_ops.h"
+#include "tensorflow/cc/ops/const_op.h"
+#include "tensorflow/cc/ops/math_ops.h"
+#include "tensorflow/cc/ops/resource_variable_ops.h"
+#include "tensorflow/cc/ops/state_ops.h"
+#include "tensorflow/cc/saved_model/loader.h"
+#include "xla/tsl/lib/core/status_test_util.h"
+#include "xla/tsl/platform/errors.h"
+#include "tensorflow/core/framework/function_testlib.h"
+#include "tensorflow/core/framework/graph.pb.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_testutil.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/framework/versions.pb.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/platform/types.h"
+#include "tensorflow/core/protobuf/meta_graph.pb.h"
+#include "tensorflow/core/public/session.h"
+#include "tensorflow/core/public/session_options.h"
+
+namespace tensorflow {
+namespace {
+
+class FreezeTest : public ::testing::Test {
+ protected:
+  void GraphDefEqual(const GraphDef& actual, const GraphDef& expected) {
+    EXPECT_EQ(actual.ShortDebugString(), expected.ShortDebugString());
+  }
+
+  // Builds a SignatureDef with the provided `inputs` and `outputs`.
+  SignatureDef BuildSignatureDef(
+      const std::unordered_set<std::string>& inputs,
+      const std::unordered_set<std::string>& outputs) {
+    SignatureDef signature_def;
+    for (const std::string& input : inputs) {
+      (*signature_def.mutable_inputs())[input].set_name(input);
+    }
+    for (const std::string& output : outputs) {
+      (*signature_def.mutable_outputs())[output].set_name(output);
+    }
+    return signature_def;
+  }
+
+  // Adds `signature_def` to `saved_model_bundle` under `key`.
+  void AddSignatureDefToSavedModelBundle(const SignatureDef& signature_def,
+                                         const std::string& key,
+                                         SavedModelBundle* saved_model_bundle) {
+    MetaGraphDef* meta_graph_def = &saved_model_bundle->meta_graph_def;
+    (*meta_graph_def->mutable_signature_def())[key] = signature_def;
+  }
+
+  // Adds an initialized session to `saved_model_bundle` using `graph_def` and
+  // initializing with `init_node`.
+  absl::Status InitializeSavedModelBundleSession(
+      const GraphDef& graph_def, const std::string& init_node,
+      SavedModelBundle* saved_model_bundle) {
+    SessionOptions session_options;
+    saved_model_bundle->session.reset(NewSession(session_options));
+    TF_RETURN_IF_ERROR(saved_model_bundle->session->Create(graph_def));
+    if (!init_node.empty()) {
+      std::vector<Tensor> outputs;
+      return saved_model_bundle->session->Run(
+          /* inputs */ {}, /* output_tensors */ {}, {init_node}, &outputs);
+    }
+    return absl::OkStatus();
+  }
+
+  // Adds `graph_def` to `saved_model_bundle` and initializes a session with
+  // `init_node`.
+  absl::Status AddGraphDefToSavedModelBundle(
+      const GraphDef& graph_def, const std::string& init_node,
+      SavedModelBundle* saved_model_bundle) {
+    MetaGraphDef* meta_graph_def = &saved_model_bundle->meta_graph_def;
+    *meta_graph_def->mutable_graph_def() = graph_def;
+    return InitializeSavedModelBundleSession(graph_def, init_node,
+                                             saved_model_bundle);
+  }
+
+  // Adds `graph_def` and `outputs` as the GraphDef and SignatureDef in
+  // `saved_model_bundle` and initializes a session with `init_node`.
+  absl::Status AddGraphDefWithOutputsToSavedModelBundle(
+      const GraphDef& graph_def, const std::unordered_set<std::string>& outputs,
+      const std::string& init_node, SavedModelBundle* saved_model_bundle) {
+    SignatureDef signature_def =
+        BuildSignatureDef(std::unordered_set<std::string>(), outputs);
+    AddSignatureDefToSavedModelBundle(signature_def, "signature_def",
+                                      saved_model_bundle);
+    return AddGraphDefToSavedModelBundle(graph_def, init_node,
+                                         saved_model_bundle);
+  }
+
+  // Runs and compares the outputs of `tensor_name` on both the
+  // `unfrozen_session` and the `frozen_graph_def.
+  void RunAndCompareFrozenAndUnfrozenGraphs(Session* unfrozen_session,
+                                            const GraphDef& frozen_graph_def,
+                                            const std::string& tensor_name) {
+    std::vector<Tensor> unfrozen_outputs;
+    TF_ASSERT_OK(unfrozen_session->Run(/* inputs */ {}, {tensor_name},
+                                       /* targets */ {}, &unfrozen_outputs));
+
+    SessionOptions session_options;
+    std::unique_ptr<Session> frozen_session(NewSession(session_options));
+    TF_ASSERT_OK(frozen_session->Create(frozen_graph_def));
+    std::vector<Tensor> frozen_outputs;
+    TF_ASSERT_OK(frozen_session->Run(/* inputs */ {}, {tensor_name},
+                                     /* targets */ {}, &frozen_outputs));
+
+    test::ExpectTensorEqual<float>(unfrozen_outputs[0], frozen_outputs[0]);
+  }
+
+  void TestFreezeGraphWithoutDependentVariables(bool use_resource) {
+    // Test freezing a graph with variables that are not needed by the outputs
+    // in the SignatureDef. The resulting graph shouldn't be frozen, but
+    // non-dependent nodes should be pruned.
+    SavedModelBundle saved_model_bundle;
+    GraphDef graph_def;
+    Scope scope = Scope::NewRootScope();
+    Output a = ops::Const(scope.WithOpName("a"), 10.0f, {});
+    Output b = ops::Const(scope.WithOpName("b"), 10.0f, {});
+    Output c = ops::Mul(scope.WithOpName("c"), a, b);
+    if (use_resource) {
+      Output var =
+          ops::VarHandleOp(scope.WithOpName("var"), DataType::DT_FLOAT, {});
+      Output read_var = ops::ReadVariableOp(
+          scope.WithOpName("var/Read/ReadVariableOp"), var, DataType::DT_FLOAT);
+      auto assign = ops::AssignVariableOp(scope.WithOpName("assign"), var, a);
+    } else {
+      Output var =
+          ops::Variable(scope.WithOpName("var"), {}, DataType::DT_FLOAT);
+      Output assign = ops::Assign(scope.WithOpName("assign"), var, a);
+    }
+
+    TF_ASSERT_OK(scope.ToGraphDef(&graph_def));
+    // "c" isn't dependent on the variable, so nothing should be frozen.
+    TF_ASSERT_OK(AddGraphDefWithOutputsToSavedModelBundle(
+        graph_def, {"c:0"}, "assign", &saved_model_bundle));
+
+    GraphDef frozen_graph_def;
+    std::unordered_set<std::string> inputs;
+    std::unordered_set<std::string> outputs;
+    TF_ASSERT_OK(FreezeSavedModel(saved_model_bundle, &frozen_graph_def,
+                                  &inputs, &outputs));
+
+    GraphDef expected_graph_def;
+    Scope expected_scope = Scope::NewRootScope();
+    Output expected_a = ops::Const(expected_scope.WithOpName("a"), 10.0f, {});
+    Output expected_b = ops::Const(expected_scope.WithOpName("b"), 10.0f, {});
+    Output expected_c =
+        ops::Mul(expected_scope.WithOpName("c"), expected_a, expected_b);
+    TF_ASSERT_OK(expected_scope.ToGraphDef(&expected_graph_def));
+
+    GraphDefEqual(frozen_graph_def, expected_graph_def);
+
+    RunAndCompareFrozenAndUnfrozenGraphs(saved_model_bundle.session.get(),
+                                         frozen_graph_def, "c:0");
+  }
+
+  void TestFreezeGraphWithDependentVariables(bool use_resource,
+                                             bool use_identity = false) {
+    // Test freezing a graph with variables that are needed by outputs in the
+    // SignatureDef. The variables should be frozen.
+    SavedModelBundle saved_model_bundle;
+    GraphDef graph_def;
+    Scope scope = Scope::NewRootScope();
+    Output a = ops::Const(scope.WithOpName("a"), 10.0f, {});
+    Output read_var;
+    if (use_resource) {
+      Output var =
+          ops::VarHandleOp(scope.WithOpName("var"), DataType::DT_FLOAT, {});
+      if (use_identity) {
+        Output identity = ops::Identity(scope.WithOpName("identity"), var);
+        read_var =
+            ops::ReadVariableOp(scope.WithOpName("var/Read/ReadVariableOp"),
+                                identity, DataType::DT_FLOAT);
+      } else {
+        read_var =
+            ops::ReadVariableOp(scope.WithOpName("var/Read/ReadVariableOp"),
+                                var, DataType::DT_FLOAT);
+      }
+      auto assign = ops::AssignVariableOp(scope.WithOpName("assign"), var, a);
+    } else {
+      Output read_var =
+          ops::Variable(scope.WithOpName("var"), {}, DataType::DT_FLOAT);
+      Output assign = ops::Assign(scope.WithOpName("assign"), read_var, a);
+    }
+    Output c = ops::Mul(scope.WithOpName("c"), a, read_var);
+    TF_ASSERT_OK(scope.ToGraphDef(&graph_def));
+    // "c" isn't dependent on the variable, so nothing should be frozen.
+    TF_ASSERT_OK(AddGraphDefWithOutputsToSavedModelBundle(
+        graph_def, {"c:0"}, "assign", &saved_model_bundle));
+
+    GraphDef frozen_graph_def;
+    std::unordered_set<std::string> inputs;
+    std::unordered_set<std::string> outputs;
+    TF_ASSERT_OK(FreezeSavedModel(saved_model_bundle, &frozen_graph_def,
+                                  &inputs, &outputs));
+
+    // If using normal variables there should be 3 nodes in the resulting
+    // graph_def. If using resource variables there should be 4 nodes in the
+    // resulting graph_def if use_identity == false, otherwise 5 variables.
+    // In both cases, none should be variables.
+    size_t expected_nodes = use_resource ? (use_identity ? 5 : 4) : 3;
+
+    EXPECT_EQ(frozen_graph_def.node_size(), expected_nodes);
+    for (const NodeDef& node : frozen_graph_def.node()) {
+      EXPECT_NE(node.op(), "Variable") << node.name();
+      EXPECT_NE(node.op(), "VariableV2") << node.name();
+      EXPECT_NE(node.op(), "VarHandleOp") << node.name();
+      EXPECT_NE(node.op(), "ReadVariableOp") << node.name();
+    }
+
+    RunAndCompareFrozenAndUnfrozenGraphs(saved_model_bundle.session.get(),
+                                         frozen_graph_def, "c:0");
+  }
+
+  void TestFreezeGraphWithAndWithoutDependentVariables(bool use_resource) {
+    // Test freezing a graph with some variables that are needed and not needed
+    // by
+    // the outputs in the SignatureDef. The resulting graph should only freeze
+    // dependent variables.
+    SavedModelBundle saved_model_bundle;
+    GraphDef graph_def;
+    Scope scope = Scope::NewRootScope();
+    Output a = ops::Const(scope.WithOpName("a"), 10.0f, {});
+    Output read_var;
+
+    if (use_resource) {
+      Output var =
+          ops::VarHandleOp(scope.WithOpName("var"), DataType::DT_FLOAT, {});
+      read_var = ops::ReadVariableOp(
+          scope.WithOpName("var/Read/ReadVariableOp"), var, DataType::DT_FLOAT);
+      auto assign = ops::AssignVariableOp(scope.WithOpName("assign"), var, a);
+      Output var_1 =
+          ops::VarHandleOp(scope.WithOpName("var_1"), DataType::DT_FLOAT, {});
+      Output read_var_1 =
+          ops::ReadVariableOp(scope.WithOpName("var_1/Read/ReadVariableOp"),
+                              var, DataType::DT_FLOAT);
+      auto assign_1 =
+          ops::AssignVariableOp(scope.WithOpName("assign_1"), var_1, a);
+    } else {
+      read_var = ops::Variable(scope.WithOpName("var"), {}, DataType::DT_FLOAT);
+      Output assign = ops::Assign(scope.WithOpName("assign"), read_var, a);
+      Output var_1 =
+          ops::Variable(scope.WithOpName("var_1"), {}, DataType::DT_FLOAT);
+      Output assign_1 = ops::Assign(scope.WithOpName("assign_1"), var_1, a);
+    }
+
+    Output c = ops::Mul(scope.WithOpName("c"), a, read_var);
+    TF_ASSERT_OK(scope.ToGraphDef(&graph_def));
+    // "c" isn't dependent on the variable, so nothing should be frozen.
+    TF_ASSERT_OK(AddGraphDefWithOutputsToSavedModelBundle(
+        graph_def, {"c:0"}, "assign", &saved_model_bundle));
+
+    GraphDef frozen_graph_def;
+    std::unordered_set<std::string> inputs;
+    std::unordered_set<std::string> outputs;
+    TF_ASSERT_OK(FreezeSavedModel(saved_model_bundle, &frozen_graph_def,
+                                  &inputs, &outputs));
+
+    // There should be 3 nodes in the resulting graph_def, and none should be
+    // variables.
+    size_t expected_nodes = use_resource ? 4 : 3;
+    EXPECT_EQ(frozen_graph_def.node_size(), expected_nodes);
+    for (const NodeDef& node : frozen_graph_def.node()) {
+      EXPECT_NE(node.op(), "Variable") << node.name();
+      EXPECT_NE(node.op(), "VariableV2") << node.name();
+      EXPECT_NE(node.op(), "VarHandleOp") << node.name();
+      EXPECT_NE(node.op(), "ReadVariableOp") << node.name();
+    }
+
+    RunAndCompareFrozenAndUnfrozenGraphs(saved_model_bundle.session.get(),
+                                         frozen_graph_def, "c:0");
+  }
+};
+
+TEST_F(FreezeTest, InputsAndOutputsSingleSignatureDef) {
+  // Test that inputs and outputs get correctly populated for a single
+  // SignatureDef.
+  SavedModelBundle saved_model_bundle;
+  std::unordered_set<std::string> expected_inputs = {"input0:0", "input1:0"};
+  std::unordered_set<std::string> expected_outputs = {"output0:0", "output1:0"};
+  SignatureDef signature_def =
+      BuildSignatureDef(expected_inputs, expected_outputs);
+  AddSignatureDefToSavedModelBundle(signature_def, "signature_def",
+                                    &saved_model_bundle);
+  GraphDef frozen_graph_def;
+  std::unordered_set<std::string> inputs;
+  std::unordered_set<std::string> outputs;
+  TF_ASSERT_OK(FreezeSavedModel(saved_model_bundle, &frozen_graph_def, &inputs,
+                                &outputs));
+  EXPECT_EQ(expected_inputs, inputs);
+  EXPECT_EQ(expected_outputs, outputs);
+}
+
+TEST_F(FreezeTest, InputsAndOutputsMultipleSignatureDefs) {
+  // Test that inputs and outputs get correctly merged and populated when
+  // multiple SignatureDefs are provided.
+  SavedModelBundle saved_model_bundle;
+  SignatureDef signature_def_0 = BuildSignatureDef({"input0:0"}, {"output0:0"});
+  SignatureDef signature_def_1 = BuildSignatureDef({"input1:0"}, {"output1:0"});
+  AddSignatureDefToSavedModelBundle(signature_def_0, "signature_def_0",
+                                    &saved_model_bundle);
+  AddSignatureDefToSavedModelBundle(signature_def_1, "signature_def_1",
+                                    &saved_model_bundle);
+  GraphDef frozen_graph_def;
+  std::unordered_set<std::string> inputs;
+  std::unordered_set<std::string> outputs;
+  TF_ASSERT_OK(FreezeSavedModel(saved_model_bundle, &frozen_graph_def, &inputs,
+                                &outputs));
+  std::unordered_set<std::string> expected_inputs = {"input0:0", "input1:0"};
+  std::unordered_set<std::string> expected_outputs = {"output0:0", "output1:0"};
+  EXPECT_EQ(expected_inputs, inputs);
+  EXPECT_EQ(expected_outputs, outputs);
+}
+
+TEST_F(FreezeTest, GraphDefVersionsAndLibrary) {
+  // Test that GraphDef versions and library are copied correctly into the
+  // frozen graph.
+  SavedModelBundle saved_model_bundle;
+  GraphDef graph_def;
+  graph_def.mutable_versions()->set_producer(1234);
+  graph_def.mutable_versions()->set_min_consumer(1234);
+  *graph_def.mutable_library()->add_function() = test::function::NonZero();
+  TF_ASSERT_OK(
+      AddGraphDefToSavedModelBundle(graph_def, "", &saved_model_bundle));
+
+  GraphDef frozen_graph_def;
+  std::unordered_set<std::string> inputs;
+  std::unordered_set<std::string> outputs;
+  TF_ASSERT_OK(FreezeSavedModel(saved_model_bundle, &frozen_graph_def, &inputs,
+                                &outputs));
+
+  GraphDefEqual(frozen_graph_def, graph_def);
+}
+
+TEST_F(FreezeTest, GraphDefWithNoVariables) {
+  // Test freezing a graph with no variables.
+  SavedModelBundle saved_model_bundle;
+  GraphDef graph_def;
+  Scope scope = Scope::NewRootScope();
+  Output a = ops::Const(scope.WithOpName("a"), 10.0f, {});
+  Output b = ops::Const(scope.WithOpName("b"), 10.0f, {});
+  Output c = ops::Mul(scope.WithOpName("c"), a, b);
+  TF_ASSERT_OK(scope.ToGraphDef(&graph_def));
+  TF_ASSERT_OK(AddGraphDefWithOutputsToSavedModelBundle(graph_def, {"c:0"}, "",
+                                                        &saved_model_bundle));
+
+  GraphDef frozen_graph_def;
+  std::unordered_set<std::string> inputs;
+  std::unordered_set<std::string> outputs;
+  TF_ASSERT_OK(FreezeSavedModel(saved_model_bundle, &frozen_graph_def, &inputs,
+                                &outputs));
+
+  GraphDefEqual(frozen_graph_def, graph_def);
+}
+
+TEST_F(FreezeTest, GraphDefWithMultiOutputOperation) {
+  // Tensors from operations with multiple outputs get tensor suffixes when used
+  // in input fields of following nodes, i.e. split:0, split:1.
+  // Test that we traverse those correctly.
+  SavedModelBundle saved_model_bundle;
+  GraphDef graph_def;
+  Scope scope = Scope::NewRootScope();
+  Output a = ops::Const(scope.WithOpName("a"), {10.0f, 10.0f}, {2});
+  Output axis = ops::Const(scope.WithOpName("axis"), 0, {});
+  OutputList split = ops::Split(scope.WithOpName("split"), axis, a, 2).output;
+  Output b = ops::Const(scope.WithOpName("b"), 10.0f, {});
+  Output c = ops::Mul(scope.WithOpName("c"), split[1], b);
+  TF_ASSERT_OK(scope.ToGraphDef(&graph_def));
+  TF_ASSERT_OK(AddGraphDefWithOutputsToSavedModelBundle(graph_def, {"c:0"}, "",
+                                                        &saved_model_bundle));
+
+  GraphDef frozen_graph_def;
+  std::unordered_set<std::string> inputs;
+  std::unordered_set<std::string> outputs;
+  TF_ASSERT_OK(FreezeSavedModel(saved_model_bundle, &frozen_graph_def, &inputs,
+                                &outputs));
+
+  GraphDefEqual(frozen_graph_def, graph_def);
+}
+
+TEST_F(FreezeTest, GraphDefWithControlDependency) {
+  // Inputs that are control dependencies get tensor prefixes,
+  // i.e. ^control_dependency.
+  // Test that we traverse those correctly.
+  SavedModelBundle saved_model_bundle;
+  GraphDef graph_def;
+  Scope scope = Scope::NewRootScope();
+  Output source = ops::Const(scope.WithOpName("source"), 10.0f, {});
+  Output a = ops::Const(scope.WithOpName("a").WithControlDependencies(source),
+                        {10.0f, 10.0f}, {2});
+  Output b = ops::Const(scope.WithOpName("b"), 10.0f, {});
+  Output c = ops::Mul(scope.WithOpName("c"), a, b);
+  TF_ASSERT_OK(scope.ToGraphDef(&graph_def));
+  TF_ASSERT_OK(AddGraphDefWithOutputsToSavedModelBundle(graph_def, {"c:0"}, "",
+                                                        &saved_model_bundle));
+
+  GraphDef frozen_graph_def;
+  std::unordered_set<std::string> inputs;
+  std::unordered_set<std::string> outputs;
+  TF_ASSERT_OK(FreezeSavedModel(saved_model_bundle, &frozen_graph_def, &inputs,
+                                &outputs));
+
+  GraphDefEqual(frozen_graph_def, graph_def);
+}
+
+TEST_F(FreezeTest, GraphDefWithoutDependentVariables) {
+  TestFreezeGraphWithoutDependentVariables(false);
+}
+
+TEST_F(FreezeTest, GraphDefWithoutDependentResourceVariables) {
+  TestFreezeGraphWithoutDependentVariables(true);
+}
+
+TEST_F(FreezeTest, GraphDefWithDependentVariables) {
+  TestFreezeGraphWithDependentVariables(false);
+}
+
+TEST_F(FreezeTest, GraphDefWithDependentResourceVariables) {
+  TestFreezeGraphWithDependentVariables(true);
+}
+
+TEST_F(FreezeTest, GraphDefWithDependentResourceVariablesAndIdentity) {
+  TestFreezeGraphWithDependentVariables(true, true);
+}
+
+TEST_F(FreezeTest, GraphDefWithAndWithoutDependentVariables) {
+  TestFreezeGraphWithAndWithoutDependentVariables(false);
+}
+
+TEST_F(FreezeTest, GraphDefWithAndWithoutDependentResourceVariables) {
+  TestFreezeGraphWithAndWithoutDependentVariables(true);
+}
+
+TEST_F(FreezeTest, InputsAndOutputsCompositeTensorSignatureDef) {
+  // Test that inputs and outputs get correctly populated for a
+  // SignatureDef containing composite tensor inputs and outputs.
+  SavedModelBundle saved_model_bundle;
+  SignatureDef signature_def;
+
+  TensorInfo& in = (*signature_def.mutable_inputs())["input_arg"];
+  in.mutable_composite_tensor()->add_components()->set_name("input1:0");
+  in.mutable_composite_tensor()->add_components()->set_name("input2:0");
+
+  TensorInfo& out = (*signature_def.mutable_outputs())["output_arg"];
+  out.mutable_composite_tensor()->add_components()->set_name("output2:0");
+  out.mutable_composite_tensor()->add_components()->set_name("output1:0");
+
+  AddSignatureDefToSavedModelBundle(signature_def, "signature_def",
+                                    &saved_model_bundle);
+  GraphDef frozen_graph_def;
+  std::unordered_set<std::string> inputs;
+  std::unordered_set<std::string> outputs;
+  TF_ASSERT_OK(FreezeSavedModel(saved_model_bundle, &frozen_graph_def, &inputs,
+                                &outputs));
+  std::unordered_set<std::string> expected_inputs = {"input1:0", "input2:0"};
+  std::unordered_set<std::string> expected_outputs = {"output1:0", "output2:0"};
+  EXPECT_EQ(expected_inputs, inputs);
+  EXPECT_EQ(expected_outputs, outputs);
+}
+
+TEST_F(FreezeTest, InputsAndOutputsSparseCooSignatureDef) {
+  // Test that inputs and outputs get correctly populated for a
+  // SignatureDef containing composite tensor inputs and outputs.
+  SavedModelBundle saved_model_bundle;
+  SignatureDef signature_def;
+
+  TensorInfo& in = (*signature_def.mutable_inputs())["input_arg"];
+  in.mutable_coo_sparse()->set_values_tensor_name("input1:0");
+  in.mutable_coo_sparse()->set_indices_tensor_name("input2:0");
+  in.mutable_coo_sparse()->set_dense_shape_tensor_name("input3:0");
+
+  TensorInfo& out = (*signature_def.mutable_outputs())["output_arg"];
+  out.mutable_coo_sparse()->set_values_tensor_name("output1:0");
+  out.mutable_coo_sparse()->set_indices_tensor_name("output2:0");
+  out.mutable_coo_sparse()->set_dense_shape_tensor_name("output3:0");
+
+  AddSignatureDefToSavedModelBundle(signature_def, "signature_def",
+                                    &saved_model_bundle);
+  GraphDef frozen_graph_def;
+  std::unordered_set<std::string> inputs;
+  std::unordered_set<std::string> outputs;
+  TF_ASSERT_OK(FreezeSavedModel(saved_model_bundle, &frozen_graph_def, &inputs,
+                                &outputs));
+  std::unordered_set<std::string> expected_inputs = {"input1:0", "input2:0",
+                                                     "input3:0"};
+  std::unordered_set<std::string> expected_outputs = {"output1:0", "output2:0",
+                                                      "output3:0"};
+  EXPECT_EQ(expected_inputs, inputs);
+  EXPECT_EQ(expected_outputs, outputs);
+}
+
+}  // namespace
+}  // namespace tensorflow