nirtorch 0.0.2

Neuromorphic Intermediate Representation

NIRTorch

PyTorch helpers for the Neuromorphic Intermediate Representation (NIR). This is a no frills python package to enable torch based libraries to translate to NIR.

Installation

pip install nirtorch

Usage

NIRTorch is typically only interfaced by library/hardwarae developers. NIRTorch provides the extract_nir_graph function that takes as input a torch.nn.Module and a means to map Torch modules into NIR nodes. An NIR node is an element in the NIR compute graph, corresponding to neuromorphic ODEs.

Here is an example from the Norse library:

def _extract_norse_module(module: torch.nn.Module) -> Optional[nir.NIRNode]:
    if isinstance(module, LIFBoxCell):
        return nir.LIF(
            tau=module.p.tau_mem_inv,
            v_th=module.p.v_th,
            v_leak=module.p.v_leak,
            r=torch.ones_like(module.p.v_leak),
        )
    elif isinstance(module, torch.nn.Linear):
        return nir.Linear(module.weight, module.bias)
    elif ...

    return None

def to_nir(
    module: torch.nn.Module, sample_data: torch.Tensor, model_name: str = "norse"
) -> nir.NIRNode:
    return extract_nir_graph(
        module, _extract_norse_module, sample_data, model_name=model_name
    )

Acknowledgements

If you use NIR torch in your work, please cite the following Zenodo reference

@software{nir2023,
  author       = {Abreu, Steven and
                  Bauer, Felix and
                  Eshraghian, Jason and
                  Jobst, Matthias and
                  Lenz, Gregor and
                  Pedersen, Jens Egholm and
                  Sheik, Sadique},
  title        = {Neuromorphic Intermediate Representation},
  month        = jul,
  year         = 2023,
  publisher    = {Zenodo},
  version      = {0.0.1},
  doi          = {10.5281/zenodo.8105042},
  url          = {https://doi.org/10.5281/zenodo.8105042}
}