Minor improvements over cl/163390862: avoid storing unused data, and add a

little bit of documentation for future reference.

PiperOrigin-RevId: 163483408

tensorflow/core/grappler/optimizers/graph_rewriter.cc

@@ -37,10 +37,9 @@ GraphRewriter::GraphRewriter(const GrapplerItem& item) {
     Status s = op_registry->LookUp(node.op(), &op_reg_data);
     // TODO(bsteiner): make this not a best-effort lookup and evaluation?
     if (s.ok()) {
-      s = InOutTypesForNode(node, op_reg_data->op_def, &info->inputs,
-                            &info->outputs);
+      DataTypeVector inputs;
+      s = InOutTypesForNode(node, op_reg_data->op_def, &inputs, &info->outputs);
       if (!s.ok()) {
-        info->inputs.clear();
         info->outputs.clear();
       }
     }

tensorflow/core/grappler/optimizers/graph_rewriter.h

@@ -71,7 +71,6 @@ class GraphRewriter {
 
     // These are filled in when the NodeInfo is built, but not that they
     // may be empty - if the op could not be loaded from the registry.
-    DataTypeVector inputs;
     DataTypeVector outputs;
   };
 

tensorflow/core/grappler/optimizers/model_pruner.cc

@@ -85,7 +85,9 @@ Status ModelPruner::Optimize(Cluster* cluster, const GrapplerItem& item,
     // - Don't prune nodes that are driven by another device since these could
     //   be used to reduce cross device communication.
     // - Don't remove nodes that receive reference values, as those can be
-    //   converting references to non-references.
+    //   converting references to non-references. It is important to preserve
+    //   these non-references since the partitioner will avoid sending
+    //   non-references accross partitions more than once.
     if (!rewriter.DrivesControlDependency(node) &&
         !rewriter.IsDrivenByControlDependency(node) &&
         !rewriter.IsConnectedToFunction(node) &&