Makes the private all_graph_edges() function more memory efficient.

PiperOrigin-RevId: 294959528

Author: Neural-Link Team

neural_structured_learning/tools/graph_utils.py

@@ -89,11 +89,13 @@ def add_undirected_edges(graph):
     `None`. Instead, this function has a side-effect on the `graph` argument.
   """
   def all_graph_edges():
-    edges = []
-    for s, t_dict in six.iteritems(graph):
-      for t, w in six.iteritems(t_dict):
-        edges.append((s, t, w))
-    return edges
+    # Make a copy of all source IDs to avoid concurrent iteration failure.
+    sources = list(graph.keys())
+    for source in sources:
+      # Make a copy of source's out-edges to avoid concurrent iteration failure.
+      out_edges = dict(graph[source])
+      for target, weight in six.iteritems(out_edges):
+        yield (source, target, weight)
 
   start_time = time.time()
   logging.info('Making all edges bi-directional...')

neural_structured_learning/tools/graph_utils_test.py

@@ -22,7 +22,7 @@
 from absl.testing import absltest
 from neural_structured_learning.tools import graph_utils
 
-GRAPH = {'A': {'B': 0.5, 'C': 0.9}, 'B': {'A': 0.4, 'C': 1.0}}
+GRAPH = {'A': {'B': 0.5, 'C': 0.9}, 'B': {'A': 0.4, 'C': 1.0}, 'D': {'A': 0.75}}
 
 
 class GraphUtilsTest(absltest.TestCase):
@@ -35,6 +35,7 @@ def testAddEdge(self):
     graph_utils.add_edge(graph, ['A', 'C', 0.8])  # ...is used.
     graph_utils.add_edge(graph, ('B', 'A', '0.4'))
     graph_utils.add_edge(graph, ('B', 'C'))  # Tests default weight
+    graph_utils.add_edge(graph, ('D', 'A', 0.75))
     self.assertDictEqual(graph, GRAPH)
 
   def testAddUndirectedEdges(self):
@@ -44,15 +45,19 @@ def testAddUndirectedEdges(self):
         g_actual, {
             'A': {
                 'B': 0.5,
-                'C': 0.9
+                'C': 0.9,
+                'D': 0.75
             },
             'B': {
                 'A': 0.5,  # Note, changed from 0.4 to 0.5
                 'C': 1.0
             },
-            'C': {
+            'C': {         # Added
                 'A': 0.9,  # Added
                 'B': 1.0   # Added
+            },
+            'D': {
+                'A': 0.75
             }
         })