Ensured fx nodes cannot be named after keywords (#32)

Fixes a bug where nodes like nir.IF would be called if in the generated Python code

Author: Jegp

nirtorch/nir_interpreter.py

@@ -1,4 +1,5 @@
 import collections
+import keyword
 import typing
 
 import nir
@@ -68,6 +69,14 @@ def _default_map_linear(linear: nir.Linear) -> torch.nn.Linear:
 }
 
 
+def _sanitize_name(name: str) -> str:
+    """Sanitize module name to ensure torch.fx doesn't write any keywords in code"""
+    if keyword.iskeyword(name):
+        return "nir_node_" + name
+    else:
+        return name
+
+
 def _map_nir_node_to_torch(
     node: nir.NIRNode, node_map: NodeMapType
 ) -> typing.Optional[torch.nn.Module]:
@@ -86,11 +95,13 @@ def _construct_module_dict_recursive(
     for name, node in nir_graph.nodes.items():
         # Recurse into subgraphs
         if isinstance(node, nir.NIRGraph):
-            owning_module[name] = _construct_module_dict_recursive(node, node_map)
+            owning_module[_sanitize_name(name)] = _construct_module_dict_recursive(
+                node, node_map
+            )
         else:
             mapped_module = _map_nir_node_to_torch(node, node_map=node_map)
             if mapped_module is not None:
-                owning_module[name] = mapped_module
+                owning_module[_sanitize_name(name)] = mapped_module
     return owning_module
 
 
@@ -138,6 +149,9 @@ def _construct_fx_graph(
 ) -> torch.fx.GraphModule:
     node_outputs = {}
     recursion_counter = collections.Counter(nir_graph.nodes.keys())
+    sanitized_edges = [
+        (_sanitize_name(a), _sanitize_name(b)) for a, b in nir_graph.edges
+    ]
     # The maximum iterations per node (see https://github.com/neuromorphs/NIRTorch/pull/28#discussion_r1959343951)
     max_iterations = min(3, len(recursion_counter))
     torch_graph = torch.fx.Graph(owning_module)
@@ -151,6 +165,10 @@ def _construct_fx_graph(
     # Loop through all the nodes in the queue
     while module_queue:
         module_name, module = module_queue.popleft()
+        # Sanitize the module name to avoid writing keywords in the generated Python code
+        module_name = _sanitize_name(module_name)
+
+        # Test for number of recursions
         if recursion_counter[module_name] > max_iterations:
             raise RecursionError(
                 f"Module {module_name} has been traversed multiple times"
@@ -185,7 +203,7 @@ def _construct_fx_graph(
                 for input_name, output in module.input_type.items():
                     # First fetch the required input nodes
                     module_input_nodes = _find_input_nodes(
-                        module_name, edges=nir_graph.edges, node_outputs=node_outputs
+                        module_name, edges=sanitized_edges, node_outputs=node_outputs
                     )
                     # If the module uses input that is not yet defined, set the inputs to some dummy value
                     # and enqueue the module again for processing (where it's hopefully defined)
@@ -239,7 +257,7 @@ def _construct_fx_graph(
                     # Add the raw output to the graph
                     node_outputs[f"{module_name}_raw"] = output
                     # Add the module state to the graph for use as the input to the next module
-                    node_outputs[f"{module_name}"] = torch_graph.call_method(
+                    node_outputs[module_name] = torch_graph.call_method(
                         "__getitem__", (output, 0)
                     )
                     # Add the state to the state dictionary
@@ -273,6 +291,7 @@ def nir_to_torch(
     Finally, we wrap the execution in a StatefulInterpreter, to ensure that the internal state of modules are handled correctly.
 
     Example:
+
     >>> # First, we describe the NIR graph
     >>> nir_avgpool = nir.AvgPool2d(kernel_size=np.array([2, 2]), stride=np.array([1]), padding=np.array([0, 0]))
     >>> nir_linear = nir.Linear(weight=np.ones((5, 5), dtype=np.float32))

nirtorch/torch_tracer.py

@@ -71,7 +71,9 @@ def call_module(self, target, args, kwargs):
 def torch_to_nir(
     module: torch.nn.Module,
     module_map: Dict[torch.nn.Module, Callable[[torch.nn.Module], nir.NIRNode]],
-    default_dict: Dict[torch.nn.Module, Callable[[torch.nn.Module], nir.NIRNode]] = DEFAULT_MAP,
+    default_dict: Dict[
+        torch.nn.Module, Callable[[torch.nn.Module], nir.NIRNode]
+    ] = DEFAULT_MAP,
 ) -> nir.NIRGraph:
     """
     Traces a PyTorch module and converts it to a NIR graph using the specified module map.

tests/test_nir_interpreter.py

@@ -115,6 +115,35 @@ def test_map_conv2d_node():
     assert torch_conv(torch.ones(1, 3, 10, 11)).shape == (1, 2, 5, 3)
 
 
+def test_map_if_node():
+    # Generating code with the "if" keyword can be sensitive
+    # This tests ensures that it works properly
+    v_th = np.random.random(1)
+    r = np.random.random(1)
+    nir_node = nir.IF(r=r, v_threshold=v_th)
+
+    class MyIF(torch.nn.Module):
+        def __init__(self, r, v_th):
+            super().__init__()
+            self.r = r
+            self.v_th = v_th
+
+        def forward(self, x, state):
+            return x * self.r - self.v_th, state
+
+    node_map = {
+        nir.IF: lambda node: MyIF(
+            torch.from_numpy(node.r), torch.from_numpy(node.v_threshold)
+        )
+    }
+    torch_if = nir_interpreter._map_nir_node_to_torch(nir_node, node_map)
+    data = torch.rand(1)
+    assert torch.allclose(data * r - v_th, torch_if(data, None)[0])
+
+    torch_if = nir_interpreter.nir_to_torch(nir.NIRGraph.from_list(nir_node), node_map)
+    assert isinstance(torch_if.get_submodule("nir_node_if"), MyIF)
+
+
 def test_map_leaky_stateful_graph_single_module():
     # Test that the graph can handle a single stateful module
     tau = np.random.random(1)