`NIR` export of `LIF` node might be out of date, `v_reset` parameter is missing

[jeromehue]

Trying to run the following code snippet

import snntorch as snn
import torch
import nir

from snntorch.export_nir import export_to_nir

lif1 = snn.Leaky(beta=0.9, init_hidden=True)
lif2 = snn.Leaky(beta=0.9, init_hidden=True, output=True)

# Create a network
snntorch_network = torch.nn.Sequential(
    torch.nn.Flatten(),
    torch.nn.Linear(784, 500),
    lif1,
    torch.nn.Linear(500, 10),
    lif2
)

sample_data = torch.randn(1, 784)
nir_model = export_to_nir(snntorch_network, sample_data)
nir.write("nir_model.nir", nir_model)

with nir 1.0.6, nirtorch 2.0.4 and snntorch built from sources (commit a2ce8cd) results in the following error :

Traceback (most recent call last):
  File "/home/Tasks/task-nir/true.py", line 20, in <module>
    nir_model = export_to_nir(snntorch_network, sample_data)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/nobackup/documents/jenv/lib/python3.12/site-packages/snntorch/export_nir.py", line 244, in export_to_nir
    nir_graph = nirtorch.extract_nir_graph(
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/nobackup/documents/jenv/lib/python3.12/site-packages/nirtorch/to_nir.py", line 79, in extract_nir_graph
    mapped_node = model_map(node.elem)
                  ^^^^^^^^^^^^^^^^^^^^
  File "/nobackup/documents/jenv/lib/python3.12/site-packages/snntorch/export_nir.py", line 77, in _extract_snntorch_module
    return nir.LIF(
           ^^^^^^^^
TypeError: LIF.__init__() missing 1 required positional argument: 'v_reset'

It seems indeed that the function export_nir does not take into account the addition of the v_reset parameter in nir (added in commit neuromorphs/NIR@6189cac).

[jeromehue]

@jeshraghian this is surely related to #382

[fabio-innatera]

I have a branch that fixes this issue: https://github.com/fabio-innatera/snntorch/tree/nir-missing-v_reset

However, this will break support for NIR version 1.0.5 and older, so I am not sure if we want to merge this right away,
especially because NIR 1.0.6 might break snnTorch NIR support in other ways.

@jeromehue Downgrading NIR to 1.0.5 or maybe 1.0.4 should temporarily fix your problem.

Some more notes on NIR 1.0.6 support in neuromorphs/NIR#157.

[jeshraghian]

Sorry this took me a while to look at. I haven't kept up with the latest changes in NIR - the refactor in NIR seemed to kill a lot of backwards compatibility.

I'd be happy to let go of older NIR support if @fabio-innatera you were to make a PR? I commented out most of the tests as they're probably outdated now too.

[jeshraghian]

I went ahead and merged a PR from @fabio-innatera 's branch and almost everything looks good now!

There was the recent addition of the RLeaky neuron as well, and I added the v_reset parameter to that although now the RLeaky tests are failing.

[fabio-innatera]

#389 fixes the RLeaky related NIR tests and #390 bumps the NIR version in documentation and tests.

[fabio-innatera]

One issue I discovered recently is that snntorch exports NIR graphs that contained scalar parameters for neurons, like in the example network above:

lif1 = snn.Leaky(beta=0.9, init_hidden=True)

As per the type annotation here https://github.com/neuromorphs/NIR/blob/3109a90ebe3e1e44182d71fd5609b8c9e50e5a33/nir/ir/neuron.py#L30 this is not valid in NIR.

In the discussion with @Jegp here neuromorphs/NIR#157 we decided that this should stay this way, so we can guarantee a well formed NIR graph.

Newer versions of NIR reject nodes with scalar parameters due to the new type inference and checking (right now with an ugly error message, PR to fix that here neuromorphs/NIR#162), and I think this is the main reason for some backwards incompatibility issues like in examples/tutorial_snntorch_to_nir.ipynb.

TL;DR: Older versions of snntorch exported technically invalid NIR graphs, that can no longer be loaded with newer versions of NIR.

@jeromehue To make your network exportable to NIR, change your Leaky neurons to something like this:

lif1 = snn.Leaky(beta=0.9 * torch.ones(500), threshold=torch.ones(500), init_hidden=True)
lif2 = snn.Leaky(beta=0.9 * torch.ones(10), threshold=torch.ones(10), init_hidden=True, output=True)

[jeromehue]

With snntorch built from source (commit bd56874) , nir version 1.0.6 and nirtorch version 2.0.5 (following requirements from a20deca), there are still some errors :

import snntorch as snn
import torch
from snntorch.export_nir import export_to_nir
from snntorch.import_nir import import_from_nir

lif1 = snn.Leaky(beta=0.9 * torch.ones(500), threshold=torch.ones(500), init_hidden=True)
lif2 = snn.Leaky(beta=0.9 * torch.ones(10), threshold=torch.ones(10), init_hidden=True, output=True)

net = torch.nn.Sequential(
    torch.nn.Flatten(),
    torch.nn.Linear(784, 500),
    lif1,
    torch.nn.Linear(500, 10),
    lif2
)

sample_data = torch.randn(1, 784)
nir_graph = export_to_nir(net, sample_data, model_name="snntorch")

runs into :

WARNING:root:removing self-connection ('0', '0')
Traceback (most recent call last):
  File "/home/Experiments/snns_training_import_export/flatten.py", line 18, in <module>
    nir_graph = export_to_nir(net, sample_data, model_name="snntorch")
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/jenv/lib/python3.12/site-packages/snntorch/export_nir.py", line 299, in export_to_nir
    nir_graph.infer_types()
  File "/home/jenv/lib/python3.12/site-packages/nir/ir/graph.py", line 328, in infer_types
    raise ValueError(
ValueError: Type inference error: type mismatch: 4.output: [[10]] -> output.input: [[ 1 10]]

So at least the error related to nir is fixed, but there are still problems with flatten.
Looks like using ignore_dims=[0] fixes the problem.

nir_graph = export_to_nir(net, sample_data, model_name="snntorch", ignore_dims=[0])

So my question now, is why does the flatten export change the start dim ?

Even looking at this example tutorial, the imported network fails with a different batch size :

import snntorch as snn
import torch
from snntorch.export_nir import export_to_nir
from snntorch.import_nir import import_from_nir


lif1 = snn.Leaky(beta=0.9 * torch.ones(500), threshold=torch.ones(500), init_hidden=True)
lif2 = snn.Leaky(beta=0.9 * torch.ones(10), threshold=torch.ones(10), init_hidden=True, output=True)

net = torch.nn.Sequential(
    torch.nn.Flatten(),
    torch.nn.Linear(784, 500),
    lif1,
    torch.nn.Linear(500, 10),
    lif2
)

sample_data = torch.randn(1, 784)
nir_graph = export_to_nir(net, sample_data, model_name="snntorch", ignore_dims=[0])

net2 = import_from_nir(nir_graph)
sample_data2 = torch.randn(16, 784)
out = net2(sample_data2)
print(out)

WARNING:root:removing self-connection ('0', '0')
replace rnn subgraph with nirgraph
Traceback (most recent call last):
  File "/home/documents/NIR_TORCH/tuto.py", line 23, in <module>
    out = net2(sample_data2)
          ^^^^^^^^^^^^^^^^^^
  File "/home/jenv/lib/python3.12/site-packages/torch/fx/graph_module.py", line 830, in call_wrapped
    return self._wrapped_call(self, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/jenv/lib/python3.12/site-packages/torch/fx/graph_module.py", line 406, in __call__
    raise e
  File "/home/jenv/lib/python3.12/site-packages/torch/fx/graph_module.py", line 393, in __call__
    return super(self.cls, obj).__call__(*args, **kwargs)  # type: ignore[misc]
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/jenv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1751, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/jenv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1762, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<eval_with_key>.0", line 8, in forward
  File "/home/jenv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1751, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/jenv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1762, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/jenv/lib/python3.12/site-packages/torch/nn/modules/linear.py", line 125, in forward
    return F.linear(input, self.weight, self.bias)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: mat1 and mat2 shapes cannot be multiplied (1x12544 and 784x500)

Obviously caused by the flatten start dim being set to zero during export.

[EvaReid12]

I try to use nir with snntorch and I end up with the same errors ? is there anything new on it ?