Add Full NIR Export/Import Support for RLeaky Layers #380
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Summary
This PR introduces comprehensive support for exporting and importing recurrent RLeaky neuron layers between snntorch and the Neuromorphic Intermediate Representation (NIR) format. With this change, models containing RLeaky cells can now be seamlessly round-tripped between snntorch and other NIR-compatible frameworks, ensuring correctness of parameters and numerical fidelity. The PR also updates the documentation and adds tests to guarantee functionality and stability.
Motivation / Previous Problem
Previously, the NIR export and import functionalities in snntorch only supported feedforward layers and a limited subset of recurrent cell types. Specifically:
This gap hindered the interoperability of advanced SNN models, particularly those using recurrent memory, across neuromorphic platforms and toolchains.
What’s Changed
Export Logic for RLeaky (snntorch/export_nir.py)
The exporter now detects snn.RLeaky modules and builds the correct NIR subgraph:
Import Logic for RLeaky (snntorch/import_nir.py)
The importer now parses NIR subgraphs that implement LIF-based recurrent motifs and reconstructs the corresponding RLeaky cell in snntorch:
Documentation Updates
Both export_nir and import_nir module docs have been updated:
Comprehensive Testing (tests/test_nir.py)
Adds fixtures for constructing test networks with RLeaky layers.
Adds tests for:
How Does This Help?
Backward Compatibility
Additional Notes
Please let me know if any additional documentation, examples, or edge-case tests would be helpful!