Exportation to NIR format

[dylanperdigao]

The following issue comports two parts:

1. snntorch deployed on pypi/conda is outdated
2. problems with nodes dtype/shape/data

Current versions of the software:

python version 3.12
snntorch version: 0.9.1
nir version: 1.0.4

1) Outdated deployed version

As I mentioned this week in the OpenNeuromorphic Discord server, the actual deployed version of snntorch on pypi and conda is outdated, and version 0.9.1 does not support some layers for the neural networks:

model = nn.Sequential(

    nn.Conv2d(in_channels=2, out_channels=8, kernel_size=(3, 3), padding=(1, 1), bias=False),
    snntorch.Leaky(beta=0.1, threshold=1.0, spike_grad=sigmoid(25), init_hidden=True),
    nn.AvgPool2d(2, 2),

    nn.Conv2d(in_channels=8, out_channels=16, kernel_size=(3, 3), padding=(1, 1), bias=False),
    snntorch.Leaky(beta=0.1, threshold=1.0, spike_grad=sigmoid(25), init_hidden=True),
    nn.AvgPool2d(2, 2),

    nn.Conv2d(in_channels=16, out_channels=16, kernel_size=(3, 3), padding=(1, 1), stride=(2, 2),  bias=False),
    snntorch.Leaky(beta=0.1, threshold=1.0, spike_grad=sigmoid(25), init_hidden=True),

    nn.Flatten(),
    nn.Linear(16 * 4 * 4, 10, bias=False),
    snntorch.Leaky(beta=0.1, threshold=1.0, spike_grad=sigmoid(25), init_hidden=True, output=True)
)
sample_data = torch.rand((batch_size, 2, 34, 34))
nir_model = snntorch.export_to_nir(model, sample_data)

Will produce:

[WARNING] module not implemented: Conv2d
[WARNING] module not implemented: AvgPool2d
[WARNING] module not implemented: Conv2d
[WARNING] module not implemented: AvgPool2d
[WARNING] module not implemented: Conv2d
[WARNING] module not implemented: Flatten

2) Problems with NIR nodes

Then, if I use the current version of the code that is on GitHub, I'm able to run without the "not implemented" warning. However, if I try to save the model in NIR format:

nir.write("models/snntorch_model-csnn.nir", nir_model)

I get the following error:

TypeError: One of data, shape or dtype must be specified

A possible solution to mitigate this is to manually set the nodes:

# SET DTYPES AND INPUT SHAPE FOR EACH NODE
# conv2d
nir_model.nodes['0'].input_type = {'input': np.array([], dtype=np.float64)}
nir_model.nodes['0'].output_type = {'output': np.array([], dtype=np.float64)}
nir_model.nodes['0'].input_shape = (batch_size, 2, 34, 34)
# maxpool2d
nir_model.nodes['2'].input_type = {'input': np.array([], dtype=np.float64)}
nir_model.nodes['2'].output_type = {'output': np.array([], dtype=np.float64)}
# conv2d
nir_model.nodes['3'].input_type = {'input': np.array([], dtype=np.float64)}
nir_model.nodes['3'].output_type = {'output': np.array([], dtype=np.float64)}
nir_model.nodes['3'].input_shape = (batch_size, 8, 17, 17)
# maxpool2d
nir_model.nodes['5'].input_type = {'input': np.array([], dtype=np.float64)}
nir_model.nodes['5'].output_type = {'output': np.array([], dtype=np.float64)}
# conv2d
nir_model.nodes['6'].input_type = {'input': np.array([], dtype=np.float64)}
nir_model.nodes['6'].output_type = {'output': np.array([], dtype=np.float64)}
nir_model.nodes['6'].input_shape = (batch_size, 16, 9, 9)
# flatten
nir_model.nodes['8'].input_type = {'input': np.array([], dtype=np.float64)}
nir_model.nodes['8'].output_type = {'output': np.array([], dtype=np.float64)}

But is it supposed to do this like that?
It could be related to the previous point of the convolutional, pooling, and flattening layers that were not implemented. Is something missing in the implementation?