From a0a2f47465ea5858b949eebc3a08cc6920b7161f Mon Sep 17 00:00:00 2001
From: "Jens E. Pedersen" <jensegholm@protonmail.com>
Date: Tue, 10 Oct 2023 21:49:39 +0200
Subject: [PATCH 01/27] Added option to execute stateful submodules

---
 nirtorch/from_nir.py   | 41 +++++++++++++++++++++++++++++++++++------
 nirtorch/graph.py      |  4 +++-
 tests/test_from_nir.py | 23 +++++++++++++++++++++--
 3 files changed, 59 insertions(+), 9 deletions(-)

diff --git a/nirtorch/from_nir.py b/nirtorch/from_nir.py
index 6bdc381..773ba11 100644
--- a/nirtorch/from_nir.py
+++ b/nirtorch/from_nir.py
@@ -1,4 +1,5 @@
-from typing import Callable, Dict, List, Optional
+import inspect
+from typing import Callable, Dict, List, Optional, NamedTuple, Any
 
 import nir
 import torch
@@ -49,11 +50,17 @@ class GraphExecutor(nn.Module):
     def __init__(self, graph: Graph) -> None:
         super().__init__()
         self.graph = graph
+        self.stateful_modules = {}
         self.instantiate_modules()
         self.execution_order = self.get_execution_order()
         if len(self.execution_order) == 0:
             raise ValueError("Graph is empty")
 
+    def _is_module_stateful(self, module: torch.nn.Module) -> bool:
+        signature = inspect.signature(module.forward)
+        arguments = len(signature.parameters)
+        return arguments > 1
+
     def get_execution_order(self) -> List[Node]:
         """Evaluate the execution order and instantiate that as a list."""
         execution_order = []
@@ -67,13 +74,32 @@ def get_execution_order(self) -> List[Node]:
 
     def instantiate_modules(self):
         for mod, name in self.graph.module_names.items():
-            self.add_module(sanitize_name(name), mod)
+            if mod is not None:
+                self.add_module(sanitize_name(name), mod)
+                self.stateful_modules[sanitize_name(name)] = self._is_module_stateful(
+                    mod
+                )
 
     def get_input_nodes(self) -> List[Node]:
         # NOTE: This is a hack. Should use the input nodes from NIR graph
         return self.graph.get_root()
 
-    def forward(self, data: torch.Tensor):
+    def _apply_module(
+        self, node: Node, x: torch.Tensor, state: Optional[Dict[str, Any]]
+    ):
+        """Applies a module and keeps track of its state.
+        TODO: Use pytree to recursively construct the state
+        """
+        if node.name in self.stateful_modules and node.name in state:
+            out = node.elem(x, *state[node.name])
+        else:
+            out = node.elem(x)
+        if self.stateful_modules[node.name]:
+            state[node.name] = out[1:]
+            out = out[0]
+        return out
+
+    def forward(self, data: torch.Tensor, state: Optional[Dict[str, Any]] = {}):
         outs = {}
         # NOTE: This logic is not yet consistent for models with multiple input nodes
         for node in self.execution_order:
@@ -82,11 +108,14 @@ def forward(self, data: torch.Tensor):
                 continue
             if len(input_nodes) == 0 or len(outs) == 0:
                 # This is the root node
-                outs[node.name] = node.elem(data)
+                outs[node.name] = self._apply_module(node, data, state)
             else:
                 # Intermediate nodes
-                input_data = (outs[node.name] for node in input_nodes)
-                outs[node.name] = node.elem(*input_data)
+                input_data = [outs[node.name] for node in input_nodes]
+                input_data = torch.stack(input_data).sum(
+                    0
+                )  # Multiple inputs are summed
+                outs[node.name] = self._apply_module(node, input_data, state)
         return outs[node.name]
 
 
diff --git a/nirtorch/graph.py b/nirtorch/graph.py
index ac5c7c1..aae3e61 100644
--- a/nirtorch/graph.py
+++ b/nirtorch/graph.py
@@ -5,6 +5,8 @@
 import torch
 import torch.nn as nn
 
+from .utils import sanitize_name
+
 
 def named_modules_map(
     model: nn.Module, model_name: Optional[str] = "model"
@@ -41,7 +43,7 @@ def __init__(
         outgoing_nodes: Optional[Dict["Node", torch.Tensor]] = None,
     ) -> None:
         self.elem = elem
-        self.name = name
+        self.name = sanitize_name(name)
         if not outgoing_nodes:
             self.outgoing_nodes = {}
         else:
diff --git a/tests/test_from_nir.py b/tests/test_from_nir.py
index d6d09ae..717d7d6 100644
--- a/tests/test_from_nir.py
+++ b/tests/test_from_nir.py
@@ -1,7 +1,7 @@
 import nir
 import numpy as np
-import pytest
 import torch
+import pytest
 
 from nirtorch.from_nir import load
 
@@ -29,7 +29,7 @@ def test_extract_empty():
 
 
 def test_extract_illegal_name():
-    graph = nir.NIRGraph({"a.b": nir.Input(np.ones(1))}, [])
+    graph = nir.NIRGraph({"a.b": nir.Linear(np.ones((1, 1)))}, [])
     torch_graph = load(graph, _torch_model_map)
     assert "a_b" in torch_graph._modules
 
@@ -62,3 +62,22 @@ def test_extrac_recurrent():
     m = load(g, _torch_model_map)
     data = torch.randn(1, 1, dtype=torch.float64)
     torch.allclose(m(data), l2(l1(data)))
+
+
+def test_execute_stateful():
+    class StatefulModel(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+
+        def forward(self, x, state=None):
+            if state is None:
+                state = 1
+            return x + state, state
+
+    g = nir.NIRGraph(
+        nodes={"li": nir.Flatten(np.array([1])), "li2": nir.Flatten(np.array([1]))},
+        edges=[("li", "li2")],
+    )  # Mock node
+    m = load(g, lambda m: StatefulModel())
+    out = m(torch.ones(10))
+    assert torch.allclose(out, torch.ones(10) * 3)

From 1465f2985538ad2ee8c47e4c1076f51559a92bc8 Mon Sep 17 00:00:00 2001
From: "Jens E. Pedersen" <jensegholm@protonmail.com>
Date: Tue, 10 Oct 2023 21:53:27 +0200
Subject: [PATCH 02/27] Returned state if stateful module

---
 nirtorch/from_nir.py   | 5 ++++-
 tests/test_from_nir.py | 4 +++-
 2 files changed, 7 insertions(+), 2 deletions(-)

diff --git a/nirtorch/from_nir.py b/nirtorch/from_nir.py
index 773ba11..e867129 100644
--- a/nirtorch/from_nir.py
+++ b/nirtorch/from_nir.py
@@ -116,7 +116,10 @@ def forward(self, data: torch.Tensor, state: Optional[Dict[str, Any]] = {}):
                     0
                 )  # Multiple inputs are summed
                 outs[node.name] = self._apply_module(node, input_data, state)
-        return outs[node.name]
+        if len(state) > 0:
+            return outs[node.name], state
+        else:
+            return outs[node.name]
 
 
 def _mod_nir_to_graph(nir_graph: nir.NIRNode) -> Graph:
diff --git a/tests/test_from_nir.py b/tests/test_from_nir.py
index 717d7d6..aad54d0 100644
--- a/tests/test_from_nir.py
+++ b/tests/test_from_nir.py
@@ -79,5 +79,7 @@ def forward(self, x, state=None):
         edges=[("li", "li2")],
     )  # Mock node
     m = load(g, lambda m: StatefulModel())
-    out = m(torch.ones(10))
+    out, state = m(torch.ones(10))
     assert torch.allclose(out, torch.ones(10) * 3)
+    assert state["li"] == (1, )
+    assert state["li"] == (1, )

From bbe54a0602899515a61072ccb031dae9fecb1bb9 Mon Sep 17 00:00:00 2001
From: "Jens E. Pedersen" <jensegholm@protonmail.com>
Date: Tue, 10 Oct 2023 21:55:55 +0200
Subject: [PATCH 03/27] Ruff

---
 nirtorch/from_nir.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/nirtorch/from_nir.py b/nirtorch/from_nir.py
index e867129..dff21fc 100644
--- a/nirtorch/from_nir.py
+++ b/nirtorch/from_nir.py
@@ -1,5 +1,5 @@
 import inspect
-from typing import Callable, Dict, List, Optional, NamedTuple, Any
+from typing import Callable, Dict, List, Optional, Any
 
 import nir
 import torch

From a64c23d473561d4e8947cc84864b4cada38f937c Mon Sep 17 00:00:00 2001
From: "Jens E. Pedersen" <jensegholm@protonmail.com>
Date: Wed, 11 Oct 2023 23:55:01 +0200
Subject: [PATCH 04/27] Added recurrent execution

---
 nirtorch/from_nir.py   | 81 ++++++++++++++++++++++++++++++------------
 tests/test_from_nir.py | 42 ++++++++++++++++------
 2 files changed, 89 insertions(+), 34 deletions(-)

diff --git a/nirtorch/from_nir.py b/nirtorch/from_nir.py
index dff21fc..4b3302e 100644
--- a/nirtorch/from_nir.py
+++ b/nirtorch/from_nir.py
@@ -1,3 +1,4 @@
+import dataclasses
 import inspect
 from typing import Callable, Dict, List, Optional, Any
 
