use structural comparison in topological order for rebuilt computation graph equality tests

Author: TimeDelta

graph_builder.py

@@ -1,6 +1,3 @@
-"""
-Generated by ChatGPT
-"""
 from typing import Dict, List, Tuple
 
 import torch
@@ -147,9 +144,9 @@ def forward(
         self,
         loss: torch.Tensor,
         prev_loss: torch.Tensor,
-        named_params: List[Tuple[str, torch.Tensor]],
+        named_parameters: List[Tuple[str, torch.Tensor]],
     ) -> Dict[str, torch.Tensor]:
-        params = [p for _, p in named_params]
+        params = [p for _, p in named_parameters]
         all_inputs = [loss, prev_loss] + params
         features = torch.stack(all_inputs, 0)
 
@@ -198,8 +195,6 @@ def rebuild_and_script(graph_dict, config, key) -> DynamicOptimizerModule:
 
     # --- build a Python module and script it ---
     if genome.connections:
-        module = DynamicOptimizerModule(
-            genome, config.input_keys, config.output_keys, graph_dict
-        )
+        module = DynamicOptimizerModule(genome, config.input_keys, config.output_keys, graph_dict)
         return torch.jit.script(module)
     return None

population.py

@@ -292,7 +292,8 @@ def evaluate_optimizer(self, optimizer, model, task, steps=10):
           task: The task on which to evaluate the optimizer.
           steps: Number of update iterations.
         """
-        # TODO: find way to correct for time improvements that are solely due to RAM cache tiers
+        # TODO: clear all levels of RAM caches in between every run to create fair starting point
+        # for comparison
         tracemalloc.start()
         start = time.perf_counter()
         prev_metrics_values = torch.tensor([0.0] * len(task.metrics))

tests/test_graph_builder.py

@@ -6,6 +6,7 @@
 import neat
 import pytest
 import torch
+from torch.fx.passes.utils.matcher_utils import SubgraphMatcher
 
 # allow imports from repo root
 sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[1]))
@@ -27,6 +28,104 @@ def make_config():
     )
 
 
+def get_node_signature(node):
+    # simple signature includes kind (operator name), types of inputs, and output type
+    # TODO: for robust comparison, also need to compare attributes and potentially canonicalize constant values
+    input_kinds = [inp.node().kind() for inp in node.inputs()]
+
+    attributes = {}
+    if node.kind() == "prim::Constant":
+        if node.hasAttribute("value"):
+            attributes["value"] = node.t("value")
+        elif node.hasAttribute("i"):
+            attributes["value"] = node.i("i")
+        elif node.hasAttribute("f"):
+            attributes["value"] = node.f("f")
+    # TODO: finish
+
+    return (node.kind(), tuple(input_kinds), node.output().type(), tuple(sorted(attributes.items())))
+
+
+def compare_jit_graphs_structural(original: torch.jit.ScriptModule, rebuilt: torch.jit.ScriptModule) -> bool:
+    original_inputs = list(original.graph.inputs())
+    rebuilt_inputs = list(rebuilt.graph.inputs())
+    original_outputs = list(original.graph.outputs())
+    rebuilt_outputs = list(rebuilt.graph.outputs())
+    if len(original_inputs) != len(rebuilt_inputs) or len(original_outputs) != len(rebuilt_outputs):
+        print(
+            f"Input/output counts differ: original.graph inputs={len(original_inputs)}, outputs={len(original_outputs)} vs rebuilt inputs={len(rebuilt_inputs)}, outputs={len(rebuilt_outputs)}",
+            file=sys.stderr,
+        )
+        return False
+
+    # default iterator for graph.nodes() is typically a topological sort
+    original_nodes = list(original.graph.nodes())
+    rebuilt_nodes = list(rebuilt.graph.nodes())
+
+    if len(original_nodes) != len(rebuilt_nodes):
+        print(
+            f"Number of nodes differ: original.graph has {len(original_nodes)} nodes, rebuilt has {len(rebuilt_nodes)} nodes",
+            file=sys.stderr,
+        )
+        return False
+
+    # create mapping from nodes to canonical representation based on signature + inputs
+    original_node_map = {}
+    rebuilt_node_map = {}
+    for i, (original_node, rebuilt_node) in enumerate(zip(original_nodes, rebuilt_nodes)):
+        signature1 = get_node_signature(original_node)
+        signature2 = get_node_signature(rebuilt_node)
+
+        if signature1 != signature2:
+            print(f"Signatures differ at node {i}:", file=sys.stderr)
+            print(f"  original.graph Node Kind: {original_node.kind()}", file=sys.stderr)
+            print(f"  rebuilt Node Kind: {rebuilt_node.kind()}", file=sys.stderr)
+            # TODO: add more detailed diffing here
+            return False
+
+        # assumes a consistent order of inputs and that corresponding inputs have corresponding nodes
+        for input_idx, (original_input_val, rebeuilt_input_val) in enumerate(
+            zip(original_node.inputs(), rebuilt_node.inputs())
+        ):
+            if original_input_val.node().kind() != rebeuilt_input_val.node().kind():
+                print(f"Input kind differs for node {i}, input {input_idx}", file=sys.stderr)
+                return False
+            # TODO: need to further compare value properties if they are constants or recursively
+            # check if the input nodes themselves are structurally equivalent up to that point
+
+    original_params = dict(original.named_parameters())
+    rebuilt_params = dict(rebuilt.named_parameters())
+    if len(original_params) != len(rebuilt_params):
+        print("Parameter counts differ", file=sys.stderr)
+        return False
+    for name, original_param in original_params.items():
+        if name not in rebuilt_params:
+            print(f"Parameter '{name}' missing in rebuilt graph", file=sys.stderr)
+            return False
+        rebuilt_param = rebuilt_params[name]
+        if not torch.equal(original_param, rebuilt_param):
+            print(f"Parameter '{name}' values differ", file=sys.stderr)
+            return False
+
+    if not compare_custom_data(original, rebuilt):
+        print("Custom data attributes differ", file=sys.stderr)
+        return False
+
+    return True
+
+
+def compare_custom_data(original: torch.jit.ScriptModule, rebuilt: torch.jit.ScriptModule) -> bool:
+    if hasattr(original, "node_types") and hasattr(rebuilt, "node_types"):
+        if original.node_types != rebuilt.node_types:
+            print("node_types differ", file=sys.stderr)
+            return False
+    if hasattr(original, "edge_index") and hasattr(rebuilt, "edge_index"):
+        if not torch.equal(original.edge_index, rebuilt.edge_index):
+            print("edge_index differ", file=sys.stderr)
+            return False
+    return True
+
+
 @pytest.mark.parametrize("pt_path", glob.glob(os.path.join("computation_graphs", "optimizers", "*.pt")))
 def test_graph_builder_rebuilds_pt(pt_path):
     original = torch.jit.load(pt_path)
@@ -51,11 +150,5 @@ def test_graph_builder_rebuilds_pt(pt_path):
     assert len(list(rebuilt.parameters())) == len(expected_edges)
     assert len(rebuilt.node_types) == len(data.node_types)
 
-    # Verify that the rebuilt computation graph is identical to the original
-    if str(rebuilt.graph) != str(original.graph):
-        print("Original graph:\n", original.graph)
-        print("Rebuilt graph:\n", rebuilt.graph)
-    assert str(rebuilt.graph) == str(original.graph), (
-        "\nOriginal graph:\n" + str(original.graph) +
-        "\nRebuilt graph:\n" + str(rebuilt.graph)
-    )
+    # Verify that the rebuilt computation graph is structurally identical to the original
+    assert compare_jit_graphs_structural(rebuilt, original)