done sorting dependencies logic; verified with temp testing

mattport/nerf/graph.py

@@ -3,6 +3,7 @@ The Graph module contains all trainable parameters.
 """
 import importlib
 from dataclasses import dataclass
+from typing import Optional
 
 from torch import nn
 
@@ -20,6 +21,7 @@ class Node:
 
     name: str
     children: dict
+    visited: Optional[bool] = False
 
     def __hash__(self):
         return hash(self.name)
@@ -72,6 +74,19 @@ class Graph(nn.Module):
                     roots.add(curr_module)
         return roots
 
+    def topological_sort(self, curr_node: "Node", ordering_stack: list) -> None:
+        """utility function to sort the call order in the dependency graph
+
+        Args:
+            curr_node (Node): pointer to current node in process
+            ordering_stack (list): cumulative ordering of graph nodes
+        """
+        curr_node.visited = True
+        for child_node in curr_node.children.values():
+            if not child_node.visited:
+                self.topological_sort(child_node, ordering_stack)
+        ordering_stack.append(curr_node.name)
+
     def get_module_order(self) -> list:
         """Generates a graph and determines order of module operations using topological sorting
 
@@ -81,10 +96,13 @@ class Graph(nn.Module):
         Returns:
             list: ordering of the module names that should be executed
         """
+        roots = self.construct_graph()
+        ordering_stack = []
+        for root in roots:
+            if not root.visited:
+                self.topological_sort(root, ordering_stack)
 
-        # call construct graph
-        # function to do topological sort to get ordering
-        return self.modules_config.keys()
+        return ordering_stack[::-1]
 
     def forward(self, x):
         """_summary_

tests/nerf/test_graph.py

@@ -1,5 +1,5 @@
+import numpy as np
 import pytest
-
 from mattport.nerf.graph import Graph
 
 
@@ -135,7 +135,6 @@ def test_construct_graph_repeat():
 
     test_graph.modules_config["mlp_0"]["inputs"][1] = "encoder_1"
     test_graph.modules_config["encoder_1"] = {"class_name": "Encoding", "inputs": ["x"], "meta_data": {"out_dim": 8}}
-    print(test_graph.modules_config["mlp_0"]["inputs"][1])
     roots = test_graph.construct_graph()
     assert len(roots) == 2
 
@@ -157,3 +156,30 @@ def test_construct_graph_complex():
 
     targets_p3 = [("encoder_0", 3), ("mlp_2", 0)]
     check_consistency(roots[0], targets_p3)
+
+
+def test_ordering():
+    """test dependency ordering of constructed graph"""
+    test_basic_graph = Graph(create_basic_graph())
+    _ = test_basic_graph.construct_graph()
+    order = test_basic_graph.get_module_order()
+    target = ["encoder_0", "mlp_0", "mlp_1"]
+    assert np.array_equal(target, order)
+
+    test_repeated_graph = Graph(create_repeat_graph())
+    _ = test_repeated_graph.construct_graph()
+    order = test_repeated_graph.get_module_order()
+    target = ["encoder_0", "mlp_0", "mlp_1"]
+    assert np.array_equal(target, order)
+
+    test_repeated_graph.modules_config["mlp_0"]["inputs"][1] = "encoder_1"
+    test_repeated_graph.modules_config["encoder_1"] = {"class_name": "Encoding", "inputs": ["x"], "meta_data": {"out_dim": 8}}
+    order = test_repeated_graph.get_module_order()
+    assert order.index('mlp_0') > order.index('encoder_0') and order.index('mlp_0') > order.index('encoder_1')
+
+    test_complex_graph = Graph(create_complex_graph())
+    _ = test_complex_graph.construct_graph()
+    order = test_complex_graph.get_module_order()
+    assert order.index('mlp_2') > order.index('mlp_0')
+    assert order.index('mlp_2') > order.index('encoder_0')
+    assert order.index('mlp_3') > order.index('encoder_0')