@@ -46,6 +47,16 @@ def execution_order_up_to_node(
         return execution_order + [node]
 
 
+@dataclasses.dataclass
+class GraphExecutorState:
+    """State for the GraphExecutor that keeps track of both the state of hidden
+    units and caches the output of previous modules, for use in (future) recurrent
+    computations."""
+
+    state: Dict[str, Any] = dataclasses.field(default_factory=dict)
+    cache: Dict[str, Any] = dataclasses.field(default_factory=dict)
+
+
 class GraphExecutor(nn.Module):
     def __init__(self, graph: Graph) -> None:
         super().__init__()
@@ -85,41 +96,65 @@ def get_input_nodes(self) -> List[Node]:
         return self.graph.get_root()
 
     def _apply_module(
-        self, node: Node, x: torch.Tensor, state: Optional[Dict[str, Any]]
+        self,
+        node: Node,
+        input_nodes: List[Node],
+        old_state: GraphExecutorState,
+        new_state: GraphExecutorState,
+        data: Optional[torch.Tensor] = None,
     ):
         """Applies a module and keeps track of its state.
         TODO: Use pytree to recursively construct the state
         """
-        if node.name in self.stateful_modules and node.name in state:
-            out = node.elem(x, *state[node.name])
-        else:
-            out = node.elem(x)
+        inputs = []
+        # Append state if needed
+        if node.name in self.stateful_modules and node.name in old_state.state:
+            inputs.extend(old_state.state[node.name])
+
+        # Sum recurrence if needed
+        summed_inputs = [] if data is None else [data]
+        for input_node in input_nodes:
+            if (
+                input_node.name not in new_state.cache
+                and input_node.name in old_state.cache
+            ):
+                summed_inputs.append(old_state.cache[input_node.name])
+            elif input_node.name in new_state.cache:
+                summed_inputs.append(new_state.cache[input_node.name])
+
+        inputs.insert(
+            0, torch.stack(summed_inputs).sum(0)
+        )  # Insert input, sum if multiple
+
+        out = node.elem(*inputs)
+        # If the module is stateful, we know the output is (at least) a tuple
         if self.stateful_modules[node.name]:
-            state[node.name] = out[1:]
+            new_state.state[node.name] = out[1:]  # Store the new state
             out = out[0]
-        return out
+        return out, new_state
 
-    def forward(self, data: torch.Tensor, state: Optional[Dict[str, Any]] = {}):
-        outs = {}
+    def forward(
+        self, data: torch.Tensor, old_state: Optional[GraphExecutorState] = None
+    ):
+        if old_state is None:
+            old_state = GraphExecutorState()
+        new_state = GraphExecutorState()
+        first_node = True
         # NOTE: This logic is not yet consistent for models with multiple input nodes
         for node in self.execution_order:
             input_nodes = self.graph.find_source_nodes_of(node)
             if node.elem is None:
                 continue
-            if len(input_nodes) == 0 or len(outs) == 0:
-                # This is the root node
-                outs[node.name] = self._apply_module(node, data, state)
-            else:
-                # Intermediate nodes
-                input_data = [outs[node.name] for node in input_nodes]
-                input_data = torch.stack(input_data).sum(
-                    0
-                )  # Multiple inputs are summed
-                outs[node.name] = self._apply_module(node, input_data, state)
-        if len(state) > 0:
-            return outs[node.name], state
-        else:
-            return outs[node.name]
+            out, new_state = self._apply_module(
+                node, input_nodes, new_state, old_state, data if first_node else None
+            )
+            new_state.cache[node.name] = out
+            first_node = False
+
+        # If the output node is a dummy nir.Output node, use the second-to-last node
+        if node.name not in new_state.cache:
+            node = self.execution_order[-2]
+        return new_state.cache[node.name], new_state
 
 
 def _mod_nir_to_graph(nir_graph: nir.NIRNode) -> Graph:
diff --git a/tests/test_from_nir.py b/tests/test_from_nir.py
index aad54d0..4ae9a33 100644
--- a/tests/test_from_nir.py
+++ b/tests/test_from_nir.py
@@ -9,12 +9,12 @@
 def _torch_model_map(m: nir.NIRNode, device: str = "cpu") -> torch.nn.Module:
     if isinstance(m, nir.Affine):
         lin = torch.nn.Linear(*m.weight.shape[-2:])
-        lin.weight.data = torch.nn.Parameter(torch.tensor(m.weight).to(device))
-        lin.bias.data = torch.nn.Parameter(torch.tensor(m.bias).to(device))
+        lin.weight.data = torch.nn.Parameter(torch.tensor(m.weight).to(device).float())
+        lin.bias.data = torch.nn.Parameter(torch.tensor(m.bias).to(device).float())
         return lin
     elif isinstance(m, nir.Linear):
         lin = torch.nn.Linear(*m.weight.shape[-2:], bias=False)
-        lin.weight.data = torch.nn.Parameter(torch.tensor(m.weight).to(device))
+        lin.weight.data = torch.nn.Parameter(torch.tensor(m.weight).to(device).float())
         return lin
     elif isinstance(m, nir.Input) or isinstance(m, nir.Output):
         return None
@@ -46,22 +46,25 @@ def test_extract_lin():
     assert isinstance(m.execution_order[0].elem, torch.nn.Linear)
     assert torch.allclose(m.execution_order[0].elem.weight, lin.weight)
     assert torch.allclose(m.execution_order[0].elem.bias, lin.bias)
-    assert torch.allclose(m(x), y)
+    assert isinstance(m.execution_order[1].elem, torch.nn.Linear)
+    assert torch.allclose(m.execution_order[1].elem.weight, lin.weight)
+    assert torch.allclose(m.execution_order[1].elem.bias, lin.bias)
+    assert torch.allclose(m(x)[0], y)
 
 
-def test_extrac_recurrent():
+def test_extract_recurrent():
     w = np.random.randn(1, 1)
     g = nir.NIRGraph(
         nodes={"in": nir.Input(np.ones(1)), "a": nir.Linear(w), "b": nir.Linear(w)},
         edges=[("in", "a"), ("a", "b"), ("b", "a")],
     )
     l1 = torch.nn.Linear(1, 1, bias=False)
-    l1.weight.data = torch.tensor(w)
+    l1.weight.data = torch.tensor(w).float()
     l2 = torch.nn.Linear(1, 1, bias=False)
-    l2.weight.data = torch.tensor(w)
+    l2.weight.data = torch.tensor(w).float()
     m = load(g, _torch_model_map)
-    data = torch.randn(1, 1, dtype=torch.float64)
-    torch.allclose(m(data), l2(l1(data)))
+    data = torch.randn(1, 1, dtype=torch.float32)
+    torch.allclose(m(data)[0], l2(l1(data)))
 
 
 def test_execute_stateful():
@@ -81,5 +84,22 @@ def forward(self, x, state=None):
     m = load(g, lambda m: StatefulModel())
     out, state = m(torch.ones(10))
     assert torch.allclose(out, torch.ones(10) * 3)
-    assert state["li"] == (1, )
-    assert state["li"] == (1, )
+    assert state.state["li"] == (1,)
+    assert state.state["li"] == (1,)
+
+
+def test_execute_recurrent():
+    w = np.ones((1, 1))
+    g = nir.NIRGraph(
+        nodes={"in": nir.Input(np.ones(1)), "a": nir.Linear(w), "b": nir.Linear(w)},
+        edges=[("in", "a"), ("a", "b"), ("b", "a")],
+    )
+    m = load(g, _torch_model_map)
+    data = torch.ones(1, 1)
+
+    # Same execution without reusing state should yield the same result
+    y1 = m(data)
+    y2 = m(data)
+    assert torch.allclose(y1[0], y2[0])
+    out, s = m(*m(data))
+    assert torch.allclose(out, torch.tensor(2.0))

From 0a1d97d0fd79f0d9d0f1a6cc953f58070835f7c7 Mon Sep 17 00:00:00 2001
From: "Jens E. Pedersen" <jensegholm@protonmail.com>
Date: Thu, 12 Oct 2023 22:32:29 +0200
Subject: [PATCH 05/27] Added tests for recurrent execution

---
 nirtorch/from_nir.py   |  11 +++++++----
 tests/braille.nir      | Bin 0 -> 41896 bytes
 tests/test_from_nir.py |  33 +++++++++++++++++++++++++++++++++
 3 files changed, 40 insertions(+), 4 deletions(-)
 create mode 100644 tests/braille.nir

diff --git a/nirtorch/from_nir.py b/nirtorch/from_nir.py
index 4b3302e..349a1d5 100644
--- a/nirtorch/from_nir.py
+++ b/nirtorch/from_nir.py
@@ -1,6 +1,6 @@
 import dataclasses
 import inspect
-from typing import Callable, Dict, List, Optional, Any
+from typing import Callable, Dict, List, Optional, Any, Union
 
 import nir
 import torch
@@ -173,18 +173,21 @@ def _switch_models_with_map(
 
 
 def load(
-    nir_graph: nir.NIRNode, model_map: Callable[[nir.NIRNode], nn.Module]
+    nir_graph: Union[nir.NIRNode, str], model_map: Callable[[nir.NIRNode], nn.Module]
 ) -> nn.Module:
-    """Load a NIR object and convert it to a torch module using the given model map.
+    """Load a NIR graph and convert it to a torch module using the given model map.
 
     Args:
-        nir_graph (nir.NIRNode): NIR object
+        nir_graph (Union[nir.NIRNode, str]): The NIR object to load, or a string representing
+            the path to the NIR object.
         model_map (Callable[[nn.NIRNode], nn.Module]): A method that returns the a torch
             module that corresponds to each NIR node.
 
     Returns:
         nn.Module: The generated torch module
     """
+    if isinstance(nir_graph, str):
+        nir_graph = nir.read(nir_graph)
     # Map modules to the target modules using th emodel map
     nir_module_graph = _switch_models_with_map(nir_graph, model_map)
     # Build a nirtorch.Graph based on the nir_graph
diff --git a/tests/braille.nir b/tests/braille.nir
new file mode 100644
index 0000000000000000000000000000000000000000..ed8f57e7f4909bd10ea0e92cf2fba67f6c7c4efe
GIT binary patch
literal 41896
zcmeHQ3tUah_dh*Fl0+({Bd-t@A+={trAQ(q@-7uc5{U>UQj&KmkCZ$jg^H;5tT{vv
zk@qV^@`{jm-v5qs_Prbz_ul`#_vik8?LM6~vu5_pTHig-HM3_gv9sx;(m=ZbPjo3M
z@#J}pKRLOdkCJQ07BPqObK$NMhm3?qNazk(u|tL@`s495C44nW`qJ@+4zac6)yCw~
z=cwd1p0T{xrNlAn=>HS~cD7c7L<+k|wC2iFp`@H~V9fNK?(aR-m*+dx!&AiL%B3Y2
z;SW+7o`N_HNkTF_aS7zom8&h;BDa=|%%_qqsy*T`aXf|E_#(foteCANsR?piLQ?gO
z$BX-e68^K%od!DE$na!;kvdR8oIi(0qQu`=@_@Md%9s0R<^0)m`&PRsjo**j#W*do
zrX0>cyWVy3{j~^us~p|XmgJ-R{rp7LwXvk|i0W|P-i{8R>at`Qz}0I>wdNUj>DtA(
zHXhgPa^-Gdzaf36yZL!>)trmp&(|+tMs0ja_svtRb&D>m2@|}1J$bdQtN4=UbN+gg
z@ZJIA-5mP$;c<THd^tat9+z)@$^G1YTsq<b5AQ2^l%?cx44Cmrfrb)3r=O3v?<bDR
zSNsW|IE}u+SC#NN{iaINm%z>YHGOe9YG0)z&bP@|d~tkr37^BkTY`U6seC!Tnn|TI
z!M*Dzj>b3m%_V&9dHQ%y_;dr;A4~Hk&C@S9bNX_{lf%Kw+rz`NR<#zA^riVaU&&+g
ziNh6d4z~&JVtz|Wd}$mcY#yI`a2%tK)*(=bKpg^g2-G1^hd><ybqLfUP=`Pr0{>$O
zNZVQY68(+237*71W452C$SJb3`qxrVkBOfsXD=|5#QVME@3XTu=8I+EX4srWah&6)
zq+&a($hs`P;*!2ECa0ZHV(8$WgR~upvy(({34cG?e|_kW?57b1ViiR;t-9-f4*}`+
zj%&C0oMA}<31>&+>`R=TjI)1n_AAc5#nnIlHz56%-8NB@Z(~Wh{89P6(NUZMSFgX3
zT^;+s6oGG*BeTw8g<Xuq^jq=lOvT||J;YR0j=1#zk&RzuKjr))8!KlA7ujbyzR2cl
z=ot(pvmzn!pc>tBs2rj)`82%r0t6{Fqg$O*;cdud7*M1`g@K`P;0g<k_Q?o7KA!_@
z|1g+c!pHp<pT?%g!%=9pJniFY4JVr1$H_$xa*XaIfcte>dUZl;p+})BbhcU0Yu&CY
z;=?>Z1sC=49oB)w$2bTVI{4BIdw>_SjB&=8du;ZZN<aq*+4gKPyf*p(=j`&qE20J-
zl_IF=y@I?8*$Fbwm1&!P417??hhU}pg7x<`p~MBz*FHVyffhsI)-5CA0}g^^12SP|
zuq~B$dO?;M?t#2;TjK4K3_cs%(EX_dl!v~AD@qw)miYwMEenDq4|O_z)O+Z*?=s9U
zj)a-Uqo7F7R?z&lIz6>V3+O~Qq7i0}<&=GhmBSzyRHYB|Tnq*4Bm8NP<xSzplz6P~
zB1h%M?WUdDPp1Pq9HFOWw$T%t7t_+f2s-3QG7U?qPuHG4$qX;gpzh~~&|F&q&24*&
z9-e=Lj_x2M6uc~_XPcj<efC_V_q`Q_hc*_|<!LWzM$c+$zb=s$hF8#yPH$+2aSE-s
zVlS;yy+S9KD+skV+@Shiujt~HdO}Sjp0H$ytnm1(P?~wLnD!qQhi44mg<HC;$8mWF
z;mp<s*zv+$vg^({l%y6-qmRy~IUZT$WY<h)IqD=Fu{Q&*JZVek?OQ}w+*TA$^*TpG
z<8Q*rk~A8VU`<aS*A;rE`-3o0N!T=`n^3h{4xCmLV4Z>u%yuhX9QnZ(+UYxj)!tLM
z%hoXhepLZ}@3(_7T0-cgweHC2kv(i@ond8UK6q<hC;3PXRF^q`k<SH?X|xA=%{dOK
zX=ecf!tk5%eFdds-@?toI`A}d6!_Y&!AtX01U)Bhf^osGi9;+BruA^8FOKTvI-yI{
zDd;GuN2fwqO)-oVG@{dDazVY~Fxb4yf@eW5;Qh_(5SOD3V_V&aQ-e=}Lh1sz*dzuL
ztY3nMrV^chydJ%GPliseDFoTH3gFK&q)G*{bZ_evP*5OXoJHZnSp#Z$v^8C<j_3;W
z8jw2{3v(O!11xt3SM~NZ*h!6cO_HbXX7zIIXAGtnqM8bC=Del*wQK13DNpDl;UQZ3
zPDg0;wv^8AM`%Y}4V_B&((+M5sny|Zsw3-9BR!SriehshhtJ}93oDS!HXS<XWE&hA
zp)H6!=#DEVJOJy8J(R#Kx}oqWZJnA$k1ab#ug=&{&5JJ6CwGde$>r6w-|0S7_sLRN
z*f^DDHc6+Bf}1qp<7<khe4tVFWP}I#l*(>=NOP=<XtCiDnxL6R*F7~89t|m=jg!=c
zlkeq`IG;fRnUoUph75%>!zk_1u_b8~^%&dkhyruD%j{U@4cgdsHM(lM6U;~Sp)#xb
z!t~>p@Uuz!w7k3lot9GDrfld)^v^dFB=mny$GpwOIr($wf(<5u<)$a`suq4&(dr&q
z(mo8!9MFZiR`$ZC!))-9xp{D<Z#0_SWE38E(}}$8dJ&#gw5BVPvgq;7DRf=hc{;S`
z2|6KrFm2_upT2KCjRtqzK<CJ4QQM(rbYiR!UZ|_l?gi;^?9Cezv@Qo;9CCv|&EC*;
z@B<jVc^8b5ZAcT>?|`#SL!hSoGTC=F5Jhx2OSZ=b(f*;0AuYE)ousQql@m5&pYh6Y
z(x5S2pl&GCQp_gZ6jJGk=O3xX+>Nwq7fbv2ZKtpH=FpV9T)J83EFHBgpK33#r5RRH
z^hG~ws?k$HIAf;_9q`&8qSj3XZL?^Y7i<K(^EF{nfG+i(xewm;iG=*>)3En|0=1Dj
zL+1BiMupLbsjABfY9O<k8n)D?>5dC%`rQy}XtkMoW~I{<`RVlhByHONXawkM$Wi{b
zqhNnWl{T8DPe+|ehU<A=@Os`X!VGUsx9G(IPe+@&+J=D2`9|o)VmVsXR0W%cEuq0n
zfgY-Iz#m7|r*8tXkpmdOoNk4%e6YVjZT1=%b>KN_s~JZNvx10mXnnH7r8DdvXa`U2
z)X*N=QVWM7MZ`}&Kp!4hKo#4rh7t3t$efF7VPN!f{5UQF@}u-2`NRcmFm51D))Qj6
znLD7H^&r9WU5ltvaw4j>eL#=xSD+@UgX!($DKy~BW$*}(23F9Dx_;<OC+!cXO2#8;
zbh{JKZ*@FG+&u+t8r{Zm{!>Z!-2-T!G2L+RNfX>~@?pAd-*Gg(#vDf^cZ0CS?P*l=
zayT2WLid$N!-fH~>2$}{bi&PjU@=yWjx=3Q_YLu&)7{c(>Wc03`FK+bMe=lIea+e?
zUzf)4!ts^66)5paOWMJqEh_0)35T2N<B!(YXnN8$N`}b`N9W9i$_5MQ`>QkYUd>}<
zOwm>>zd8n%mAS*UG)3Cso}n;4Wh*)qSw`Ok>InBL4TF;DNoaomAS~PCBvn%HMpd_N
zMIQ?e(gOq3=@$Js8ZT1`vz{8_u!thG?{X-af4miSyO(Uy?TS79ydM|!_0r>EZklK;
z8IRUR`E+sfe}6BxPlf<)ylW{5|D!}a-&Oo&Uy)ExoC24w^mvAo>z-6R>G2FFw$S;@
z^zKRgN-knPH+~M3jK8^jrN`rZade(b?=RB}k@ypB#e8l)X&3xuyts{G|G`aS`d9z{
zeRdQ+Uo0zUY2YM};~YOFmENyrVEBD8xpHu-!}sI;uZ4bHuEaG^rgq<>=;G=NN9Ag3
zvdAr(BZx7z6h*6@PjLGoWj{szbw8vr|5F4>Y_VJPyYg4(PvrXHyj;=%NY|F1NQF94
z{ze3(`%%6muRzIsoXeW)M`w9@PxPwYC^z2QO=M30r9PQ)5q~HBx}J(_->>!Z&+b<?
z<BLU*F46yn_WxHqx>>jH1Nlq)FZG<H@8VD}5O0^04*sw6e_gH&KIteUZpL}-BpP$}
zA!$FiPhO3$C)o?i=QUU2HJ1#JzSA$*@x@8=xh<TMoB}2EKSRF8<IDb#)1RY%UVpFl
zVhyG1S;T)v@o(h%qyDhEr}##$-h3mOI`)4l0^jNnBP4cbPVIkBf5;OzfWORGISr)i
zDaZdgsPwwBo%Q$S;`+lIQqM{HF0Nji^^y+x@AQA&{(spYa-9{|{!9D+e1BkPE=JUV
z+iEWX_&?hpen`wEeJ?f_6f~@$`8l2O*d1DMDX;|!F$@;;$-aXs9oHc@xq;-Nnx1e9
z4kPYuJ!wcbAI}=!9eV1{L@9R9k#K=7N#3Lj`zEfWm$S~3&O=yax?+g%nN~XXZLo~U
z-+9Z<I~k2;F36xa=htT|o$N@0Lk(*^dKucXpc(45#*(fbC!}w@qAf<Q?nz^{O4(4~
zhs0B^n!2e{!J*cBNmlD(Xgy2|n|LVEs>DU)c9tq0ka~!`@h!$GZ=W(pjo!o1W<kW2
zZH5P(At+W}4_C^`;CR#h#I}nT)pwmkwB#R=YIYJfO^z4D*uBD$)*Wy!(-5>X+l6ji
z_m*t7Q-;OEccYo>cA*o*Lhy8?5UcZVTTHSyM5P~G>9QV`tjVm2jQVXA7`E0QZ;*e$
zwg_|*=m|!ks+vPIPUsA&O7e`xgIHF(cO~+?8b`9m+rg{}jfk9GxWKBLJeVizqw8{=
z*moVald;d^v3D~i810~lTgfX3M)XcXa}yh)wUyi1ffWprb*v%%w#mTkkxP+s|5&zR
zq&;Iy`r<~O3HXV_Z9KVqeOPUxM#JZKfK_8pq6*stQZ#5S@=xi6_4t;!aW5Bg>}WJ-
zYjh-*7Bfj<r5eo9<OzZ_n-Wv+EqGY`QQ~aV607LTV&$RvjPWy!Rs6NsS9}-2`GQ5L
zaOsj<BbjAzXnO)QaP-G6t(8c(P(AFv=pi0@&jM{KRmE4k6{4!=Mo4$YV;o}w_-Vcy
zj=mqy^eN02%=Q=#L!*W0h>oFPWvDfB_V*!)=>vp!DngOQ$aMUALOC;TMFnbzS^$|-
zjvF_-$@a_B6s+yj30ky>B6m7WLFEaZ(GindsHII2<Feloq9$jdnnC`|@Qn}g<sJ=$
z1@0!us*e^FL|36<S$EKt3$M`3EwhowO)YSQ<{&Kdqba>&7>xi&f#H#1jNw2z$lNoH
z+#3^59?rAJQ4dsUZ=*Q$F^-3aooYy~v>b<N;8x-=vNO6NvzX*BJB)ifO=8s2`qCSh
zEU~h?IhlV#gB`3{53LZ~Bw4l>82cg<>LF7`bV`cR#X@;fwZAoT%Ug&lQV+99vPYO(
zi6@{-Vl4XD#g{nmS&crdP(ycVCk%0uahq4+baR+B(GD)5RZ}7ufq4n>=sge3IUmE^
zzH%1vV*$F3_9Hh!8{(z6&RV><IFOhqD`M61F4(oq4()YSphr!W(FcX|WWKpQdRdc>
zQlDBfC0U+qy6t{!;7iE$1-@*yMmlD)M&Rkw!my@;DRjT_gl)bj6t{eShJ7HX4<)-W
zW8?pxRC(<pd(?e!YmesCV~;7+Tv;!ud_0FW+FgYGCT+m3y4!H%>s5Hpnn=9j@CDqZ
z^;RruBe^wYMRBW^xNcp!pw?MOe;*FDP~8yWRY!j}4%HI8>ZrKfd>aSRSdgoqqPoh>
z^+or6Js$dUKFp1OIQr+uKYm!0ACJeq5h(WYcoCfAr=Z{3ZzH`<^rLyRL0_>t#y^Et
zbwd2p2uRPHEhPd7CF_iwJlwo_hFie6DV|e!BI-Zemp5G!Vdgj=Pq)cD-pp|`yrz5l
zdrkH6_!Pp~g~hd%OYOJnvv?-R#pmkpA3ygZN>YEh86%hHzhU02$m2`IU1`tB{o9n&
z&d=_*!#OU8dq_Wy^V?L6e(gCreSSNf;~GdDiTydwZ&R*4**o;N!~ZVskNOF2zRXds
zoT>2zzpNvgD)US~N?!Q>_Ilz*NjWp-imXJXh`&p<Kb7u3hsSZrIHwK|(Swqcd1onp
zV0Czi@!0F|Z@@z=W8Lo`L*UzEdT}~+;T#U4J#gF{N7Uzv<|~|Ea+Zwn`0*-HTUtkd
z9}ac4o8M=<;c%)gT6OiY7KA^5gE)5G?~f4>*;TprOmqfTSF7{jt>2xwzQ_T4Iw?`r
zPVMod!7?;EZ8-{jxsS$}eWX2l%L(~IweW-;(`nvwU2uB54~kVSm?!D4VAH0C6m7Ty
z#i^Zy*S9+oeqTz*$u<(6Yn4ZQpL9ky&-6oar;gD6&9$g{<WjQlLmC~q)C;Q1cGK?7
zG9f&98H11X1*Pu(uy^cmx=Fc)<Qk(d_9ctX<l^Q7|J-=_#sD#zKl}JlWNG^Pd{ycF
z%|Eg;tsEiNmYZjaeRaS8D+oy2nfMaL4P3rNfU`6G{O|o6N_IAImKaWte}kP#Wc3hV
z5gRzUf17gJT^cD4|HarbO6=#45mT;She$oo-%j=4&;6GDohOCMkJ?2SS1}g+e?KI3
z8vF?aq}xRUNg9@tdc`HloexlZ9>6CScRr}}c>tW<mAmg-^jqauD9QK6IPsl7D!+Tj
zi{o+i|6BRh@&7v!kS@Pmzru6<QeUoL|7+jplmF#yo;XRl<@zV-a?JH(N}gi5xOm1A
z|DcItJ}1B5tS{rmN&E+Ai}`#p;q93JWjx+mvHz5WFNyZklCS4G;+o9i!>y0}@pZVy
zj$)-b<T(i9ILA*(rQhePasR%UT)UU(_WgMOYoTA4D{;;G>)z*MB;#-H32+Fn<QzX0
z{Z@Y~-TwTjzrExoR_pin^FI{`b;A6=5s>b0EhQ=XN%}7?1FpaQ`S)w^|Hb{pulqGD
zeK_QQn{vG%f6{M<b6jrxVO+bH=6e^{>J>^ocj<c!B#c2*q(i>+bM1+%`1t=X1OG<f
zkNOGDF3VAFJo3l)YZyz)+0SlomHr;km;EM(14sWn9<GDLkF^dD5e{-PD+kx%Av#Mz
zT;AklghT4^_#6+vss9E%#4^^c@BUsKKHrZa`nHW|+aI@n*+j}O>cc+|SCjIK_N#v$
zF8Tz_=YG*QdlYN`x%5?~{GxB@i1vesuFvIaS?m8i9bGBEXs7z;>59I6^SQr;RD97l
zc|K21qt-8a*^iS;^bMZR{h}7(^YpZ8{dHFz0(A(~Ay9|F&xOFZ>`c=3>mS*fEQ7>)
zar4%nOUpWG{%!=M?M!@0hL<GkA6)*No$2R4U#i+kEVFcpj^kAMDXH{%8gcW!FD7Rf
zGM@AOcpM55bH5+&za{i*JCoQlCVjqCJBh;5^jAB}A3a~nFIb%9O3^U;vn%j3_doCa
z8V~WjpIe{#>-6VI-T;%XXA%Fr(%;GPTm7c=daiW4C`PSzN43;?FV`;qRqNf~l27`)
z*Va1=NHaerc=jICHF+AiV%iBX=^Q{z%D2JvBM(VCCwKJz{VQ}KOcgfNyMStRyy2Ml
z5NK4d6FvLtMD8&Eh0NIp24He>AbReziWtq^KuWb6fU53#Xl}Lx(zeJ_mupMOMty~x
z14|zuhaCZ|TJOFnW~ByRmpBqus`#*F7X)OpiXU?~IfOOz7r^d*t>A9vJu=dLcWx3a
zpgB9nk<BhgVcLk9<WBxps=X|~_8b>^n0QCPz8trK%HEC5$@A=xqg8D~s&-Vs;d!B`
z#RGRV$S98t-Pn{^uJ%Mq6IPLKt2dx@FAw4I3w(jLeIb706GmFE-OnWUx=5d#xry$|
z1fr!)8j%NCZcui=2%o9yO<yFsF)^1lz_zNMFnBG7(WBeIPHQu48`(vW?%15+t5jOl
z>)wVPueA?)&bCB78|#uhjd&93lZkrpZV|6Vy#xip>UgrcGAvotF?Xj=6NnvX!_wAG
zn2r$&f;|_TqM+f8v8B5lJ@P09FV_qrRgs5DuQCI&>}hwneDy9ejp|9vqc4*&p#kio
zw3c)Rbb&c-D~MfD6DZ1yrPW8TTXa6T37IM_L-$jf2(P_Ygy6-y;g02bbY<#CBH#HA
z>*P_7U9dm_Iv>rY0}N#FOkXX?(hDGa2Ds8o(Xw#8vmv>bS`Yk=Z)N2|PBG7Fze~_<
zq%9gfHV<CSoCn5tl<DgRLy@yh5(ziIM_-w0K%A^1Se@zyE!UhP@`FyYcb2$7W5a0n
z_~{-fHDey=PnCn4x*OOdZ`7%VSp%AO`2wjHXb4;<$im63hhgERDi-^tSa`<D;2t$m
z@OIl9a-vx`96sI)jR-#u-KJkMSNFDqOD8RmTyOvioVx?J$uUB=1`B#Q<ml4zb^{<#
zYcxtVt46XBa@aUhQINg%3F%^=g+9!RCQa-p^Rn0r>%TjM)E4UtcH|i5_U-nT=HH7z
z6I;u`#NBxoP5OpVlv7L&_1C9YvfmO1<pp3LP>q7rZsh8e>f;on7&tKMJ$jHKAXFeg
zL2pOXtSTXz*||R2@A86#L~bIAPn;mVb0oUARS^uXH=!?Mwdt0vd+ExdPuQUv!9;7$
zDzMlLC|++O6#Jb+0nz42ev<-Oa5MxC87zYGlw#6WPY+c0=P_tYeKd3bdi=ceJh<NY
z7^Bd0HF1gR0V8i6gVIS)Nk@y_sLEXjw6{E?lhAAwK6(p%sSu7;GWBUda0lo<M;Wzg
zYJxi!<s-fHaGL$GlzFvN7nj~Y3cIpGK;7!S#o6*Ur0kX)be+?k=5HU24)2(W6q7EH
zq<gxIK05~0H+G<Ri)4fo=3PNaO-I+B;qHxn9`d0<{tC#fYAdv3v|wD{FnWwMM6bM`
z!whnbB$d~LgI(I-aP6-2$vsuHX?;D~TfHOXPJ2pQy6z%f9(M+h!4<4#(?yJL`V8n7
zvn2O^+Cl;kw$V4QEU8TPJDL{n2bN<sQO~;(wD8pqvS$*D6kk=5<g31{zb+wjZ+?XA
z`iCIe%n{1nl3_;RDsuUVHN@1O|I;neTsY6T4}&(nqQ(PTp!SxHsLa5YAa}Gcd*N9k
zI>GBfTOM0RVwiWppM8SZoT>drK|vGd8N6be&q*K~*BnG_<YKbh#+KBKsY2Jr9-w;-
z3qTO2CY<<M76$FjClBmO=-7z`Xw@xU$hJ~KOh{X}*EpGk25p6-v)a*A?*?GC`3_7O
zm&NG$nLxeONl-K*QrIrwEx9<K&?nV>p-Q_a%^ius_USeft~#BRzFtS&HU^NZz006#
z(>j!4@|w0;QB-@rk~Q2*t6&;$*h(_o8p5hFC48!F2)a7>0$mw<0g#ivaF7OuHXZJw
zW`nb-;-ihkDK!awI0R_5S1QST$Uv0rL%3REOE<heiW=8`$Ha3;9AQIj(Xsg^bZ$^T
z`hk~)6g-B3L$3%_q#r<!7eygjG#JI@okG`3cCs@Ln&E)Sd8EV56vD3YrWUPIP{V$?
z=zWJy=5Y~a?3hbyuz65Zysfn(N*I0sdgm=AR*y$9&PuCDYM}*g_w*dNer$&4y4)dK
z2JAtjD&L_I%_h-;fq8hC<#t?X(}lQ~Hw0_-;f#)RAu5Y%M&94=B%Ew=lq}rpK`$|{
znRJ(x(A6)Rs4I7X-3z-SKh@fA4QxBg&UXmUeYh)+99poQUN%~Sbc1uq;4NLrnieha
zhx$g)HbxCvKFTKXnX4h|&=xRgJdD2Vw=GxS$uU=VR!8CVMlz5pJV5yEj*&;sC1m1r
zcWB$`APK5?1`}_zhG)%Bf^Fy##-^%03>kY91@}&(3uo;@nv1L8WK=uAT`%GWm5s@<
zW*za}r{Ro+W?Nted|*O$G{={3Fp!_12V0by;j^eYs_%b-Rh@r>8Sk(hd-TdiyvYj@
z9<rXkd+kNd-*`ZeyE&wXt`AwQ(@3z^qaN06coD67wivYuGlMqk+EI(8TS#!ny=0E+
zEaLPa7Q2nHB8H8E5Hnr|Qfx2LQuTZ!%MJsZD|d-SQGfD8;7CRX*)!(*<}g+-^k{>6
zZ%Ln9A*5^L>F_XLlP<XwigqT35bR+?_EwA|G*=I_9tz1TlWye5$_z-Kc?CuZEoqF7
zZ|?K)gV1@mVsIN*Nt!sXLJ4PIqI>aUScS1hu;cP=qN;ick{;)fr>i><!<<%R*4v0&
zb~~STZ=gqQt?rVi5A(3U+a2;?u@60&c#9O}-e*4^)Wz!-w#0Ga6Ul|3p7>7vD&nj<
z4%;-RsG_(PF*Q?SpLbYHMmg_-0@YE-bw>aidpQP@r;HKguIxh<Gy1b~gHMw9=a^ny
z+!`M&Fe4_Tve<^5#*hY?8%VFDd{k_8nzb;Ui~^J6XccTidNG5+d3ZVNV*Q>Lp2BF=
zM0;?xnvHtptRux^L(t%kSJ)xRap<n%Y??T{9U3I#2<`b-39B-Uo-QjPi(*^Clcj~k
z+Ta|atEYic;3Bre#*1i4i>**OJOM5(G@uSGLZD3NnMI+DC0)9)5hSSqd)9@ZnqIYM
zSozsOa7A18>{CUEAM1!xl7_+6=W1xug8tA-W)kXK^^`7p;Dg#7&I0>G+j8G69}o4$
zC?kKH0<vU{4%kHArG{tnk&#Y6NV8u;js~`%lW%QAgIQCwJ!TiP`^--AZd)HXb7cp+
z`|e?KOQ8guO*X;Swr%L<yr)Qrby1wme!A_f9z37B9X(yJm<*Y+9Ey5wLN)8O;B|cm
zC~(+~5?Y0TQ%ilg{cJP)KJhg<G+tkLe`z=9uX7M$i|&(i(ZR6wcu#nCZZul(@gsAr
zdJnzd?if11U>$K?m(RXFGm%bJ4`VmZyN4X#tU=UG24?mvL1iZ<veS+wppMq>q1o*e
zIOUW;2i@C)RvBo+<BkjIn0EKbpds7Ih|IoZcBge@<M9I|DQ_cb9O%mS`Pc`+KwpM=
zHkHvB_mF9m+JNpKnjl<t@(_Eqas{2UqCcqIwx>~E7s(dCyqtcsG>PK$HS~nz22^@_
zARW2cfh1mRE4&hB0j4MW!Jds>AzUqstnB#=DtERdBV${U1#fMI`x3_pj~vMs-d}tH
z)l0igblkR4_ZE?;!L8k-q<A^%nPyK$vn8ZjaV+H*4x%|3$!PwYz0mneB5Hc=80xyH
zonThlW@6*KizGy=(n|lm%<<T>c+plhunD$Ct?#Hn#lEMgSXUNn>1hgfoM{g4rz#2#
z>*+Ag$`#DVW%E$!lXqyL+hW1ATYBv7ZrxFv@gcN@&PWKzn@xi}98r~tEi+Ze3PR`J
zM9O7J*%==?!|XQhOw~bMI`ha1Vis1Pu5VKx-soI{2j>njd!m8lwKIgU&?dt6&L_~O
zpcQnMrxv&!(x$Hi8(^1|on#kWAmxF}V4iUstP^z>&NP0?x}9XH$DwxMH?%oR!gWxQ
zV_(vPX0Xwo9Uu(dL2<ilzxh-VK!en-&~C?;Wb+4B5aXDWkjc(tc|l6hC+Q+Nt$!^?
zW~n~$?r|Jm-rfya%J-NHJ@%0N1{az5_ll_0&y22iV1&F^tI)Kh%d~Lw2xx8-O@$qY
zBL4O@r0M9@%$plMX{d=AGvJOHt;tI#MLoCBgJT4cxN|4j?O)HL_gH5XUVF}5g=;J-
zexpH}7$nn?w;YA0=|MvKa|4J?kQGGecBYrS%84Cs0g2h>#mHA>lGQE`;LL;iw0QPy
zur%L{oM(j7O18Vj>am@G&r^b~^YqxU*C~>-*T5Y;HlczZ=OE3vkiEE17dL#qkQAM)
zLDm<C5*MC2Jz}OqESLB~!rRlt?MOZ;+ic8?O;&>yQ}zh7hvzW!HY_I9w+9F!+9`lc
zd<-E+){-gf&ymVG2J~@sHq$+(gj6yY$cE;r7GAx5Y4gMk91~~*`ne3WRoh9_2KImq
z*BfY0mz|i@N7%di1A5vno#|qS;ZlMQN!LqZ7PpeeXCB)eRH^-bkXpOtDDwPGqPSAI
z_FGhkks9p?sq?DQv35!%tA{%5pRGyt+nvpoHEjltLxGwl_7SM{-<iGMZ7~TbTZ(KO
z^rk9ZGA*v2d1&#-R2gpK-hc~IaoZaWAtPcP)q1-%x61w?Dop83G)})}D??)G_^h?C
zZ}0&&qB;!Eeigzz*`tPcHLW4uV=~A>wHG<r34m;#En^SNTLr#-S23mL%c;TQT%?lS
zh>UyC9HxgxBf}0V^zfZviy`fYll0lXY@3Sh5czPkU=&)xwAP=@G*eoNr!3I`q4wq6
zi2m!*hox<C!hxoOHoeQ)=IUMod$x@2wPqr?PwWEz3Z9U6<t1DgnnFL!jYD&fb*0Z$
zT9KW58`4+NO0{#7CuqYA4czE*9Q2%!P0TWdFvejOX*O^i_1H9;7<Si%_DxTqm|1NZ
zhg(hIgjEL;I$#%xjw_~1bXL+1AJl~7*PKEE-3cgba5!9-?ame*oQt~be?>Z7jn2K$
zv<GOYoFcu#tHInXnY?z127WKXv_EYGKI6uN673*-t8$Kvnl4Wr%Y2b<Kp5m{+(9!~
z6QVFRhcsH2L1AJt8Id0hZ}yCa>AJ2olV6tGu?C@oS8gI+unKm_Ur$b1_@c;yd?Ku3
zfzCJunhw2czqh1Ao~unYf45juXc`_zW;EJ{dc9IYw;Pu;_l+EgO5=o_Qs*@C{A>lY
zbM852j7bU^+ISxt*gutipz-+G^&Z$MY#(0Y{Fpej7({o+uSbcIqtMpS?&L(^HXOKA
z79S`Xh*0E|+-L(2x?@BlnYN`r(hi-;`YhT-UMDYP9AC<k*_AD^Pq~pmedZaO8NU+x
z1zK9%-kyS7c<wl3v@C?vwZ!evs@&yEvQX7ETioezQ|zCqFF5rm0XbaBC$S2$5X`oO
zh1W~ra;&Ow)T&~%vmlYezJ=sT$TXT2YRMkFo6nrKcSa!*>h$UkE!20}6L5SumAoI8
zL1X2PksB&Um}MT>?3mEC<hbWRGQZ*k8|rw7?R}$&+Beuq74`YTk?zBYOl&<S_;xR<
zy{Qx961j^#9y|=S3l7P3af*T&g$}Tw<8#=P+YA!t)#9ww7B<!7laZ)36kj(LHr%a<
zKa>=~0~~{<?ahOz`PLR`HrgoAs0u~sZKn%b#-nD}EJ(<vrOdcJJo<Qk1oFL5K}I#x
z2QP)DSWYdEY;*O_DcF+8hJ>V%8N;mLh1YuO?fZgcS{0zfcV&cked^=6IonB3`yAwa
zC=~WTRzZpnJyB=RiD;QcV|Wy%2rHdy*sW|II8k*C#<$WJ?u1k9>BJJqY@|#yb&taY
z^W7HnmM<f{w4<1`u|p_Ze;^rU+X5CY41i7hS5fO3F0d_iCwb`^PiF5Mi4@LFfs-}y
z=y9<kEDR`xzJ0BPvo1wZD-BIhNNGcYSz9_k@ggd}D33I|TY*kkB<VSAIc(@+3lAqA
z1aMzVZk^cxw~wo$!>dckxsGb+qyHip7p8!f7pdg-SzbwQPIO{z7VTq0OvBj*1<ToM
zh1IP2oR-*bU=izEzSjInXC>OSeFf9e^9d0Q+Ji=BRFTn}*AlhhbTs5pOB~$!8YIT9
z0`C?LXr=2t^xQDs!mHCFXwa*a>~t+g%Ia3QykQ%VZxoJB-5ZEbO$lO3a<?<uIv2>X
zzU}Zjn#nvGP(fTB9R%K8l+cRjS>PAZ1m<`+P`B{4#AwS}w7y3SiD?iD!|j$JJv1I~
zw_gm)PpHx=yJ*zuL=4J()DYag-ZR(aHe!!^r$~aeD(-GyjLX;aC{kEUGG$GOrm!9}
z`Sl9aep?<%O>#t0gKBbBN_>gM*>G%py%($;WGIN!&qA}BFgQuE2_Cvyo?QHp4Jj_0
zXrEOdknP8fNdMA(6r;TzI$IATTI;T$quQzD73+yETvfo{{A@5>=?!D>1iCXIgnZ03
zhr;H2$%pKIXve{kD1S&OQl4}Wt?s-5f?ajUApfOwk^Br+=IU!SU_&38u+D=m^X5am
zdl_-cK1!xe?2ipzKVx%(7t+Zxk66Lg)i82X8tE8ZADNxr2HJjk^le-jNeYw|A{RBN
zKED|*T3g`ocdHm%%dJFZUN=1C=20?z`c)X-_B~lE_XvFZR}h($W2i7|3R|>sDVx<`
zG_>YlBqt|SF?Vj*2^43Hq$|6*(gS=0XgDN<RJR)pq`J)f@$p7DR=X;<W!KhF6j24d
zNB!_}G6}nCOhL~sZ9xZKmge@z9F4~LbceLkM>MlbQ@FJ6E~u<Oik1%=1LKV2bB|V+
z=bBwRMq0leLwz!enPG-4p<j>NtW3yF7;yb0+2*~S7=^KDNK!6Z{8Ag*We8wl_)Df`
zi!Ct2-<ay(UqiF2x<OU-T=Z!DLfYzW4qBJI8{IGSL7{f5(d6>Zpb(P+5#>3cP9xFH
zf=;Azdq>dmQ%6_lY^G()1M!`(2gqVWf7q2?WzpwSZT*hs;o~J~sL)al$1=Hiz4b&y
z8oq{sOFhUd#fHLWu1gW`!)|7;P?OATUqR}Z#=`w}8rb>7X5{^%3mzF2EGU>YOt1&a
z$)fR!c)n9}L7??%6x6{HDNU*&@B7%3=$!3nVfRf~aC5Q1ve|;#kArp=9G=(}hL0AY
zsE<ovXsDbpII<~|xY?R!DI}t;8*afZwko?(UL1*ax??dXZW66my@J)Avx3-%*h6r$
zNE&gmirI0@kWn^##O_j2!@2h-gGWyuOmMCtkM0@L%I$|?)3H7D)NlhbY<DCoJTjYF
z3l;IHIaTOc(*}^({}ypJ8w+Hp0yN6I2D8_dkl>6+W}|0+_|ZDE@j9_lo?Ey2j)}f5
z=)Vwwuh*f(CE@@7{#l%4-HBW2<iL|&hvN7@2bF&REO7nz<>KlM|L0cg`%3%`p<kCP
z@k$-HzR8s@j{foQ$?cGoFYXC&O047@KNbBN4}(wgaqHjQ#~x%QQcL@}_3uCa9%!PU
zILQ!6${Y#g9RCseqkhRUQ7o49)BTSKTbJoSj=;Cp`7dqyQk1`bonQL>A8`r%wVnF2
z-`C?$6HET@?hhk2e~G`ekE^f73DP0|+y1ZH|G&CFG?tVv?g?;s|2z9b(fb#oi|Y@i
W{eQkcG?U<z_$mEgfBv)m;r{_IUiX*)

literal 0
HcmV?d00001

diff --git a/tests/test_from_nir.py b/tests/test_from_nir.py
index 4ae9a33..d16bc68 100644
--- a/tests/test_from_nir.py
+++ b/tests/test_from_nir.py
@@ -22,6 +22,33 @@ def _torch_model_map(m: nir.NIRNode, device: str = "cpu") -> torch.nn.Module:
         raise NotImplementedError(f"Unsupported module {m}")
 
 
+def _recurrent_model_map(m: nir.NIRNode, device: str = "cpu") -> torch.nn.Module:
+    class MyCubaLIF(torch.nn.Module):
+        def __init__(self, lif, lin):
+            super().__init__()
+            self.lif = lif
+            self.lin = lin
+
+        def forward(self, x, state=None):
+            if state is None:
+                state = torch.zeros_like(x)
+            z = self.lif(x + state)
+            return self.lin(z), z
+
+    try:
+        return _torch_model_map(m, device)
+    except NotImplementedError:
+        if isinstance(m, nir.CubaLIF):
+            return torch.nn.Identity()
+        elif isinstance(m, nir.NIRGraph):
+            return MyCubaLIF(
+                _recurrent_model_map(m.nodes["lif"], device),
+                _recurrent_model_map(m.nodes["lin"], device),
+            )
+        else:
+            raise NotImplementedError(f"Unsupported module {m}")
+
+
 def test_extract_empty():
     g = nir.NIRGraph({}, [])
     with pytest.raises(ValueError):
@@ -103,3 +130,9 @@ def test_execute_recurrent():
     assert torch.allclose(y1[0], y2[0])
     out, s = m(*m(data))
     assert torch.allclose(out, torch.tensor(2.0))
+
+
+def test_import_braille():
+    g = nir.read("tests/braille.nir")
+    m = load(g, _recurrent_model_map)
+    assert m(torch.empty(1, 12))[0].shape == (1, 7)

From 5cc87dee34b48519c182a9a0beb436ad1df8e2ed Mon Sep 17 00:00:00 2001
From: Steve Abreu <s.abreu@rug.nl>
Date: Fri, 13 Oct 2023 17:23:32 +0200
Subject: [PATCH 06/27] test for NIR -> NIRTorch -> NIR

---
 tests/lif_norse.nir         | Bin 0 -> 17584 bytes
 tests/test_bidirectional.py |  92 ++++++++++++++++++++++++++++++++++++
 2 files changed, 92 insertions(+)
 create mode 100644 tests/lif_norse.nir
 create mode 100644 tests/test_bidirectional.py

diff --git a/tests/lif_norse.nir b/tests/lif_norse.nir
new file mode 100644
index 0000000000000000000000000000000000000000..ce5f8e1c230f01b93edf3678479ebcfebfbfaa69
GIT binary patch
literal 17584
zcmeHO&2JM&6o2awYSMztDGd^0q+Z;Lf}yB&TM!8(;7Accw1?ahqZn)vB~jv}?g@_O
zmSc`N_SiomM-Cuy<C<fS{SRQ?zV}<Zv$nTJ`B*zMmi*@9?aX_>eQ##=?d<E)^n>$b
zm&S-wHcMwH7iRJIQPbbkQiC$umoyV7Z7TgMBbk8sNA$JUU)26xzTcnB%#h_o|Dbvq
z70yU5iwaLq6#}K1smHAFf@%$WdZTs*Ft#fjo3*u7T3uVNa2@P&5)6IPfJUVc?I<9z
zK=eCeEt~TbK@eItdyljU->Bv1awa3yS@nbwbmaIO(Qcm-9xwZEkLITX%0yB(MrC~9
zv8Uyn-XQp^8981#(>_|adC|o$X<po&5KVz|+V%GG9WDah?C9W<j&r&9c<in@H6Hfi
zhxg~_LSNQx0Q_3rnhMu$T)ScUFzv$b9?d>^xKUoO!p*@yySlzrxBR+%(^;G6H1+IR
zZM8ync9n$d=4Ky+72pG%->QcleWCo2i)G;iIo$el(j1Murv1o<2Xi%1LHWK^ihfZ4
zmDWSvn(C$1^!@^Jo_Yj40v-X6fJeY1;1Tc$cmzBG9s!SlM<D$OxN%maA{P*F)pGsM
zb^G*AVVrf-(#rDl3R65c7rC*IbaZjnoDnfJpb^P52ys@9E=$_G9|_(u1MT)TK!W?=
z#(@wgX`<tfy$w%gQv9@;7hTQpsO&ZA0|9rwQ*$8W#7z<JLp%_1N5m81Z)i^vp<CQ`
zPsh!vAW8On@SThVzfKz0w+|Nr-R!9Nz391cU1&GH(k*GP{wS2~2=4<p#?SF?lsQHZ
zJLkAO>Ls@S>aMLK=0QDQJNDt5Ux}pc-P^Z!+OOk0`fDGDW9?jzemExIUgp?+dzqKK
z_A=|xJ`Uai*lM3V0v-X6fJeY1;1TG6fV&^YV`4$~qqqdm|C`lv7zp6?m(yhU-~XSS
zD{7rS+`l4ypkH@C1G)Zj_&V)pxTl8htI?mk&o(j~HSS0z;?weexVurN6qBQiqZlJ<
zBhmm6c9BUXb<eBiul5!b&r$Ec?ag=ETH|&l9vIkjkDTBypd<E#Am@1xFPKg;n|+?Z
zdB{x2+By$eFd;*9l1EBJJIJJpCO(&IAi(vOR4L(wIWh#?b=2rc4SgQRz*tBBsMMZU
z?LoPxwK9j%BY!G_@qf~|pNjXj{PJ_pD;JH3;9B$`%>S)8dhv(eKt}VwbD!L+fI|MJ
z+cKQxxLxH#MFW|}BYTCNxfJK*FPbZQ7Szyqmd0NyT;11&k`Wf1Ak_<~fy{}}I37vs
zq4-_@`p+3r(AB?rB6@clG-*Bj_b0hhtfpHpzJ6E{aMweuYfX(iLvO5y_3~C&wV@5$
ziz}7#3!?4CdUd0+SzTLM4qGtK80B+G+)rhO$90&O`u~3Cu?~rq03#2=dN>-rS1$bm
zxjR#6SNyqiJ8A`0q66oCq8-#|X2HG{=#?-}RB!j<I=QyET^-f@zhp$qbkG0ChsUzV
z!x}{Exw-1`_zxZ{1A+&?k&S;Bz30RmllE1mDcgxnn17|@|4$17?!HP^A-fJc@oMa=
z4l7<CpYIQAU$t)J-eKo~J2?pZs@8YWJNF=WX9|BUPIhiTt+jExO4?WT8)trH1Y6tz
oAiaZ3s%RXKd?+8{F8FR1-{Z#1pBwLM8UbQSNc7)`awn1b1ebp;X8-^I

literal 0
HcmV?d00001

diff --git a/tests/test_bidirectional.py b/tests/test_bidirectional.py
new file mode 100644
index 0000000..9018d81
--- /dev/null
+++ b/tests/test_bidirectional.py
@@ -0,0 +1,92 @@
+import nir
+import numpy as np
+import torch
+import nirtorch
+
+
+use_snntorch = False
+# use_snntorch = True
+
+
+if use_snntorch:
+    import snntorch as snn
+
+
+def _nir_to_snntorch_module(node: nir.NIRNode) -> torch.nn.Module:
+    if isinstance(node, (nir.Linear, nir.Affine)):
+        return torch.nn.Linear(*node.weight.shape)
+
+    elif isinstance(node, (nir.LIF, nir.CubaLIF)):
+        return snn.Leaky(0.9, init_hidden=True)
+
+    else:
+        return None
+
+
+def _nir_to_pytorch_module(node: nir.NIRNode) -> torch.nn.Module:
+    if isinstance(node, (nir.Linear, nir.Affine)):
+        return torch.nn.Linear(*node.weight.shape)
+
+    elif isinstance(node, (nir.LIF, nir.CubaLIF)):
+        return torch.nn.Identity()
+
+    else:
+        return None
+
+
+if use_snntorch:
+    _nir_to_torch_module = _nir_to_snntorch_module
+else:
+    _nir_to_torch_module = _nir_to_pytorch_module
+
+
+def _create_torch_model() -> torch.nn.Module:
+    if use_snntorch:
+        return torch.nn.Sequential(torch.nn.Linear(1, 1), snn.Leaky(0.9, init_hidden=True))
+    else:
+        return torch.nn.Sequential(torch.nn.Linear(1, 1), torch.nn.Identity())
+
+
+def _torch_to_nir(module: torch.nn.Module) -> nir.NIRNode:
+    if isinstance(module, torch.nn.Linear):
+        return nir.Linear(np.array(module.weight.data))
+
+    else:
+        return None
+
+
+def _lif_nir_graph(from_file=True):
+    if from_file:
+        return nir.read('tests/lif_norse.nir')
+    else:
+        return nir.NIRGraph(
+            nodes={
+                '0': nir.Affine(weight=np.array([[1.]]), bias=np.array([0.])),
+                '1': nir.LIF(
+                    tau=np.array([0.1]),
+                    r=np.array([1.]),
+                    v_leak=np.array([0.]),
+                    v_threshold=np.array([0.1])
+                ),
+                'input': nir.Input(input_type={'input': np.array([1])}),
+                'output': nir.Output(output_type={'output': np.array([1])})
+            },
+            edges=[
+                ('input', '0'), ('0', '1'), ('1', 'output')
+            ]
+        )
+
+
+def test_nir_to_torch_to_nir(from_file=True):
+    graph = _lif_nir_graph(from_file=from_file)
+    assert graph is not None
+    module = nirtorch.load(graph, _nir_to_torch_module)
+    assert module is not None
+    graph2 = nirtorch.extract_nir_graph(module, _torch_to_nir, torch.zeros(1, 1))
+    print('original NIR edges', graph.edges)
+    print('converted NIR edges', graph2.edges)
+    assert graph2 is not None
+
+
+# if __name__ == '__main__':
+#     test_nir_to_torch_to_nir(from_file=False)

From 70f44474f1874af4d9bb937c0d101c88d913a14d Mon Sep 17 00:00:00 2001
From: Steve Abreu <s.abreu@rug.nl>
Date: Fri, 13 Oct 2023 19:03:03 +0200
Subject: [PATCH 07/27] refactoring + expose ignore_submodules_of

---
 nirtorch/graph.py  | 18 ++++++++++--------
 nirtorch/to_nir.py |  8 ++++++++
 2 files changed, 18 insertions(+), 8 deletions(-)

diff --git a/nirtorch/graph.py b/nirtorch/graph.py
index aae3e61..1d3f5d3 100644
--- a/nirtorch/graph.py
+++ b/nirtorch/graph.py
@@ -194,11 +194,16 @@ def populate_from(self, other_graph: "Graph"):
     def __str__(self) -> str:
         return self.to_md()
 
+    def debug_str(self) -> str:
+        debug_str = ""
+        for node in self.node_list:
+            debug_str += f"{node.name} ({node.elem.__class__.__name__})\n"
+            for outgoing, shape in node.outgoing_nodes.items():
+                debug_str += f"\t-> {outgoing.name} ({outgoing.elem.__class__.__name__})\n"
+        return debug_str.strip()
+
     def to_md(self) -> str:
-        mermaid_md = """
-```mermaid
-graph TD;
-"""
+        mermaid_md = """```mermaid\ngraph TD;\n"""
         for node in self.node_list:
             if node.outgoing_nodes:
                 for outgoing, _ in node.outgoing_nodes.items():
@@ -206,10 +211,7 @@ def to_md(self) -> str:
             else:
                 mermaid_md += f"{node.name};\n"
 
-        end = """
-```
-"""
-        return mermaid_md + end
+        return mermaid_md + "\n```\n"
 
     def leaf_only(self) -> "Graph":
         leaf_modules = self.get_leaf_modules()
diff --git a/nirtorch/to_nir.py b/nirtorch/to_nir.py
index c12511d..f680d19 100644
--- a/nirtorch/to_nir.py
+++ b/nirtorch/to_nir.py
@@ -12,9 +12,14 @@ def extract_nir_graph(
     model_map: Callable[[nn.Module], nir.NIRNode],
     sample_data: Any,
     model_name: Optional[str] = "model",
+    ignore_submodules_of=None,
 ) -> nir.NIRNode:
     """Given a `model`, generate an NIR representation using the specified `model_map`.
 
+    Assumptions and known issues:
+        - Cannot deal with layers like torch.nn.Identity(), since the input tensor and output
+            tensor will be the same object, and therefore lead to cyclic connections.
+
     Args:
         model (nn.Module): The model of interest
         model_map (Callable[[nn.Module], nir.NIRNode]): A method that converts a given
@@ -36,6 +41,9 @@ def extract_nir_graph(
         model, sample_data=sample_data, model_name=model_name
     ).ignore_tensors()
 
+    if ignore_submodules_of is not None:
+        torch_graph = torch_graph.ignore_submodules_of(ignore_submodules_of)
+
     # Get the root node
     root_nodes = torch_graph.get_root()
     if len(root_nodes) != 1:

From 175d34f8109ea1a221bfc66c67d6cdaf7be06c76 Mon Sep 17 00:00:00 2001
From: Steve Abreu <s.abreu@rug.nl>
Date: Fri, 13 Oct 2023 19:04:18 +0200
Subject: [PATCH 08/27] fix and test for issue #16

---
 nirtorch/from_nir.py        | 12 +++++++++---
 tests/test_bidirectional.py |  5 ++---
 2 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/nirtorch/from_nir.py b/nirtorch/from_nir.py
index 349a1d5..5870c5e 100644
--- a/nirtorch/from_nir.py
+++ b/nirtorch/from_nir.py
@@ -122,9 +122,15 @@ def _apply_module(
             elif input_node.name in new_state.cache:
                 summed_inputs.append(new_state.cache[input_node.name])
 
-        inputs.insert(
-            0, torch.stack(summed_inputs).sum(0)
-        )  # Insert input, sum if multiple
+        if len(input_nodes) == 0 and data is not None:
+            # single input, no need to sum (fix issue #16)
+            inputs.insert(0, data)
+        elif len(input_nodes) == 1:
+            # single input, no need to sum (fix issue #16)
+            inputs.insert(0, summed_inputs[0])
+        else:
+            # multiple inputs, sum them
+            inputs.insert(0, torch.stack(summed_inputs).sum(0))
 
         out = node.elem(*inputs)
         # If the module is stateful, we know the output is (at least) a tuple
diff --git a/tests/test_bidirectional.py b/tests/test_bidirectional.py
index 9018d81..3a8abce 100644
--- a/tests/test_bidirectional.py
+++ b/tests/test_bidirectional.py
@@ -28,7 +28,7 @@ def _nir_to_pytorch_module(node: nir.NIRNode) -> torch.nn.Module:
         return torch.nn.Linear(*node.weight.shape)
 
     elif isinstance(node, (nir.LIF, nir.CubaLIF)):
-        return torch.nn.Identity()
+        return torch.nn.Linear(1, 1)
 
     else:
         return None
@@ -83,8 +83,7 @@ def test_nir_to_torch_to_nir(from_file=True):
     module = nirtorch.load(graph, _nir_to_torch_module)
     assert module is not None
     graph2 = nirtorch.extract_nir_graph(module, _torch_to_nir, torch.zeros(1, 1))
-    print('original NIR edges', graph.edges)
-    print('converted NIR edges', graph2.edges)
+    assert sorted(graph.edges) == sorted(graph2.edges)
     assert graph2 is not None
 
 

From 37e423732c3998fa6dc765feb90f7b272a9c8282 Mon Sep 17 00:00:00 2001
From: Steve Abreu <s.abreu@rug.nl>
Date: Mon, 16 Oct 2023 07:58:52 +0200
Subject: [PATCH 09/27] fix recurrent test

---
 nirtorch/from_nir.py | 11 ++++-------
 1 file changed, 4 insertions(+), 7 deletions(-)

diff --git a/nirtorch/from_nir.py b/nirtorch/from_nir.py
index 5870c5e..17d227e 100644
--- a/nirtorch/from_nir.py
+++ b/nirtorch/from_nir.py
@@ -122,14 +122,11 @@ def _apply_module(
             elif input_node.name in new_state.cache:
                 summed_inputs.append(new_state.cache[input_node.name])
 
-        if len(input_nodes) == 0 and data is not None:
-            # single input, no need to sum (fix issue #16)
-            inputs.insert(0, data)
-        elif len(input_nodes) == 1:
-            # single input, no need to sum (fix issue #16)
+        if len(summed_inputs) == 0:
+            raise ValueError("No inputs found for node {}".format(node.name))
+        elif len(summed_inputs) == 1:
             inputs.insert(0, summed_inputs[0])
-        else:
-            # multiple inputs, sum them
+        elif len(summed_inputs) > 1:
             inputs.insert(0, torch.stack(summed_inputs).sum(0))
 
         out = node.elem(*inputs)

From c76fb6f0d4d0a567d035b4ede9e6f23849ad1743 Mon Sep 17 00:00:00 2001
From: Sadique Sheik <sadique.sheik@synsense.ai>
Date: Wed, 18 Oct 2023 14:38:42 +0300
Subject: [PATCH 10/27] remove batch froms shape spec

---
 nirtorch/to_nir.py       | 10 +++++++---
 tests/test_conversion.py |  2 +-
 tests/test_from_nir.py   |  3 ++-
 tests/test_graph.py      |  4 ++--
 4 files changed, 12 insertions(+), 7 deletions(-)

diff --git a/nirtorch/to_nir.py b/nirtorch/to_nir.py
index a6ef230..686bfa2 100644
--- a/nirtorch/to_nir.py
+++ b/nirtorch/to_nir.py
@@ -12,7 +12,7 @@ def extract_nir_graph(
     model_map: Callable[[nn.Module], nir.NIRNode],
     sample_data: Any,
     model_name: Optional[str] = "model",
-    ignore_submodules_of=None
+    ignore_submodules_of=None,
 ) -> nir.NIRNode:
     """Given a `model`, generate an NIR representation using the specified `model_map`.
 
@@ -49,7 +49,9 @@ def extract_nir_graph(
 
     # Convert the nodes and get indices
     nir_edges = []
-    nir_nodes = {"input": nir.Input(np.array(sample_data.shape))}
+    nir_nodes = {
+        "input": nir.Input(np.array(sample_data.shape[1:]))
+    }  # Remove the first dimension
 
     # Get all the NIR nodes
     for indx, node in enumerate(torch_graph.node_list):
@@ -85,7 +87,9 @@ def extract_nir_graph(
         if len(node.outgoing_nodes) == 0:
             out_name = "output"
             # Try to find shape of input to the Output node
-            output_node = nir.Output(torch_graph.module_output_types[node.elem])
+            output_node = nir.Output(
+                torch_graph.module_output_types[node.elem][1:]
+            )  # Ignore batch dimension
             nir_nodes[out_name] = output_node
             nir_edges.append((node.name, out_name))
 
diff --git a/tests/test_conversion.py b/tests/test_conversion.py
index 4e2f0db..bcbd527 100644
--- a/tests/test_conversion.py
+++ b/tests/test_conversion.py
@@ -7,7 +7,7 @@
 
 def _torch_convert(module: nn.Module) -> nir.NIRNode:
     if isinstance(module, nn.Conv1d):
-        return nir.Conv1d(module.weight, 1, 1, 1, 1, module.bias)
+        return nir.Conv1d(None, module.weight, 1, 1, 1, 1, module.bias)
     elif isinstance(module, nn.Linear):
         return nir.Affine(module.weight, module.bias)
     else:
diff --git a/tests/test_from_nir.py b/tests/test_from_nir.py
index d6d09ae..55bfc01 100644
--- a/tests/test_from_nir.py
+++ b/tests/test_from_nir.py
@@ -17,7 +17,7 @@ def _torch_model_map(m: nir.NIRNode, device: str = "cpu") -> torch.nn.Module:
         lin.weight.data = torch.nn.Parameter(torch.tensor(m.weight).to(device))
         return lin
     elif isinstance(m, nir.Input) or isinstance(m, nir.Output):
-        return None
+        return torch.nn.Identity()
     else:
         raise NotImplementedError(f"Unsupported module {m}")
 
@@ -49,6 +49,7 @@ def test_extract_lin():
     assert torch.allclose(m(x), y)
 
 
+@pytest.mark.skip("Not yet supported")
 def test_extrac_recurrent():
     w = np.random.randn(1, 1)
     g = nir.NIRGraph(
diff --git a/tests/test_graph.py b/tests/test_graph.py
index c99e1ed..b999721 100644
--- a/tests/test_graph.py
+++ b/tests/test_graph.py
@@ -279,8 +279,8 @@ def test_output_type_when_single_node():
 def test_sequential_flatten():
     d = torch.empty(2, 3, 4)
     g = extract_nir_graph(torch.nn.Flatten(1), lambda x: nir.Flatten(d.shape, 1), d)
-    g.nodes["input"].input_type["input"] == (2, 3, 4)
-    g.nodes["output"].output_type["output"] == (2, 3 * 4)
+    assert tuple(g.nodes["input"].input_type["input"]) == (3, 4)
+    assert tuple(g.nodes["output"].output_type["output"]) == (3 * 4, )
 
 
 @pytest.mark.skip(reason="Not supported yet")

From 48f9842f801bde216b974fe990e7837cb7379b9a Mon Sep 17 00:00:00 2001
From: Steve Abreu <s.abreu@rug.nl>
Date: Wed, 18 Oct 2023 22:57:41 +0200
Subject: [PATCH 11/27] bug from hell

---
 nirtorch/from_nir.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/nirtorch/from_nir.py b/nirtorch/from_nir.py
index 17d227e..5018169 100644
--- a/nirtorch/from_nir.py
+++ b/nirtorch/from_nir.py
@@ -99,8 +99,8 @@ def _apply_module(
         self,
         node: Node,
         input_nodes: List[Node],
-        old_state: GraphExecutorState,
         new_state: GraphExecutorState,
+        old_state: GraphExecutorState,
         data: Optional[torch.Tensor] = None,
     ):
         """Applies a module and keeps track of its state.

From 26242d3da708197599cc8602e40523427ce30cf2 Mon Sep 17 00:00:00 2001
From: Steve Abreu <s.abreu@rug.nl>
Date: Thu, 19 Oct 2023 17:08:06 +0200
Subject: [PATCH 12/27] from_nir hacks for snnTorch

---
 nirtorch/from_nir.py | 12 +++++++++++-
 1 file changed, 11 insertions(+), 1 deletion(-)

diff --git a/nirtorch/from_nir.py b/nirtorch/from_nir.py
index 5018169..9faeb7e 100644
--- a/nirtorch/from_nir.py
+++ b/nirtorch/from_nir.py
@@ -70,6 +70,10 @@ def __init__(self, graph: Graph) -> None:
     def _is_module_stateful(self, module: torch.nn.Module) -> bool:
         signature = inspect.signature(module.forward)
         arguments = len(signature.parameters)
+        # HACK for snntorch modules
+        if 'snntorch' in str(module.__class__):
+            if module.__class__.__name__ in ['Synaptic', 'RSynaptic', 'Leaky', 'RLeaky']:
+                return not module.init_hidden
         return arguments > 1
 
     def get_execution_order(self) -> List[Node]:
@@ -131,7 +135,13 @@ def _apply_module(
 
         out = node.elem(*inputs)
         # If the module is stateful, we know the output is (at least) a tuple
-        if self.stateful_modules[node.name]:
+        # HACK to make it work for snnTorch
+        is_rsynaptic = 'snntorch._neurons.rsynaptic.RSynaptic' in str(node.elem.__class__)
+        if is_rsynaptic and not node.elem.init_hidden:
+            assert 'lif' in node.name, "this shouldnt happen.."
+            new_state.state[node.name] = out  # snnTorch requires output inside state
+            out = out[0]
+        elif self.stateful_modules[node.name]:
             new_state.state[node.name] = out[1:]  # Store the new state
             out = out[0]
         return out, new_state

From 668e023aa58c601453ba55e3884c25110c2e478a Mon Sep 17 00:00:00 2001
From: Steve Abreu <s.abreu@rug.nl>
Date: Thu, 19 Oct 2023 17:09:00 +0200
Subject: [PATCH 13/27] + optional model.forward args for stateful modules

---
 nirtorch/graph.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/nirtorch/graph.py b/nirtorch/graph.py
index 1d3f5d3..ec1b8ae 100644
--- a/nirtorch/graph.py
+++ b/nirtorch/graph.py
@@ -386,7 +386,7 @@ def __exit__(self, exc_type, exc_value, exc_tb):
 
 
 def extract_torch_graph(
-    model: nn.Module, sample_data: Any, model_name: Optional[str] = "model"
+    model: nn.Module, sample_data: Any, model_name: Optional[str] = "model", model_args=[]
 ) -> Graph:
     """Extract computational graph between various modules in the model
     NOTE: This method is not capable of any compute happening outside of module
@@ -409,6 +409,6 @@ def extract_torch_graph(
     with GraphTracer(
         named_modules_map(model, model_name=model_name)
     ) as tracer, torch.no_grad():
-        _ = model(sample_data)
+        _ = model(sample_data, *model_args)
 
     return tracer.graph

From c555b2a0b83f2aff0d2c0732fd14f3d06ce364d2 Mon Sep 17 00:00:00 2001
From: Steve Abreu <s.abreu@rug.nl>
Date: Thu, 19 Oct 2023 17:12:44 +0200
Subject: [PATCH 14/27] change subgraphs handlign (flatten + remove I/O)

---
 nirtorch/to_nir.py | 50 +++++++++++++++++++++++++++++++++++++++++-----
 1 file changed, 45 insertions(+), 5 deletions(-)

diff --git a/nirtorch/to_nir.py b/nirtorch/to_nir.py
index f680d19..5d1c7a1 100644
--- a/nirtorch/to_nir.py
+++ b/nirtorch/to_nir.py
@@ -53,21 +53,38 @@ def extract_nir_graph(
 
     # Convert the nodes and get indices
     nir_edges = []
-    nir_nodes = {"input": nir.Input(np.array(sample_data.shape))}
-
+    input_shape = np.array(sample_data.shape)
+    # HACK: ignore dimensions
+    input_shape = np.array([e for idx, e in enumerate(sample_data.shape) if idx not in ignore_dims])
+    nir_nodes = {"input": nir.Input(input_shape)}
+
+    subgraph_keys = []
+    subgraph_input_nodekeys = []
+    subgraph_output_nodekeys = []
     # Get all the NIR nodes
     for indx, node in enumerate(torch_graph.node_list):
         # Convert the node type to NIR subgraph
         mapped_node = model_map(node.elem)
 
         if isinstance(mapped_node, nir.NIRGraph):
+            subgraph_keys.append(node.name)
             for k, v in mapped_node.nodes.items():
                 # For now, we add nodes in subgraphs to the top-level node list
-                # TODO: Parse graphs recursively
+                # TODO: support deeper nesting -> parse graphs recursively
+                assert not isinstance(v, nir.NIRGraph), "cannot handle sub-sub-graphs"
+
+                subgraph_node_key = f"{node.name}.{k}"
+
+                # keep track of subgraph input and outputs (to remove later)
+                if isinstance(v, nir.Input):
+                    subgraph_input_nodekeys.append(subgraph_node_key)
+                elif isinstance(v, nir.Output):
+                    subgraph_output_nodekeys.append(subgraph_node_key)
+
                 if isinstance(v, nir.NIRNode):
-                    nir_nodes[f"{node.name}.{k}"] = v
+                    nir_nodes[subgraph_node_key] = v
                 else:
-                    nir_nodes[v.name] = v
+                    nir_nodes[v.name] = v  # would this ever happen??
             # Add edges from graph
             for x, y in mapped_node.edges:
                 nir_edges.append((f"{node.name}.{x}", f"{node.name}.{y}"))
@@ -96,4 +113,27 @@ def extract_nir_graph(
     # Remove duplicate edges
     nir_edges = list(set(nir_edges))
 
+    # change edges to subgraph to point to either input or output of subgraph
+    for idx in range(len(nir_edges)):
+        if nir_edges[idx][0] in subgraph_keys:
+            nir_edges[idx] = (f"{nir_edges[idx][0]}.output", nir_edges[idx][1])
+        if nir_edges[idx][1] in subgraph_keys:
+            nir_edges[idx] = (nir_edges[idx][0], f"{nir_edges[idx][1]}.input")
+
+    # remove subgraph input and output nodes (& redirect edges)
+    for rm_nodekey in subgraph_input_nodekeys + subgraph_output_nodekeys:
+        in_keys = [e[0] for e in nir_edges if e[1] == rm_nodekey]
+        out_keys = [e[1] for e in nir_edges if e[0] == rm_nodekey]
+        # connect all incoming to all outgoing nodes
+        for in_key in in_keys:
+            for out_key in out_keys:
+                nir_edges.append((in_key, out_key))
+        # remove the original edges
+        for in_key in in_keys:
+            nir_edges.remove((in_key, rm_nodekey))
+        for out_key in out_keys:
+            nir_edges.remove((rm_nodekey, out_key))
+        # remove the node
+        nir_nodes.pop(rm_nodekey)
+
     return nir.NIRGraph(nir_nodes, nir_edges)

From 60c01f84af7a7529c503bb23414f6fb2ce36f316 Mon Sep 17 00:00:00 2001
From: Steve Abreu <s.abreu@rug.nl>
Date: Thu, 19 Oct 2023 17:13:31 +0200
Subject: [PATCH 15/27] model fwd args + ignore_dims arg

---
 nirtorch/to_nir.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/nirtorch/to_nir.py b/nirtorch/to_nir.py
index 5d1c7a1..a114efe 100644
--- a/nirtorch/to_nir.py
+++ b/nirtorch/to_nir.py
@@ -13,6 +13,8 @@ def extract_nir_graph(
     sample_data: Any,
     model_name: Optional[str] = "model",
     ignore_submodules_of=None,
+    model_fwd_args=[],
+    ignore_dims=[],
 ) -> nir.NIRNode:
     """Given a `model`, generate an NIR representation using the specified `model_map`.
 
@@ -38,7 +40,7 @@ def extract_nir_graph(
 
     # Extract a torch graph given the model
     torch_graph = extract_torch_graph(
-        model, sample_data=sample_data, model_name=model_name
+        model, sample_data=sample_data, model_name=model_name, model_args=model_fwd_args
     ).ignore_tensors()
 
     if ignore_submodules_of is not None:

From d4b1afb35b6104971c3c458702a75ef12b874256 Mon Sep 17 00:00:00 2001
From: Steve Abreu <s.abreu@rug.nl>
Date: Thu, 19 Oct 2023 17:18:26 +0200
Subject: [PATCH 16/27] [hack] remove wrong RNN self-connection (NIRTorch)

---
 nirtorch/to_nir.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/nirtorch/to_nir.py b/nirtorch/to_nir.py
index a114efe..500bf2a 100644
--- a/nirtorch/to_nir.py
+++ b/nirtorch/to_nir.py
@@ -138,4 +138,10 @@ def extract_nir_graph(
         # remove the node
         nir_nodes.pop(rm_nodekey)
 
+    # HACK: remove self-connections (this is a bug in the extraction of an RNN graph)
+    for edge in nir_edges:
+        if edge[0] == edge[1]:
+            print(f"[WARNING] removing self-connection {edge}")
+            nir_edges.remove(edge)
+
     return nir.NIRGraph(nir_nodes, nir_edges)

From c736c0e15233e1cfbcf07cc25adf0a84df43c0c2 Mon Sep 17 00:00:00 2001
From: "Jens E. Pedersen" <jensegholm@protonmail.com>
Date: Thu, 19 Oct 2023 22:42:58 +0200
Subject: [PATCH 17/27] Added proper graph tracing

---
 nirtorch/from_nir.py        | 107 +++++++++++++++++-------------------
 nirtorch/graph.py           |  44 ++++++++++-----
 nirtorch/graph_utils.py     |  25 +++++++++
 nirtorch/to_nir.py          |  10 +---
 tests/braille.nir           | Bin 41896 -> 34264 bytes
 tests/test_bidirectional.py |  24 ++++----
 tests/test_conversion.py    |   6 +-
 tests/test_from_nir.py      |  29 +++++++---
 tests/test_graph.py         |   6 +-
 tests/test_graph_utils.py   |  45 +++++++++++++++
 tests/test_to_nir.py        |  15 +++--
 11 files changed, 199 insertions(+), 112 deletions(-)
 create mode 100644 tests/test_graph_utils.py

diff --git a/nirtorch/from_nir.py b/nirtorch/from_nir.py
index 9faeb7e..4ad5cf6 100644
--- a/nirtorch/from_nir.py
+++ b/nirtorch/from_nir.py
@@ -7,46 +7,10 @@
 import torch.nn as nn
 
 from .graph import Graph, Node
+from .graph_utils import trace_execution
 from .utils import sanitize_name
 
 
-def execution_order_up_to_node(
-    node: Node,
-    graph: Graph,
-    execution_order: List[Node],
-    visited: Optional[Dict[Node, bool]] = None,
-) -> List[Node]:
-    """Recursive function to evaluate execution order until a given node.
-
-    Args:
-        node (Node): Execution order for the node of interest
-        graph (Graph): Graph object describing the network
-        execution_order (List[Node]): The current known execution order.
-
-    Returns:
-        List[Node]: Execution order
-    """
-    if visited is None:
-        visited = {n: False for n in graph.node_list}
-    is_recursive = False
-    if len(execution_order) == list(graph.node_list):
-        # All nodes are executed
-        return execution_order
-    for parent in graph.find_source_nodes_of(node):
-        if parent not in execution_order and not visited[parent]:
-            visited[parent] = True
-            execution_order = execution_order_up_to_node(
-                parent, graph, execution_order, visited
-            )
-        if node in parent.outgoing_nodes:
-            is_recursive = True
-    # Ensure we're not re-adding a recursive node
-    if is_recursive and node in execution_order:
-        return execution_order
-    else:  # Finally since all parents are known and executed
-        return execution_order + [node]
-
-
 @dataclasses.dataclass
 class GraphExecutorState:
     """State for the GraphExecutor that keeps track of both the state of hidden
@@ -71,21 +35,25 @@ def _is_module_stateful(self, module: torch.nn.Module) -> bool:
         signature = inspect.signature(module.forward)
         arguments = len(signature.parameters)
         # HACK for snntorch modules
-        if 'snntorch' in str(module.__class__):
-            if module.__class__.__name__ in ['Synaptic', 'RSynaptic', 'Leaky', 'RLeaky']:
+        if "snntorch" in str(module.__class__):
+            if module.__class__.__name__ in [
+                "Synaptic",
+                "RSynaptic",
+                "Leaky",
+                "RLeaky",
+            ]:
                 return not module.init_hidden
         return arguments > 1
 
     def get_execution_order(self) -> List[Node]:
         """Evaluate the execution order and instantiate that as a list."""
-        execution_order = []
-        # Then loop over all nodes and check that they are added to the execution order.
-        for node in self.graph.node_list:
-            if node not in execution_order:
-                execution_order = execution_order_up_to_node(
-                    node, self.graph, execution_order
-                )
-        return execution_order
+        # TODO: Adapt this for graphs with multiple inputs
+        inputs = self.graph.inputs
+        if len(inputs) != 1:
+            raise ValueError(
+                f"Currently, only one input is supported, but {len(inputs)} was given"
+            )
+        return trace_execution(inputs[0], lambda n: n.outgoing_nodes.keys())
 
     def instantiate_modules(self):
         for mod, name in self.graph.module_names.items():
@@ -136,9 +104,11 @@ def _apply_module(
         out = node.elem(*inputs)
         # If the module is stateful, we know the output is (at least) a tuple
         # HACK to make it work for snnTorch
-        is_rsynaptic = 'snntorch._neurons.rsynaptic.RSynaptic' in str(node.elem.__class__)
+        is_rsynaptic = "snntorch._neurons.rsynaptic.RSynaptic" in str(
+            node.elem.__class__
+        )
         if is_rsynaptic and not node.elem.init_hidden:
-            assert 'lif' in node.name, "this shouldnt happen.."
+            assert "lif" in node.name, "this shouldnt happen.."
             new_state.state[node.name] = out  # snnTorch requires output inside state
             out = out[0]
         elif self.stateful_modules[node.name]:
@@ -159,7 +129,11 @@ def forward(
             if node.elem is None:
                 continue
             out, new_state = self._apply_module(
-                node, input_nodes, new_state, old_state, data if first_node else None
+                node,
+                input_nodes,
+                new_state=new_state,
+                old_state=old_state,
+                data=data if first_node else None,
             )
             new_state.cache[node.name] = out
             first_node = False
@@ -170,18 +144,37 @@ def forward(
         return new_state.cache[node.name], new_state
 
 
-def _mod_nir_to_graph(nir_graph: nir.NIRNode) -> Graph:
-    module_names = {module: name for name, module in nir_graph.nodes.items()}
-    graph = Graph(module_names=module_names)
-    for src, dst in nir_graph.edges:
-        graph.add_edge(nir_graph.nodes[src], nir_graph.nodes[dst])
+def _mod_nir_to_graph(
+    torch_graph: nir.NIRGraph, nir_nodes: Dict[str, nir.NIRNode]
+) -> Graph:
+    module_names = {module: name for name, module in torch_graph.nodes.items()}
+    inputs = [name for name, node in nir_nodes.items() if isinstance(node, nir.Input)]
+    graph = Graph(module_names=module_names, inputs=inputs)
+    for src, dst in torch_graph.edges:
+        # Allow edges to refer to subgraph inputs and outputs
+        if not src in torch_graph.nodes and f"{src}.output" in torch_graph.nodes:
+            src = f"{src}.output"
+        if not dst in torch_graph.nodes and f"{dst}.input" in torch_graph.nodes:
+            dst = f"{dst}.input"
+        graph.add_edge(torch_graph.nodes[src], torch_graph.nodes[dst])
     return graph
 
 
+def _switch_default_models(nir_graph: nir.NIRNode) -> Optional[torch.nn.Module]:
+    if isinstance(nir_graph, nir.Input) or isinstance(nir_graph, nir.Output):
+        return torch.nn.Identity()
+
+
 def _switch_models_with_map(
     nir_graph: nir.NIRNode, model_map: Callable[[nn.Module], nn.Module]
 ) -> nir.NIRNode:
-    nodes = {name: model_map(node) for name, node in nir_graph.nodes.items()}
+    nodes = {}
+    for name, node in nir_graph.nodes.items():
+        mapped_module = model_map(node)
+        if mapped_module is None:
+            mapped_module = _switch_default_models(node)
+        nodes[name] = mapped_module
+    # nodes = {name: model_map(node) for name, node in nir_graph.nodes.items()}
     return nir.NIRGraph(nodes, nir_graph.edges)
 
 
@@ -204,6 +197,6 @@ def load(
     # Map modules to the target modules using th emodel map
     nir_module_graph = _switch_models_with_map(nir_graph, model_map)
     # Build a nirtorch.Graph based on the nir_graph
-    graph = _mod_nir_to_graph(nir_module_graph)
+    graph = _mod_nir_to_graph(nir_module_graph, nir_nodes=nir_graph.nodes)
     # Build and return a graph executor module
     return GraphExecutor(graph)
diff --git a/nirtorch/graph.py b/nirtorch/graph.py
index ec1b8ae..fde07c4 100644
--- a/nirtorch/graph.py
+++ b/nirtorch/graph.py
@@ -71,15 +71,19 @@ class Graph:
     def __init__(
         self,
         module_names: Dict[nn.Module, str],
+        inputs: List[str],
         module_output_types: Dict[nn.Module, torch.Tensor] = {},
     ) -> None:
         self.module_names = module_names
         self.node_list: List[Node] = []
         self.module_output_types = module_output_types
         self._last_used_tensor_id = None
+        self.inputs = []
         # Add modules to node_list
         for mod, name in self.module_names.items():
-            self.add_elem(mod, name)
+            node = self.add_elem(mod, name)
+            if name in inputs:
+                self.inputs.append(node)
 
     @property
     def node_map_by_id(self):
@@ -199,7 +203,9 @@ def debug_str(self) -> str:
         for node in self.node_list:
             debug_str += f"{node.name} ({node.elem.__class__.__name__})\n"
             for outgoing, shape in node.outgoing_nodes.items():
-                debug_str += f"\t-> {outgoing.name} ({outgoing.elem.__class__.__name__})\n"
+                debug_str += (
+                    f"\t-> {outgoing.name} ({outgoing.elem.__class__.__name__})\n"
+                )
         return debug_str.strip()
 
     def to_md(self) -> str:
@@ -215,7 +221,7 @@ def to_md(self) -> str:
 
     def leaf_only(self) -> "Graph":
         leaf_modules = self.get_leaf_modules()
-        filtered_graph = Graph(leaf_modules)
+        filtered_graph = Graph(leaf_modules, inputs=self.inputs)
         # Populate edges
         filtered_graph.populate_from(self)
         return filtered_graph
@@ -241,7 +247,11 @@ def ignore_submodules_of(self, classes: List[Type]) -> "Graph":
             if mod not in sub_modules_to_ignore:
                 new_named_modules[mod] = name
         # Create a new graph with the allowed modules
-        new_graph = Graph(new_named_modules, self.module_output_types)
+        new_graph = Graph(
+            new_named_modules,
+            inputs=self.inputs,
+            module_output_types=self.module_output_types,
+        )
         new_graph.populate_from(self)
         return new_graph
 
@@ -256,7 +266,7 @@ def find_source_nodes_of(self, node: Node) -> List[Node]:
         """
         source_node_list = []
         for source_node in self.node_list:
-            for outnode, shape in source_node.outgoing_nodes.items():
+            for outnode, _ in source_node.outgoing_nodes.items():
                 if node == outnode:
                     source_node_list.append(source_node)
         return source_node_list
@@ -276,7 +286,11 @@ def ignore_nodes(self, class_type: Type) -> "Graph":
         }
 
         # Generate the new graph with the filtered module names
-        graph = Graph(new_module_names, self.module_output_types)
+        graph = Graph(
+            new_module_names,
+            inputs=self.inputs,
+            module_output_types=self.module_output_types,
+        )
         # Iterate over all the nodes
         for node in self.node_list:
             if isinstance(node.elem, class_type):
@@ -308,13 +322,7 @@ def get_root(self) -> List[Node]:
         Returns:
             List[Node]: A list of root nodes for the graph.
         """
-        roots = []
-        for node in self.node_list:
-            sources = self.find_source_nodes_of(node)
-            # Append root node if it has no sources (and it isn't a sequential module)
-            if len(sources) == 0 and not isinstance(node.elem, torch.nn.Sequential):
-                roots.append(node)
-        return roots
+        return self.inputs
 
 
 _torch_module_call = torch.nn.Module.__call__
@@ -386,7 +394,10 @@ def __exit__(self, exc_type, exc_value, exc_tb):
 
 
 def extract_torch_graph(
-    model: nn.Module, sample_data: Any, model_name: Optional[str] = "model", model_args=[]
+    model: nn.Module,
+    sample_data: Any,
+    model_name: Optional[str] = "model",
+    model_args=[],
 ) -> Graph:
     """Extract computational graph between various modules in the model
     NOTE: This method is not capable of any compute happening outside of module
@@ -411,4 +422,9 @@ def extract_torch_graph(
     ) as tracer, torch.no_grad():
         _ = model(sample_data, *model_args)
 
+    # HACK: The current graph is using copy-constructors, that detaches
+    # the traced output_types from the original graph.
+    # In the future, find a way to synchronize the two representations
+    tracer.graph.module_output_types = tracer.output_types
+
     return tracer.graph
diff --git a/nirtorch/graph_utils.py b/nirtorch/graph_utils.py
index 54ccfd4..1f67d71 100644
--- a/nirtorch/graph_utils.py
+++ b/nirtorch/graph_utils.py
@@ -1,3 +1,8 @@
+from typing import Callable, List, Set, TypeVar
+
+T = TypeVar("T")
+
+
 def find_children(node, edges):
     """Given a node and the edges of a graph, find all direct children of that node."""
     return set(child for (parent, child) in edges if parent == node)
@@ -59,3 +64,23 @@ def find_all_ancestors(
 #
 #     return execution_order
 #
+
+def trace_execution(
+    node: T, edge_fn: Callable[[T], List[T]], visited: Set[T] = None
+) -> List[T]:
+    """Traces the execution of a node by listing them in order, coloring recursive
+    nodes to avoid adding the same node twice.
+    """
+    if visited is None:
+        visited = set()
+
+    if node in visited:
+        return []
+    else:
+        visited.add(node)
+
+    successors = []
+    for child in edge_fn(node):
+        if child not in visited:
+            successors += trace_execution(child, edge_fn, visited)
+    return [node] + successors
\ No newline at end of file
diff --git a/nirtorch/to_nir.py b/nirtorch/to_nir.py
index 500bf2a..df01ada 100644
--- a/nirtorch/to_nir.py
+++ b/nirtorch/to_nir.py
@@ -1,4 +1,5 @@
 from typing import Any, Callable, Optional
+import logging
 
 import nir
 import numpy as np
@@ -46,13 +47,6 @@ def extract_nir_graph(
     if ignore_submodules_of is not None:
         torch_graph = torch_graph.ignore_submodules_of(ignore_submodules_of)
 
-    # Get the root node
-    root_nodes = torch_graph.get_root()
-    if len(root_nodes) != 1:
-        raise ValueError(
-            f"Currently, only one input is supported, but {len(root_nodes)} was given"
-        )
-
     # Convert the nodes and get indices
     nir_edges = []
     input_shape = np.array(sample_data.shape)
@@ -141,7 +135,7 @@ def extract_nir_graph(
     # HACK: remove self-connections (this is a bug in the extraction of an RNN graph)
     for edge in nir_edges:
         if edge[0] == edge[1]:
-            print(f"[WARNING] removing self-connection {edge}")
+            logging.warn(f"removing self-connection {edge}")
             nir_edges.remove(edge)
 
     return nir.NIRGraph(nir_nodes, nir_edges)
diff --git a/tests/braille.nir b/tests/braille.nir
index ed8f57e7f4909bd10ea0e92cf2fba67f6c7c4efe..111b0a44eed106e15c89346c38ae988a4bea7a5d 100644
GIT binary patch
delta 2591
zcmc&$Uu+ab7{8g8c6Xs=+t}V=4(@`9Tq}3p9i|>txj%T^da+4~8nF+xfYxg@b}b-<
z8ut)Q2%vWPUYfp8h=~L>xy?p$A|^-D1Y&%+=*yY(0exsld?E2cV?yAY-C4TDo)OZ(
zKFs(1e&5V|GvCa9-<-QbYbWSbh}J$ax(5ZgF36+_jtE}TY_$UXhPnuCqL)Z;S@LdO
z?Xf?J%XdbGW08ZekH5LsCf0CXaNDs{3mYA4tz%tVhnIm!UeX3JY}?^I(F;{E2!bS&
z4#s6V`M7$k@z$A_)1GpRNmj3I(JQA-**8t-)sJ;!15{oU%<8%Xf!(=)QTdgC>=W8c
zd`v%=rJ`;OLZ#VfPOVoZj28)gitXp2u8j$&>l?O0IG*5mdKhR|1j?c4=&>D3yOcPt
zDmI^ThK?@n*T3J!vz+JMJ-)Y>_Vayc;{>iYVph9mB@#7^W(uXO0L303RI_KH_vuRz
z$w_qAcD{S)PXwWu6HX-*sOAJ(cpeQ>)Sk*Sk=2cNK~+Yn(hKv-7utn=1-)-ifAS$o
z?O{3Qezkkw`OopAZygVYlPe;%m73^2=h0SdC;zd4B=MAP-FhswjMVBG9uTk9UD<<K
zfVWO)45Npb3w_<Br3wFBnP=oCxcSnTjm^k_x-FzJ_I__+c)S0P-+Q`~_bTUp`R0)S
z3Q3_1JIeo(y0Hxw`r=FWBSLIT8qO<GI`ND-ue1_+?pa&bp0lNp1ph0Sq{Y+<{|sy*
zd2=A#N$Cv*#-IF-q+oM&4%^FnHc{mTCk@teL9;TMAyVl@JNHd>!q@SWrmCT$V)h%4
z$Apw^nuhN)-%73&KeOs_8Z6je%ERcTOs`pPXlfsGAU`1x8Fr5ZSiPZz5#3mTU_(3P
zw-%%nPvfCR^!FJYUezfrW_*~*2FNH}%nn(G#T0L5dLxYmcC_qcC+i$KtAngYI+)eX
z1{Zz=^+p%&xC^%-c3+})#4xM+JJ_V|iP`f29}Q|C>~_z}z8kZ0?{+U^75<ZAM1xSs
z<i#5)wBzT8ueSUk77yA8>(*ZcBQ!wQq~U{xi&n|Hewd#Xy*j}(ELya7h=-S|?=y{6
zu)il!6P*esz|P44R7X^F_b9tNqCagrN#H4-hI}Sru`e&no?Mol`xpCTJ|^5DaVJ#$
z$T*KzZdVbyLxXxnfY5q5Z9Sl{vHg2IcR#3x_r@&~@|H4oV840*t3TlbhIHdqD=JTA
wFe-DjeAL<f&=JFEC010G%^6x4gZY`C);ZUM^6_^tobjb4b<C!(dDun24Xz^oC;$Ke

delta 3435
zcmc&$T}&KR6rQ=TFgxXEYwZ3|pc50z><VV6rrkCTY`1|f6opnKNMqfVmcJq_{Yh=D
z%r^0XXpOALG)qfxTa6Mz43jZ&V=!40j6U$dlOaCQ#F(b>fkz*V@!q-jvIB(1)Y48e
zbM85F?mg$+d(L;y@<(X;XEd&&=`W1HX$dltODurvlE+sl&_`nQ)uxlj&d1_?5}c4c
zqLb)|G8p)~1MW$lrXr@7$LL`HkTojaB=4~#C0oc#;YX*Z(Y3|jotxy;$}^kf#0vcC
zd=Bp8kaj`R=@H%JTkBz}`HT+s_v|A>?-p_&aY!#?e+wAx>pyemEVjLS3wb${^X=&A
zuC*i;TgWS!+`(x7Q2f$GHpT=AiW!qV#iq-f5>e`I(x?O_D49#XpeSmFQHbOq(^e{A
zCQ<3{>9(R0czE35;Y$~-l|`<eG(3<k^qS)paT%YzsD6yUpQwJ2{90M|b*jd-?=xQN
zgx}N`n-#W5=h5|Q3C5=zs9a-u5q-U0YLfDqsie_iF8R~48Le?4NbL81+H@BV9hiqH
zpA2om2#!`Y(|uu4jvS(M)O?uTcf$;w!-Lh3s<V?xPfXma#~M$X>QCC>27CG=M~pXX
zh3we>qx57S!$ObfNUt^TaYXWP2h7rXzX(dbw38sCi@_s7%Od#U-3EAYTto3Xn7!~N
zsVX^F32@7?7pce0ftTF~WgFo8fhc_F?L;@~;P8P3G~Uc=8pR`c<m7OafMoL?aJ?WQ
zZzGnADj3aGkPS*`u03fCLFj@4(e5%he^5h{b>_F>5)m!dQ~GzvqVik|-5gDX*_S4b
z?0(TKsSk)c0UI6x34}z!c_*EcU0nQ=B*%78eY+O^kL&8&64{2jd`Ma9PY0>i9Mqo@
zRB{LM$$-~WiSb=PgrHJ{yNZU;*Nz=H>GneH!4b6b3jE$MFWasjTuILHH)w)NW~`<Z
zA+-|%-VCY^L$$Vm#%{A>3M_X~HGY8!rzo$y$)-~xBs!DEA6CWYG+?H&N{TXla>pJ~
znRbZHW~{((Ykb*SBHsLLylgdydHrzJ;g+*(g_XA~T$p>?m~kU0yzYXa+M(wQ3|E8{
z)8qc<=26?cQIgJ5p6h9hZ)3Q~;Us=}y<Up(nJZMT5K@8P=}|Hg1iqkflk!WWN#iR}
z^16_#6|Utipuf7wUK4;y#|)!Hv=Av$9Y%-4Q2xO;R<t>9tZ0|}C<?}V*HBPtwNNlL
z$PNxE19&8Jx3zQJ$S=J_&JUvP9U>=a7#-S(hFU8{!Dt($<0mN1o}{Q4AEpseR7_?I
zEKr1{;qvE6J4re>4u)ix?L28i^FWO6(d|;w(vmKO&nn%(KZ|(+vE;EsFX{d4>yYiY
z^Q4!wuS<GS#&_T$Z9?r6c43*FcVnKkyWBC-0^51A%XR}kb40%Fu&N=KAJ)l*-ZQ&u
z@cp%X1s+5)X!SL#+*0qe(-AdDUamr8*GMs?k6%S`9m-oiW0zDj653(5d%ifAOV}_J
ztGm(aC<N5g)Z$uCXwJ0UbVJ&qAZ@~!`3K*+IO|gFIO~2AZ}EIkQtOg#x*aFoqK`NC
z2pb^J8V`58N91&zN)rNUAz`IVL7_A$K(st%`LX0bC5~9xJ$8p1(muBtI*NZ2JLDM2
S!|kl`5NirMN9fPAW#(T;`zek9

diff --git a/tests/test_bidirectional.py b/tests/test_bidirectional.py
index 3a8abce..5345935 100644
--- a/tests/test_bidirectional.py
+++ b/tests/test_bidirectional.py
@@ -42,7 +42,9 @@ def _nir_to_pytorch_module(node: nir.NIRNode) -> torch.nn.Module:
 
 def _create_torch_model() -> torch.nn.Module:
     if use_snntorch:
-        return torch.nn.Sequential(torch.nn.Linear(1, 1), snn.Leaky(0.9, init_hidden=True))
+        return torch.nn.Sequential(
+            torch.nn.Linear(1, 1), snn.Leaky(0.9, init_hidden=True)
+        )
     else:
         return torch.nn.Sequential(torch.nn.Linear(1, 1), torch.nn.Identity())
 
@@ -57,23 +59,21 @@ def _torch_to_nir(module: torch.nn.Module) -> nir.NIRNode:
 
 def _lif_nir_graph(from_file=True):
     if from_file:
-        return nir.read('tests/lif_norse.nir')
+        return nir.read("tests/lif_norse.nir")
     else:
         return nir.NIRGraph(
             nodes={
-                '0': nir.Affine(weight=np.array([[1.]]), bias=np.array([0.])),
-                '1': nir.LIF(
+                "input": nir.Input(input_type={"input": np.array([1])}),
+                "0": nir.Affine(weight=np.array([[1.0]]), bias=np.array([0.0])),
+                "1": nir.LIF(
                     tau=np.array([0.1]),
-                    r=np.array([1.]),
-                    v_leak=np.array([0.]),
-                    v_threshold=np.array([0.1])
+                    r=np.array([1.0]),
+                    v_leak=np.array([0.0]),
+                    v_threshold=np.array([0.1]),
                 ),
-                'input': nir.Input(input_type={'input': np.array([1])}),
-                'output': nir.Output(output_type={'output': np.array([1])})
+                "output": nir.Output(output_type={"output": np.array([1])}),
             },
-            edges=[
-                ('input', '0'), ('0', '1'), ('1', 'output')
-            ]
+            edges=[("input", "0"), ("0", "1"), ("1", "output")],
         )
 
 
diff --git a/tests/test_conversion.py b/tests/test_conversion.py
index 4e2f0db..c04be4f 100644
--- a/tests/test_conversion.py
+++ b/tests/test_conversion.py
@@ -7,14 +7,12 @@
 
 def _torch_convert(module: nn.Module) -> nir.NIRNode:
     if isinstance(module, nn.Conv1d):
-        return nir.Conv1d(module.weight, 1, 1, 1, 1, module.bias)
+        return nir.Conv1d(None, module.weight, 1, 1, 1, 1, module.bias)
     elif isinstance(module, nn.Linear):
         return nir.Affine(module.weight, module.bias)
-    else:
-        raise NotImplementedError(f"Unsupported module {module}")
 
 
-def test_norse_to_sinabs():
+def test_extract_pytorch():
     model = torch.nn.Sequential(
         torch.nn.Conv1d(1, 2, 3),
         torch.nn.Linear(8, 1),
diff --git a/tests/test_from_nir.py b/tests/test_from_nir.py
index d16bc68..b119316 100644
--- a/tests/test_from_nir.py
+++ b/tests/test_from_nir.py
@@ -56,7 +56,10 @@ def test_extract_empty():
 
 
 def test_extract_illegal_name():
-    graph = nir.NIRGraph({"a.b": nir.Linear(np.ones((1, 1)))}, [])
+    graph = nir.NIRGraph(
+        {"i": nir.Input(np.ones((1, 1))), "a.b": nir.Linear(np.ones((1, 1)))},
+        [("i", "a.b")],
+    )
     torch_graph = load(graph, _torch_model_map)
     assert "a_b" in torch_graph._modules
 
@@ -68,14 +71,16 @@ def test_extract_lin():
     torchlin.weight.data = torch.nn.Parameter(lin.weight)
     torchlin.bias.data = torch.nn.Parameter(lin.bias)
     y = torchlin(torchlin(x))
-    g = nir.NIRGraph({"a": lin, "b": lin}, [("a", "b")])
+    g = nir.NIRGraph(
+        {"i": nir.Input(np.ones((1, 1))), "a": lin, "b": lin}, [("i", "a"), ("a", "b")]
+    )
     m = load(g, _torch_model_map)
-    assert isinstance(m.execution_order[0].elem, torch.nn.Linear)
-    assert torch.allclose(m.execution_order[0].elem.weight, lin.weight)
-    assert torch.allclose(m.execution_order[0].elem.bias, lin.bias)
     assert isinstance(m.execution_order[1].elem, torch.nn.Linear)
     assert torch.allclose(m.execution_order[1].elem.weight, lin.weight)
     assert torch.allclose(m.execution_order[1].elem.bias, lin.bias)
+    assert isinstance(m.execution_order[2].elem, torch.nn.Linear)
+    assert torch.allclose(m.execution_order[2].elem.weight, lin.weight)
+    assert torch.allclose(m.execution_order[2].elem.bias, lin.bias)
     assert torch.allclose(m(x)[0], y)
 
 
@@ -104,11 +109,19 @@ def forward(self, x, state=None):
                 state = 1
             return x + state, state
 
+    def _map_stateful(node):
+        if isinstance(node, nir.Flatten):
+            return StatefulModel()
+
     g = nir.NIRGraph(
-        nodes={"li": nir.Flatten(np.array([1])), "li2": nir.Flatten(np.array([1]))},
-        edges=[("li", "li2")],
+        nodes={
+            "i": nir.Input(np.array([1, 1])),
+            "li": nir.Flatten(np.array([1])),
+            "li2": nir.Flatten(np.array([1])),
+        },
+        edges=[("i", "li"), ("li", "li2")],
     )  # Mock node
-    m = load(g, lambda m: StatefulModel())
+    m = load(g, _map_stateful)
     out, state = m(torch.ones(10))
     assert torch.allclose(out, torch.ones(10) * 3)
     assert state.state["li"] == (1,)
diff --git a/tests/test_graph.py b/tests/test_graph.py
index c99e1ed..b48ae3d 100644
--- a/tests/test_graph.py
+++ b/tests/test_graph.py
@@ -109,7 +109,7 @@ def test_module_forward_wrapper():
     from nirtorch.graph import Graph, module_forward_wrapper, named_modules_map
 
     output_types = {}
-    model_graph = Graph(named_modules_map(mymodel))
+    model_graph = Graph(named_modules_map(mymodel), ["block1"])
     new_call = module_forward_wrapper(model_graph, output_types)
 
     # Override call to the new wrapped call
@@ -238,7 +238,7 @@ def test_root_has_no_source():
         len(graph.find_source_nodes_of(graph.find_node(my_branched_model.relu1))) == 0
     )
 
-
+@pytest.mark.skip(reason="Root tracing is broken")
 def test_get_root():
     graph = extract_torch_graph(my_branched_model, sample_data=data, model_name=None)
     graph = graph.ignore_tensors()
@@ -308,7 +308,7 @@ def forward(self, x, state=None):
     assert d.nodes.keys() == {"input", "l", "r", "output"}
     assert set(d.edges) == {("input", "r"), ("r", "l"), ("l", "output"), ("r", "r")}
 
-
+@pytest.mark.skip(reason="Subgraphs are currently flattened")
 def test_captures_recurrence_manually():
     def export_affine_rec_gru(module):
         if isinstance(module, torch.nn.Linear):
diff --git a/tests/test_graph_utils.py b/tests/test_graph_utils.py
new file mode 100644
index 0000000..36ec5f3
--- /dev/null
+++ b/tests/test_graph_utils.py
@@ -0,0 +1,45 @@
+from collections import defaultdict
+
+from nirtorch.graph_utils import trace_execution
+
+
+class StringNode:
+    def __init__(self, name, edges):
+        self.name = name
+        self.edges = edges
+
+    @staticmethod
+    def get_children(node):
+        return [StringNode(x, node.edges) for x in node.edges[node.name]]
+
+    @staticmethod
+    def from_string(graph):
+        edges = defaultdict(list)
+        for edge in graph.split(" "):
+            edges[edge[0]].append(edge[2])
+        return StringNode(graph[0], edges)
+
+    def __hash__(self) -> int:
+        return self.name.__hash__()
+    
+    def __eq__(self, other: object) -> bool:
+        return self.name == other.name
+
+
+def test_trace_linear():
+    graph = "a-b b-c c-d"
+    node = StringNode.from_string(graph)
+    seen = trace_execution(node, node.get_children)
+    assert "".join([x.name for x in seen]) == "abcd"
+
+
+def test_trace_recursive():
+    node = StringNode.from_string("a-b b-a")
+    seen = trace_execution(node, node.get_children)
+    assert "".join([x.name for x in seen]) == "ab"
+
+
+def test_trace_recursive_complex():
+    node = StringNode.from_string("a-b b-a b-c b-c c-d d-e")
+    seen = trace_execution(node, node.get_children)
+    assert "".join([x.name for x in seen]) == "abcde"
diff --git a/tests/test_to_nir.py b/tests/test_to_nir.py
index f0f0665..9685d64 100644
--- a/tests/test_to_nir.py
+++ b/tests/test_to_nir.py
@@ -1,18 +1,20 @@
 import nir
 import numpy as np
+import pytest
 import torch
 import torch.nn as nn
 
 from nirtorch.to_nir import extract_nir_graph
 
 
+def _node_to_affine(node):
+    if isinstance(node, torch.nn.Linear):
+        return nir.Affine(node.weight.detach().numpy(), node.bias.detach().numpy())
+
+
 def test_extract_single():
     m = nn.Linear(1, 1)
-    g = extract_nir_graph(
-        m,
-        lambda x: nir.Affine(x.weight.detach().numpy(), x.bias.detach().numpy()),
-        torch.rand(1, 1),
-    )
+    g = extract_nir_graph(m, _node_to_affine, torch.rand(1, 1))
     assert set(g.edges) == {("input", "model"), ("model", "output")}
     assert isinstance(g.nodes["input"], nir.Input)
     assert np.allclose(g.nodes["input"].input_type["input"], np.array([1, 1]))
@@ -61,6 +63,7 @@ def forward(self, x):
         return x @ self.a @ self.b
 
 
+@pytest.mark.skip(reason="Re-implement with correct recursive graph parsing")
 def test_extract_multiple_explicit():
     model = nn.Sequential(BranchedModel(1, 2, 3), nn.Linear(3, 4))
 
@@ -80,7 +83,7 @@ def extractor(module: nn.Module):
 
     g = extract_nir_graph(model, extractor, torch.rand(1))
     print([type(n) for n in g.nodes])
-    assert len(g.nodes) == 7
+    assert len(g.nodes) == 4
     assert len(g.edges) == 8  # in + 5 + 1 + out
 
 

From fe7188af95919d0b69ea3493c21729e08ba86f31 Mon Sep 17 00:00:00 2001
From: Steve Abreu <s.abreu@rug.nl>
Date: Fri, 20 Oct 2023 15:28:17 +0200
Subject: [PATCH 18/27] + arg to ignore dims in to_nir

---
 nirtorch/to_nir.py | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/nirtorch/to_nir.py b/nirtorch/to_nir.py
index 686bfa2..0074cf5 100644
--- a/nirtorch/to_nir.py
+++ b/nirtorch/to_nir.py
@@ -1,4 +1,4 @@
-from typing import Any, Callable, Optional
+from typing import Any, Callable, Optional, Sequence
 
 import nir
 import numpy as np
@@ -13,6 +13,7 @@ def extract_nir_graph(
     sample_data: Any,
     model_name: Optional[str] = "model",
     ignore_submodules_of=None,
+    ignore_dims: Sequence[int] = [],
 ) -> nir.NIRNode:
     """Given a `model`, generate an NIR representation using the specified `model_map`.
 
@@ -49,9 +50,10 @@ def extract_nir_graph(
 
     # Convert the nodes and get indices
     nir_edges = []
-    nir_nodes = {
-        "input": nir.Input(np.array(sample_data.shape[1:]))
-    }  # Remove the first dimension
+    input_shape = np.array(sample_data.shape)
+    input_shape = np.delete(input_shape, ignore_dims)
+    nir_nodes = {"input": nir.Input(input_shape)}
+    nir_edges = []
 
     # Get all the NIR nodes
     for indx, node in enumerate(torch_graph.node_list):

From a21819f12197fd1c1c91ddbc424da3785613b31f Mon Sep 17 00:00:00 2001
From: Steve Abreu <s.abreu@rug.nl>
Date: Fri, 20 Oct 2023 15:28:21 +0200
Subject: [PATCH 19/27] add tests

---
 tests/test_to_nir.py | 32 ++++++++++++++++++++++++++++++++
 1 file changed, 32 insertions(+)

diff --git a/tests/test_to_nir.py b/tests/test_to_nir.py
index f0f0665..26fb0a1 100644
--- a/tests/test_to_nir.py
+++ b/tests/test_to_nir.py
@@ -105,6 +105,38 @@ def extractor(m):
     }
 
 
+def test_ignore_batch_dim():
+    model = nn.Linear(3, 1)
+
+    def extractor(module: nn.Module):
+        return nir.Affine(module.weight, module.bias)
+
+    raw_input_shape = (1, 3)
+    g = extract_nir_graph(model, extractor, torch.ones(raw_input_shape), ignore_dims=[0])
+    exp_input_shape = (3,)
+    assert np.alltrue(g.nodes["input"].input_type["input"] == np.array(exp_input_shape))
+    assert g.nodes["model"].weight.shape == (1, 3)
+    assert np.alltrue(g.nodes["output"].output_type["output"] == np.array([1]))
+
+
+def test_ignore_time_and_batch_dim():
+    model = nn.Linear(3, 1)
+
+    def extractor(module: nn.Module):
+        return nir.Affine(module.weight, module.bias)
+
+    raw_input_shape = (1, 10, 3)
+    g = extract_nir_graph(model, extractor, torch.ones(raw_input_shape), ignore_dims=[0, -2])
+    exp_input_shape = (3,)
+    assert np.alltrue(g.nodes["input"].input_type["input"] == np.array(exp_input_shape))
+    assert g.nodes["model"].weight.shape == (1, 3)
+
+    raw_input_shape = (1, 10, 3)
+    g = extract_nir_graph(model, extractor, torch.ones(raw_input_shape), ignore_dims=[0, 1])
+    exp_input_shape = (3,)
+    assert np.alltrue(g.nodes["input"].input_type["input"] == np.array(exp_input_shape))
+
+
 # def test_extract_stateful():
 #     model = norse.SequentialState(norse.LIFBoxCell(), nn.Linear(3, 1))
 

From bef454b141249058601f1dd83bd3756c07523ff1 Mon Sep 17 00:00:00 2001
From: Sadique Sheik <sadique.sheik@synsense.ai>
Date: Fri, 20 Oct 2023 16:02:41 +0200
Subject: [PATCH 20/27] output_shape also uses ignore_dims

---
 nirtorch/to_nir.py  | 26 ++++++++++++++++++--------
 tests/test_graph.py |  6 ++++--
 2 files changed, 22 insertions(+), 10 deletions(-)

diff --git a/nirtorch/to_nir.py b/nirtorch/to_nir.py
index 0074cf5..ac4c635 100644
--- a/nirtorch/to_nir.py
+++ b/nirtorch/to_nir.py
@@ -12,8 +12,8 @@ def extract_nir_graph(
     model_map: Callable[[nn.Module], nir.NIRNode],
     sample_data: Any,
     model_name: Optional[str] = "model",
-    ignore_submodules_of=None,
-    ignore_dims: Sequence[int] = [],
+    ignore_submodules_of: Optional[Sequence[nn.Module]] = None,
+    ignore_dims: Optional[Sequence[int]] = None,
 ) -> nir.NIRNode:
     """Given a `model`, generate an NIR representation using the specified `model_map`.
 
@@ -24,7 +24,11 @@ def extract_nir_graph(
         sample_data (Any): Sample input data to be used for model extraction
         model_name (Optional[str], optional): The name of the top level module.
             Defaults to "model".
-
+        ignore_submodules_of (Optional[Sequence[nn.Module]]): If specified,
+            the corresponding module's children will not be traversed for graph.
+        ignore_dims (Optional[Sequence[int]]): Dimensions of data to be ignored for
+            type/shape inference. Typically the dimensions that you will want to ignore
+            are for batch and time.
     Returns:
         nir.NIR: Returns the generated NIR graph representation.
     """
@@ -51,8 +55,10 @@ def extract_nir_graph(
     # Convert the nodes and get indices
     nir_edges = []
     input_shape = np.array(sample_data.shape)
-    input_shape = np.delete(input_shape, ignore_dims)
-    nir_nodes = {"input": nir.Input(input_shape)}
+    if ignore_dims:
+        nir_nodes = {"input": nir.Input(np.delete(input_shape, ignore_dims))}
+    else:
+        nir_nodes = {"input": nir.Input(input_shape)}
     nir_edges = []
 
     # Get all the NIR nodes
@@ -89,9 +95,13 @@ def extract_nir_graph(
         if len(node.outgoing_nodes) == 0:
             out_name = "output"
             # Try to find shape of input to the Output node
-            output_node = nir.Output(
-                torch_graph.module_output_types[node.elem][1:]
-            )  # Ignore batch dimension
+            if ignore_dims:
+                out_shape = np.delete(
+                    torch_graph.module_output_types[node.elem], ignore_dims
+                )
+            else:
+                out_shape = torch_graph.module_output_types[node.elem]
+            output_node = nir.Output(out_shape)
             nir_nodes[out_name] = output_node
             nir_edges.append((node.name, out_name))
 
diff --git a/tests/test_graph.py b/tests/test_graph.py
index b999721..0f66f62 100644
--- a/tests/test_graph.py
+++ b/tests/test_graph.py
@@ -278,9 +278,11 @@ def test_output_type_when_single_node():
 
 def test_sequential_flatten():
     d = torch.empty(2, 3, 4)
-    g = extract_nir_graph(torch.nn.Flatten(1), lambda x: nir.Flatten(d.shape, 1), d)
+    g = extract_nir_graph(
+        torch.nn.Flatten(1), lambda x: nir.Flatten(d.shape, 1), d, ignore_dims=[0]
+    )
     assert tuple(g.nodes["input"].input_type["input"]) == (3, 4)
-    assert tuple(g.nodes["output"].output_type["output"]) == (3 * 4, )
+    assert tuple(g.nodes["output"].output_type["output"]) == (3 * 4,)
 
 
 @pytest.mark.skip(reason="Not supported yet")

From b95ad5c478ac64b269db44e9426ef68515e749b1 Mon Sep 17 00:00:00 2001
From: "Jens E. Pedersen" <jensegholm@protonmail.com>
Date: Fri, 20 Oct 2023 16:06:45 +0200
Subject: [PATCH 21/27] Added test for flatten

---
 tests/test_graph.py | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/tests/test_graph.py b/tests/test_graph.py
index 0f66f62..60c9189 100644
--- a/tests/test_graph.py
+++ b/tests/test_graph.py
@@ -277,10 +277,13 @@ def test_output_type_when_single_node():
 
 
 def test_sequential_flatten():
+    d = torch.empty(3, 4)
+    g = extract_nir_graph(torch.nn.Flatten(1), lambda x: nir.Flatten(d.shape, 1), d)
+    assert tuple(g.nodes["input"].input_type["input"]) == (3, 4)
+    assert tuple(g.nodes["output"].output_type["output"]) == (3 * 4,)
+
     d = torch.empty(2, 3, 4)
-    g = extract_nir_graph(
-        torch.nn.Flatten(1), lambda x: nir.Flatten(d.shape, 1), d, ignore_dims=[0]
-    )
+    g = extract_nir_graph(torch.nn.Flatten(1), lambda x: nir.Flatten(d.shape, 1), d, ignore_dims=[0])
     assert tuple(g.nodes["input"].input_type["input"]) == (3, 4)
     assert tuple(g.nodes["output"].output_type["output"]) == (3 * 4,)
 

From 3bc8bd2e839b4f204840805a6b0a7bf0810952d4 Mon Sep 17 00:00:00 2001
From: Sadique Sheik <sadique.sheik@synsense.ai>
Date: Fri, 20 Oct 2023 22:25:29 +0200
Subject: [PATCH 22/27] minor correction to default value

---
 nirtorch/to_nir.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/nirtorch/to_nir.py b/nirtorch/to_nir.py
index 8cec0fe..47e7d03 100644
--- a/nirtorch/to_nir.py
+++ b/nirtorch/to_nir.py
@@ -1,4 +1,4 @@
-from typing import Any, Callable, Optional
+from typing import Any, Callable, Optional, Sequence
 import logging
 
 import nir
@@ -15,7 +15,7 @@ def extract_nir_graph(
     model_name: Optional[str] = "model",
     ignore_submodules_of=None,
     model_fwd_args=[],
-    ignore_dims=[],
+    ignore_dims: Optional[Sequence[int]]=None,
 ) -> nir.NIRNode:
     """Given a `model`, generate an NIR representation using the specified `model_map`.
 

From 84e3cc85a5175afef6f474a84c9f8fb51f82fa24 Mon Sep 17 00:00:00 2001
From: "Jens E. Pedersen" <jensegholm@protonmail.com>
Date: Sat, 21 Oct 2023 22:24:50 +0200
Subject: [PATCH 23/27] Added ability to ignore state in executor

---
 nirtorch/from_nir.py        | 33 +++++++++++++++++++++++++--------
 tests/test_bidirectional.py |  6 ++++--
 tests/test_from_nir.py      |  7 ++++++-
 3 files changed, 35 insertions(+), 11 deletions(-)

diff --git a/nirtorch/from_nir.py b/nirtorch/from_nir.py
index 4ad5cf6..d2324c0 100644
--- a/nirtorch/from_nir.py
+++ b/nirtorch/from_nir.py
@@ -22,10 +22,25 @@ class GraphExecutorState:
 
 
 class GraphExecutor(nn.Module):
-    def __init__(self, graph: Graph) -> None:
+    """Executes the NIR graph in PyTorch.
+
+    By default the graph executor is stateful, since there may be recurrence or stateful modules in the graph.
+    Specifically, that means accepting and returning a state object (`GraphExecutorState`).
+    If that is not desired, set `return_state=False` in the constructor.
+
+    Arguments:
+        graph (Graph): The graph to execute
+        return_state (bool, optional): Whether to return the state object. Defaults to True.
+
+    Raises:
+        ValueError: If there are no edges in the graph
+    """
+
+    def __init__(self, graph: Graph, return_state: bool = True) -> None:
         super().__init__()
         self.graph = graph
-        self.stateful_modules = {}
+        self.stateful_modules = set()
+        self.return_state = return_state
         self.instantiate_modules()
         self.execution_order = self.get_execution_order()
         if len(self.execution_order) == 0:
@@ -43,7 +58,7 @@ def _is_module_stateful(self, module: torch.nn.Module) -> bool:
                 "RLeaky",
             ]:
                 return not module.init_hidden
-        return arguments > 1
+        return "state" in signature.parameters and arguments > 1
 
     def get_execution_order(self) -> List[Node]:
         """Evaluate the execution order and instantiate that as a list."""
@@ -59,9 +74,8 @@ def instantiate_modules(self):
         for mod, name in self.graph.module_names.items():
             if mod is not None:
                 self.add_module(sanitize_name(name), mod)
-                self.stateful_modules[sanitize_name(name)] = self._is_module_stateful(
-                    mod
-                )
+                if self._is_module_stateful(mod):
+                    self.stateful_modules.add(sanitize_name(name))
 
     def get_input_nodes(self) -> List[Node]:
         # NOTE: This is a hack. Should use the input nodes from NIR graph
@@ -111,7 +125,7 @@ def _apply_module(
             assert "lif" in node.name, "this shouldnt happen.."
             new_state.state[node.name] = out  # snnTorch requires output inside state
             out = out[0]
-        elif self.stateful_modules[node.name]:
+        elif node.name in self.stateful_modules:
             new_state.state[node.name] = out[1:]  # Store the new state
             out = out[0]
         return out, new_state
@@ -141,7 +155,10 @@ def forward(
         # If the output node is a dummy nir.Output node, use the second-to-last node
         if node.name not in new_state.cache:
             node = self.execution_order[-2]
-        return new_state.cache[node.name], new_state
+        if self.return_state:
+            return new_state.cache[node.name], new_state
+        else:
+            return new_state.cache[node.name]
 
 
 def _mod_nir_to_graph(
diff --git a/tests/test_bidirectional.py b/tests/test_bidirectional.py
index 5345935..42e3df6 100644
--- a/tests/test_bidirectional.py
+++ b/tests/test_bidirectional.py
@@ -83,8 +83,10 @@ def test_nir_to_torch_to_nir(from_file=True):
     module = nirtorch.load(graph, _nir_to_torch_module)
     assert module is not None
     graph2 = nirtorch.extract_nir_graph(module, _torch_to_nir, torch.zeros(1, 1))
-    assert sorted(graph.edges) == sorted(graph2.edges)
-    assert graph2 is not None
+    edges1 = sorted(graph.edges)
+    edges2 = sorted(graph2.edges)
+    for e1, e2 in zip(edges1, edges2):
+        assert e1 == e2
 
 
 # if __name__ == '__main__':
diff --git a/tests/test_from_nir.py b/tests/test_from_nir.py
index 121d352..5e5942a 100644
--- a/tests/test_from_nir.py
+++ b/tests/test_from_nir.py
@@ -123,11 +123,16 @@ def _map_stateful(node):
         edges=[("i", "li"), ("li", "li2")],
     )  # Mock node
     m = load(g, _map_stateful)
-    out, state = m(torch.ones(10))
+    out = m(torch.ones(10))
+    assert isinstance(out, tuple)
+    out, state = out
     assert torch.allclose(out, torch.ones(10) * 3)
     assert state.state["li"] == (1,)
     assert state.state["li"] == (1,)
 
+    # Test that the model can return zero state
+    m.return_state = False
+    assert not isinstance(m(torch.ones(10)), tuple)
 
 def test_execute_recurrent():
     w = np.ones((1, 1))

From 8278437ae6b80fd3e017df2dca9b46d4565388b0 Mon Sep 17 00:00:00 2001
From: "Jens E. Pedersen" <jensegholm@protonmail.com>
Date: Sat, 21 Oct 2023 22:29:58 +0200
Subject: [PATCH 24/27] Added flag in nirtorch parsing

---
 nirtorch/from_nir.py | 20 ++++++++++++++++++--
 1 file changed, 18 insertions(+), 2 deletions(-)

diff --git a/nirtorch/from_nir.py b/nirtorch/from_nir.py
index d2324c0..419a264 100644
--- a/nirtorch/from_nir.py
+++ b/nirtorch/from_nir.py
@@ -196,15 +196,31 @@ def _switch_models_with_map(
 
 
 def load(
-    nir_graph: Union[nir.NIRNode, str], model_map: Callable[[nir.NIRNode], nn.Module]
+    nir_graph: Union[nir.NIRNode, str],
+    model_map: Callable[[nir.NIRNode], nn.Module],
+    return_state: bool = True,
 ) -> nn.Module:
     """Load a NIR graph and convert it to a torch module using the given model map.
 
+    Because the graph can contain recurrence and stateful modules, the execution accepts
+    a secondary state argument and returns a tuple of [output, state], instead of just the output as follows
+
+    >>> executor = nirtorch.load(nir_graph, model_map)
+    >>> old_state = None
+    >>> output, state = executor(input, old_state) # Notice the second argument and output
+    >>> output, state = executor(input, state) # This can go on for many (time)steps
+
+    If you do not wish to operate with state, set `return_state=False`.
+
     Args:
         nir_graph (Union[nir.NIRNode, str]): The NIR object to load, or a string representing
             the path to the NIR object.
         model_map (Callable[[nn.NIRNode], nn.Module]): A method that returns the a torch
             module that corresponds to each NIR node.
+        return_state (bool): If True, the execution of the loaded graph will return a tuple
+            of [output, state], where state is a GraphExecutorState object. If False, only
+            the NIR graph output will be returned. Note that state is required for recurrence
+            to work in the graphs.
 
     Returns:
         nn.Module: The generated torch module
@@ -216,4 +232,4 @@ def load(
     # Build a nirtorch.Graph based on the nir_graph
     graph = _mod_nir_to_graph(nir_module_graph, nir_nodes=nir_graph.nodes)
     # Build and return a graph executor module
-    return GraphExecutor(graph)
+    return GraphExecutor(graph, return_state=return_state)

From ec8cded0027f991796eb2f8d23f79107603a3902 Mon Sep 17 00:00:00 2001
From: "Jens E. Pedersen" <jensegholm@protonmail.com>
Date: Sat, 21 Oct 2023 22:53:40 +0200
Subject: [PATCH 25/27] Added flag in nirtorch parsing

---
 tests/test_from_nir.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/test_from_nir.py b/tests/test_from_nir.py
index 5e5942a..5269eb8 100644
--- a/tests/test_from_nir.py
+++ b/tests/test_from_nir.py
@@ -130,8 +130,8 @@ def _map_stateful(node):
     assert state.state["li"] == (1,)
     assert state.state["li"] == (1,)
 
-    # Test that the model can return zero state
-    m.return_state = False
+    # Test that the model can avoid returning state
+    m = load(g, _map_stateful, return_state=False)
     assert not isinstance(m(torch.ones(10)), tuple)
 
 def test_execute_recurrent():

From 0325c80a2b591e92ce911376a91c5b189005a70f Mon Sep 17 00:00:00 2001
From: Sadique Sheik <sadique.sheik@synsense.ai>
Date: Tue, 5 Dec 2023 15:25:22 +0100
Subject: [PATCH 26/27] minor changes to the doc strings

---
 nirtorch/from_nir.py | 31 +++++++++++++++++--------------
 1 file changed, 17 insertions(+), 14 deletions(-)

diff --git a/nirtorch/from_nir.py b/nirtorch/from_nir.py
index 419a264..c8ec7d9 100644
--- a/nirtorch/from_nir.py
+++ b/nirtorch/from_nir.py
@@ -24,13 +24,15 @@ class GraphExecutorState:
 class GraphExecutor(nn.Module):
     """Executes the NIR graph in PyTorch.
 
-    By default the graph executor is stateful, since there may be recurrence or stateful modules in the graph.
-    Specifically, that means accepting and returning a state object (`GraphExecutorState`).
-    If that is not desired, set `return_state=False` in the constructor.
+    By default the graph executor is stateful, since there may be recurrence or 
+    stateful modules in the graph. Specifically, that means accepting and returning a 
+    state object (`GraphExecutorState`). If that is not desired, 
+    set `return_state=False` in the constructor.
 
     Arguments:
         graph (Graph): The graph to execute
-        return_state (bool, optional): Whether to return the state object. Defaults to True.
+        return_state (bool, optional): Whether to return the state object. 
+            Defaults to True.
 
     Raises:
         ValueError: If there are no edges in the graph
@@ -169,9 +171,9 @@ def _mod_nir_to_graph(
     graph = Graph(module_names=module_names, inputs=inputs)
     for src, dst in torch_graph.edges:
         # Allow edges to refer to subgraph inputs and outputs
-        if not src in torch_graph.nodes and f"{src}.output" in torch_graph.nodes:
+        if src not in torch_graph.nodes and f"{src}.output" in torch_graph.nodes:
             src = f"{src}.output"
-        if not dst in torch_graph.nodes and f"{dst}.input" in torch_graph.nodes:
+        if dst not in torch_graph.nodes and f"{dst}.input" in torch_graph.nodes:
             dst = f"{dst}.input"
         graph.add_edge(torch_graph.nodes[src], torch_graph.nodes[dst])
     return graph
@@ -203,24 +205,25 @@ def load(
     """Load a NIR graph and convert it to a torch module using the given model map.
 
     Because the graph can contain recurrence and stateful modules, the execution accepts
-    a secondary state argument and returns a tuple of [output, state], instead of just the output as follows
+    a secondary state argument and returns a tuple of [output, state], instead of just 
+    the output as follows
 
     >>> executor = nirtorch.load(nir_graph, model_map)
     >>> old_state = None
-    >>> output, state = executor(input, old_state) # Notice the second argument and output
+    >>> output, state = executor(input, old_state) # Notice second argument and output
     >>> output, state = executor(input, state) # This can go on for many (time)steps
 
     If you do not wish to operate with state, set `return_state=False`.
 
     Args:
-        nir_graph (Union[nir.NIRNode, str]): The NIR object to load, or a string representing
-            the path to the NIR object.
+        nir_graph (Union[nir.NIRNode, str]): The NIR object to load, or a string 
+            representing the path to the NIR object.
         model_map (Callable[[nn.NIRNode], nn.Module]): A method that returns the a torch
             module that corresponds to each NIR node.
-        return_state (bool): If True, the execution of the loaded graph will return a tuple
-            of [output, state], where state is a GraphExecutorState object. If False, only
-            the NIR graph output will be returned. Note that state is required for recurrence
-            to work in the graphs.
+        return_state (bool): If True, the execution of the loaded graph will return a 
+            tuple of [output, state], where state is a GraphExecutorState object. 
+            If False, only the NIR graph output will be returned. Note that state is 
+            required for recurrence to work in the graphs.
 
     Returns:
         nn.Module: The generated torch module

From 53109c328a564a9dbf9b7deb538edf8ca4e00cff Mon Sep 17 00:00:00 2001
From: Sadique Sheik <sadique.sheik@synsense.ai>
Date: Tue, 5 Dec 2023 17:05:24 +0100
Subject: [PATCH 27/27] formatting fixes

---
 nirtorch/__init__.py        |  2 +-
 nirtorch/from_nir.py        | 25 +++++++++++++------------
 nirtorch/graph_utils.py     |  8 ++++----
 nirtorch/to_nir.py          |  9 +++++----
 tests/test_bidirectional.py |  2 +-
 tests/test_conversion.py    |  2 +-
 tests/test_from_nir.py      |  8 ++++++--
 tests/test_graph.py         |  8 ++++++--
 tests/test_graph_utils.py   |  2 +-
 tests/test_to_nir.py        | 12 +++++++++---
 10 files changed, 47 insertions(+), 31 deletions(-)

diff --git a/nirtorch/__init__.py b/nirtorch/__init__.py
index e4ae538..87eff9d 100644
--- a/nirtorch/__init__.py
+++ b/nirtorch/__init__.py
@@ -1,5 +1,5 @@
-from .graph import extract_torch_graph  # noqa F401
 from .from_nir import load  # noqa F401
+from .graph import extract_torch_graph  # noqa F401
 from .to_nir import extract_nir_graph  # noqa F401
 
 __version__ = version = "0.2.1"
diff --git a/nirtorch/from_nir.py b/nirtorch/from_nir.py
index c8ec7d9..77a5fb6 100644
--- a/nirtorch/from_nir.py
+++ b/nirtorch/from_nir.py
@@ -1,6 +1,6 @@
 import dataclasses
 import inspect
-from typing import Callable, Dict, List, Optional, Any, Union
+from typing import Any, Callable, Dict, List, Optional, Union
 
 import nir
 import torch
@@ -13,8 +13,8 @@
 
 @dataclasses.dataclass
 class GraphExecutorState:
-    """State for the GraphExecutor that keeps track of both the state of hidden
-    units and caches the output of previous modules, for use in (future) recurrent
+    """State for the GraphExecutor that keeps track of both the state of hidden units
+    and caches the output of previous modules, for use in (future) recurrent
     computations."""
 
     state: Dict[str, Any] = dataclasses.field(default_factory=dict)
@@ -24,14 +24,14 @@ class GraphExecutorState:
 class GraphExecutor(nn.Module):
     """Executes the NIR graph in PyTorch.
 
-    By default the graph executor is stateful, since there may be recurrence or 
-    stateful modules in the graph. Specifically, that means accepting and returning a 
-    state object (`GraphExecutorState`). If that is not desired, 
+    By default the graph executor is stateful, since there may be recurrence or
+    stateful modules in the graph. Specifically, that means accepting and returning a
+    state object (`GraphExecutorState`). If that is not desired,
     set `return_state=False` in the constructor.
 
     Arguments:
         graph (Graph): The graph to execute
-        return_state (bool, optional): Whether to return the state object. 
+        return_state (bool, optional): Whether to return the state object.
             Defaults to True.
 
     Raises:
@@ -92,6 +92,7 @@ def _apply_module(
         data: Optional[torch.Tensor] = None,
     ):
         """Applies a module and keeps track of its state.
+
         TODO: Use pytree to recursively construct the state
         """
         inputs = []
@@ -205,7 +206,7 @@ def load(
     """Load a NIR graph and convert it to a torch module using the given model map.
 
     Because the graph can contain recurrence and stateful modules, the execution accepts
-    a secondary state argument and returns a tuple of [output, state], instead of just 
+    a secondary state argument and returns a tuple of [output, state], instead of just
     the output as follows
 
     >>> executor = nirtorch.load(nir_graph, model_map)
@@ -216,13 +217,13 @@ def load(
     If you do not wish to operate with state, set `return_state=False`.
 
     Args:
-        nir_graph (Union[nir.NIRNode, str]): The NIR object to load, or a string 
+        nir_graph (Union[nir.NIRNode, str]): The NIR object to load, or a string
             representing the path to the NIR object.
         model_map (Callable[[nn.NIRNode], nn.Module]): A method that returns the a torch
             module that corresponds to each NIR node.
-        return_state (bool): If True, the execution of the loaded graph will return a 
-            tuple of [output, state], where state is a GraphExecutorState object. 
-            If False, only the NIR graph output will be returned. Note that state is 
+        return_state (bool): If True, the execution of the loaded graph will return a
+            tuple of [output, state], where state is a GraphExecutorState object.
+            If False, only the NIR graph output will be returned. Note that state is
             required for recurrence to work in the graphs.
 
     Returns:
diff --git a/nirtorch/graph_utils.py b/nirtorch/graph_utils.py
index 1f67d71..d807d49 100644
--- a/nirtorch/graph_utils.py
+++ b/nirtorch/graph_utils.py
@@ -65,12 +65,12 @@ def find_all_ancestors(
 #     return execution_order
 #
 
+
 def trace_execution(
     node: T, edge_fn: Callable[[T], List[T]], visited: Set[T] = None
 ) -> List[T]:
-    """Traces the execution of a node by listing them in order, coloring recursive
-    nodes to avoid adding the same node twice.
-    """
+    """Traces the execution of a node by listing them in order, coloring recursive nodes
+    to avoid adding the same node twice."""
     if visited is None:
         visited = set()
 
@@ -83,4 +83,4 @@ def trace_execution(
     for child in edge_fn(node):
         if child not in visited:
             successors += trace_execution(child, edge_fn, visited)
-    return [node] + successors
\ No newline at end of file
+    return [node] + successors
diff --git a/nirtorch/to_nir.py b/nirtorch/to_nir.py
index 47e7d03..a819c21 100644
--- a/nirtorch/to_nir.py
+++ b/nirtorch/to_nir.py
@@ -1,5 +1,5 @@
-from typing import Any, Callable, Optional, Sequence
 import logging
+from typing import Any, Callable, Optional, Sequence
 
 import nir
 import numpy as np
@@ -15,13 +15,14 @@ def extract_nir_graph(
     model_name: Optional[str] = "model",
     ignore_submodules_of=None,
     model_fwd_args=[],
-    ignore_dims: Optional[Sequence[int]]=None,
+    ignore_dims: Optional[Sequence[int]] = None,
 ) -> nir.NIRNode:
     """Given a `model`, generate an NIR representation using the specified `model_map`.
 
     Assumptions and known issues:
-        - Cannot deal with layers like torch.nn.Identity(), since the input tensor and output
-            tensor will be the same object, and therefore lead to cyclic connections.
+        - Cannot deal with layers like torch.nn.Identity(), since the input tensor and
+          output tensor will be the same object, and therefore lead to cyclic
+          connections.
 
     Args:
         model (nn.Module): The model of interest
diff --git a/tests/test_bidirectional.py b/tests/test_bidirectional.py
index 42e3df6..52f8f6d 100644
--- a/tests/test_bidirectional.py
+++ b/tests/test_bidirectional.py
@@ -1,8 +1,8 @@
 import nir
 import numpy as np
 import torch
-import nirtorch
 
+import nirtorch
 
 use_snntorch = False
 # use_snntorch = True
diff --git a/tests/test_conversion.py b/tests/test_conversion.py
index c04be4f..ee48b5d 100644
--- a/tests/test_conversion.py
+++ b/tests/test_conversion.py
@@ -1,7 +1,7 @@
+import nir
 import torch
 import torch.nn as nn
 
-import nir
 import nirtorch
 
 
diff --git a/tests/test_from_nir.py b/tests/test_from_nir.py
index d2680b9..b2907f4 100644
--- a/tests/test_from_nir.py
+++ b/tests/test_from_nir.py
@@ -1,7 +1,7 @@
 import nir
 import numpy as np
-import torch
 import pytest
+import torch
 
 from nirtorch.from_nir import load
 
@@ -56,7 +56,10 @@ def test_extract_empty():
 
 
 def test_extract_illegal_name():
-    graph = nir.NIRGraph({"a.b": nir.Input(np.ones(1)), "a.c": nir.Linear(np.array([[1.]]))}, [("a.b", "a.c")])
+    graph = nir.NIRGraph(
+        {"a.b": nir.Input(np.ones(1)), "a.c": nir.Linear(np.array([[1.0]]))},
+        [("a.b", "a.c")],
+    )
     torch_graph = load(graph, _torch_model_map)
     assert "a_c" in torch_graph._modules
 
@@ -131,6 +134,7 @@ def _map_stateful(node):
     m = load(g, _map_stateful, return_state=False)
     assert not isinstance(m(torch.ones(10)), tuple)
 
+
 def test_execute_recurrent():
     w = np.ones((1, 1))
     g = nir.NIRGraph(
diff --git a/tests/test_graph.py b/tests/test_graph.py
index 20d17f5..6e10625 100644
--- a/tests/test_graph.py
+++ b/tests/test_graph.py
@@ -1,10 +1,10 @@
+import nir
 import pytest
 import torch
 import torch.nn as nn
 from norse.torch import LIBoxCell, LIFCell, SequentialState
 from sinabs.layers import Merge
 
-import nir
 from nirtorch import extract_nir_graph, extract_torch_graph
 
 
@@ -238,6 +238,7 @@ def test_root_has_no_source():
         len(graph.find_source_nodes_of(graph.find_node(my_branched_model.relu1))) == 0
     )
 
+
 @pytest.mark.skip(reason="Root tracing is broken")
 def test_get_root():
     graph = extract_torch_graph(my_branched_model, sample_data=data, model_name=None)
@@ -282,7 +283,9 @@ def test_sequential_flatten():
     assert tuple(g.nodes["input"].input_type["input"]) == (3, 4)
 
     d = torch.empty(2, 3, 4)
-    g = extract_nir_graph(torch.nn.Flatten(1), lambda x: nir.Flatten(d.shape, 1), d, ignore_dims=[0])
+    g = extract_nir_graph(
+        torch.nn.Flatten(1), lambda x: nir.Flatten(d.shape, 1), d, ignore_dims=[0]
+    )
     assert tuple(g.nodes["input"].input_type["input"]) == (3, 4)
 
 
@@ -311,6 +314,7 @@ def forward(self, x, state=None):
     assert d.nodes.keys() == {"input", "l", "r", "output"}
     assert set(d.edges) == {("input", "r"), ("r", "l"), ("l", "output"), ("r", "r")}
 
+
 @pytest.mark.skip(reason="Subgraphs are currently flattened")
 def test_captures_recurrence_manually():
     def export_affine_rec_gru(module):
diff --git a/tests/test_graph_utils.py b/tests/test_graph_utils.py
index 36ec5f3..95b51b9 100644
--- a/tests/test_graph_utils.py
+++ b/tests/test_graph_utils.py
@@ -21,7 +21,7 @@ def from_string(graph):
 
     def __hash__(self) -> int:
         return self.name.__hash__()
-    
+
     def __eq__(self, other: object) -> bool:
         return self.name == other.name
 
diff --git a/tests/test_to_nir.py b/tests/test_to_nir.py
index 3ea108d..d35d217 100644
--- a/tests/test_to_nir.py
+++ b/tests/test_to_nir.py
@@ -115,7 +115,9 @@ def extractor(module: nn.Module):
         return nir.Affine(module.weight, module.bias)
 
     raw_input_shape = (1, 3)
-    g = extract_nir_graph(model, extractor, torch.ones(raw_input_shape), ignore_dims=[0])
+    g = extract_nir_graph(
+        model, extractor, torch.ones(raw_input_shape), ignore_dims=[0]
+    )
     exp_input_shape = (3,)
     assert np.alltrue(g.nodes["input"].input_type["input"] == np.array(exp_input_shape))
     assert g.nodes["model"].weight.shape == (1, 3)
@@ -129,13 +131,17 @@ def extractor(module: nn.Module):
         return nir.Affine(module.weight, module.bias)
 
     raw_input_shape = (1, 10, 3)
-    g = extract_nir_graph(model, extractor, torch.ones(raw_input_shape), ignore_dims=[0, -2])
+    g = extract_nir_graph(
+        model, extractor, torch.ones(raw_input_shape), ignore_dims=[0, -2]
+    )
     exp_input_shape = (3,)
     assert np.alltrue(g.nodes["input"].input_type["input"] == np.array(exp_input_shape))
     assert g.nodes["model"].weight.shape == (1, 3)
 
     raw_input_shape = (1, 10, 3)
-    g = extract_nir_graph(model, extractor, torch.ones(raw_input_shape), ignore_dims=[0, 1])
+    g = extract_nir_graph(
+        model, extractor, torch.ones(raw_input_shape), ignore_dims=[0, 1]
+    )
     exp_input_shape = (3,)
     assert np.alltrue(g.nodes["input"].input_type["input"] == np.array(exp_input_shape))