Skip to content

Add Full NIR Export/Import Support for RLeaky Layers #380

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: master
Choose a base branch
from

Conversation

KAVYANSHTYAGI
Copy link

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:

Export: Attempting to export a model with an RLeaky cell would raise a NotImplementedError, making it impossible to convert such models to NIR for deployment or cross-library research.

Import: Similarly, importing NIR graphs representing LIF-based recurrent subgraphs (used to encode RLeaky) was unsupported, with attempts failing early in the conversion process.

Documentation: The limitations were not clearly documented, which could confuse users about the current state of recurrent support in NIR workflows.

This gap hindered the interoperability of advanced SNN models, particularly those using recurrent memory, across neuromorphic platforms and toolchains.

What’s Changed

  1. Export Logic for RLeaky (snntorch/export_nir.py)

    The exporter now detects snn.RLeaky modules and builds the correct NIR subgraph:

     Supports both all-to-all and diagonal (vector) recurrent connections.
    
     Extracts recurrent weight matrices, neuron parameters (beta, threshold, etc.), and correctly encodes them in the NIR graph using nir.LIF and nir.Linear nodes.
    
     Ensures parameter shapes and connectivity are faithfully represented for round-trip equivalence.
    
  2. 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:

     Extracts all neuron and recurrent parameters.
    
     Handles both diagonal and dense cases.
    
     Performs assert checks for mathematical correctness and parameter compatibility.
    
     Restores all relevant attributes (weights, thresholds, biases, etc.).
    
  3. Documentation Updates

    Both export_nir and import_nir module docs have been updated:

     Explicitly state support for exporting and recreating networks with recurrent layers, including RLeaky and RSynaptic.
    
     Removes outdated warnings about unsupported features.
    
  4. Comprehensive Testing (tests/test_nir.py)

    Adds fixtures for constructing test networks with RLeaky layers.

    Adds tests for:

     Correct export of networks with recurrent cells.
    
     Accurate graph structure in exported NIR.
    
     Round-trip equivalence: numerical outputs of the network remain unchanged after export/import.
    
     Edge cases (all-to-all and diagonal recurrent matrices).
    

How Does This Help?

Unlocks Full Interoperability: Enables users to export and import models with RLeaky cells for simulation, deployment, or research across all NIR-compatible neuromorphic frameworks.

Ensures Research Reproducibility: Researchers can now share, publish, or reproduce SNN models with recurrent memory reliably.

Facilitates Hardware Deployment: Models built in snntorch can be deployed to neuromorphic hardware or other simulators via NIR without modification.

Promotes Community Adoption: Lowers the barrier for collaborative development and tool integration in the neuromorphic and SNN community.

Improved User Experience: Users can now confidently build complex models, knowing that export/import will just work, with clear documentation and robust tests as guarantees.

Backward Compatibility

No breaking changes for users who do not use RLeaky.

Existing export/import for other layers remains unchanged.

Additional Notes

This PR lays the foundation for further recurrent cell support in the future. Any new cell types can follow the pattern established here.

Comprehensive tests ensure that future refactors will not silently break recurrent export/import.

Please let me know if any additional documentation, examples, or edge-case tests would be helpful!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

1 participant