Add a helper method Node::input_edges() that populates a vector of all input edges to a node, indexed by input number.

Change: 145301512

1 files changed, 25 insertions, 0 deletions

diff --git a/tensorflow/core/graph/graph.cc b/tensorflow/core/graph/graph.cc
index 8776e1a281..509c67c11f 100644
--- a/tensorflow/core/graph/graph.cc
+++ b/tensorflow/core/graph/graph.cc
@@ -173,6 +173,31 @@ Status Node::input_edge(int idx, const Edge** e) const {
   return errors::NotFound("Could not find input edge ", idx, " for ", name());
 }
 
+// Returns a vector of the non-control input edges to a node, indexed by ID.
+Status Node::input_edges(std::vector<const Edge*>* input_edges) const {
+  input_edges->clear();
+  input_edges->resize(num_inputs(), nullptr);
+
+  for (const Edge* edge : in_edges()) {
+    if (edge->IsControlEdge()) continue;
+    if (edge->dst_input() < 0 || edge->dst_input() >= num_inputs()) {
+      return errors::Internal("Invalid edge input number ", edge->dst_input());
+    }
+    if ((*input_edges)[edge->dst_input()] != nullptr) {
+      return errors::Internal("Duplicate edge input number: ",
+                              edge->dst_input());
+    }
+    (*input_edges)[edge->dst_input()] = edge;
+  }
+
+  for (int i = 0; i < num_inputs(); ++i) {
+    if ((*input_edges)[i] == nullptr) {
+      return errors::InvalidArgument("Missing edge input number: ", i);
+    }
+  }
+  return Status::OK();
+}
+
 Status Node::input_node(int idx, const Node** n) const {
   const Edge* e;
   TF_RETURN_IF_ERROR(input_edge(idx, &e));