Refactor fold_constants to make reuse easier.

Change: 138588267

1 files changed, 6 insertions, 0 deletions

diff --git a/tensorflow/tools/graph_transforms/transform_utils.h b/tensorflow/tools/graph_transforms/transform_utils.h
index 8f6433fa34..7fb885b1ac 100644
--- a/tensorflow/tools/graph_transforms/transform_utils.h
+++ b/tensorflow/tools/graph_transforms/transform_utils.h
@@ -45,6 +45,12 @@ void FilterGraphDef(const GraphDef& input_graph_def,
                     std::function<bool(const NodeDef&)> selector,
                     GraphDef* output_graph_def);
 
+// Creates a copy of the input graph, with all occurences of the attributes with
+// the names in the argument removed from the node defs.
+void RemoveAttributes(const GraphDef& input_graph_def,
+                      const std::vector<string>& attributes,
+                      GraphDef* output_graph_def);
+
 }  // namespace graph_transforms
 }  // namespace tensorflow