Update Argument/Fetch nodes in graphdef exporter to make sure node na…

…me is canonicalized (e.g. no :0 in suffix). PiperOrigin-RevId: 400069029 Change-Id: I40a686f715f9c61d38f0d71ed91f843987002c2e
bas-aarts · Oct 1, 2021 · 328d945 · 328d945
1 parent 137658a
commit 328d945
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 33 deletions.
diff --git a/tensorflow/compiler/mlir/tensorflow/tests/mlir2graphdef/fetch_feed_names.mlir b/tensorflow/compiler/mlir/tensorflow/tests/mlir2graphdef/fetch_feed_names.mlir
diff --git a/tensorflow/compiler/mlir/tensorflow/translate/export_graphdef.cc b/tensorflow/compiler/mlir/tensorflow/translate/export_graphdef.cc
@@ -175,7 +175,7 @@ StatusOr<std::unique_ptr<NodeDef>> Exporter::GetArgumentNode(
 
   auto node_def = absl::make_unique<NodeDef>();
   if (!name.empty())
-    node_def->set_name(std::string(ParseTensorName(name.str()).node()));
+    node_def->set_name(name.str());
   else
     node_def->set_name(
         std::string(op_to_name_.GetUniqueName(func.getName().str())));
@@ -235,7 +235,7 @@ StatusOr<std::unique_ptr<NodeDef>> Exporter::GetReturnNode(
     FuncOp function, Value operand, unsigned index, llvm::StringRef name) {
   auto node_def = absl::make_unique<NodeDef>();
   if (!name.empty())
-    node_def->set_name(std::string(ParseTensorName(name.str()).node()));
+    node_def->set_name(name.str());
   else
     node_def->set_name(
         std::string(op_to_name_.GetUniqueName(function.getName().str())));