<picture>
<source media="(prefers-color-scheme: dark)" srcset="https://github.com/neuromorphs/NIR/raw/main/docs/logo_dark.png">
<img alt="NIR Logo" src="https://github.com/neuromorphs/NIR/raw/main/docs/logo_light.png">
</picture>
# NIRTorch - Torch for the Neuromorphic Intermediate Representation
[](https://doi.org/10.1038/s41467-024-52259-9)
[](https://pypi.org/project/nirtorch/)
[](https://github.com/neuromorphs/NIRTorch/releases)
[](https://discord.gg/JRMRGP9h3c)
PyTorch helpers for the [Neuromorphic Intermediate Representation (NIR)](https://github.com/neuromorphs/nir).
This is a no frills python package to enable torch based libraries to translate to NIR.
## Installation
```shell
pip install nirtorch
```
## Usage
> [!TIP]
> Read the full documentation at [neuroir.org/docs](https://neuroir.org/docs/dev_pytorch.html).
NIRTorch is typically only interfaced by library/hardwarae developers.
NIRTorch provides two functions `nir_to_torch` and `torch_to_nir` that maps NIR graphs to Torch and vice versa.
### Translating from NIR to Torch
NIRTorch automatically creates the Torch graph and module behind the scenes.
The only thing you have to do as a developer is to provide a mapping between NIR nodes to PyTorch modules.
That is, a dictionary of nodes with associated functions that produces Torch modules, given a NIR Node (`Dict[nir.NIRNode, Callable[[nir.NIRNode], torch.nn.Module]]`), where a NIR node is an element in the NIR compute graph, corresponding to neuromorphic ODEs.
Here is one example taken from the [Norse](https://github.com/norse/norse) library:
```python
# Define NIR -> Norse mapping
my_node_dictionary = {}
def _map_linear(linear: nir.Linear) -> torch.nn.Module:
output_shape, input_shape = linear.weight.shape[-2:]
l = torch.nn.Linear(input_shape, output_shape, bias=False)
l.weight.data = linear.weight
return l
_my_node_dictionary[nir.Linear] = _map_linear
... # And so on for other modules
# Declare a user-facing function to output Norse modules
def from_nir(module: nir.NIRGraph) -> torch.nn.Module:
return nirtorch.nir_to_torch(
module, # The NIR Graph to map
_my_node_dictionary # The dictionary used to map modules to Norse
)
```
### Translating from Torch to NIR
NIRTorch helps construct a NIR graph by tracing through the PyTorch graph module by module.
Our job as a developer is now to provide a mapping that helps translate PyTorch modules into NIR nodes, where a NIR node is an element in the NIR compute graph.
That is, a `Dict[torch.nn.Module, Callable[[torch.nn.Module], nir.NIRNode]]`.
Note that the output node may be a subgraph.
def torch_to_nir(
module: torch.nn.Module,
module_map: Dict[torch.nn.Module, Callable[[torch.nn.Module], nir.NIRNode]],
default_dict: Optional[
Dict[torch.nn.Module, Callable[[torch.nn.Module], nir.NIRNode]]
] = None,
) -> nir.NIRGraph:
```python
# Define Norse -> NIR mapping
_my_module_dict = {}
def _extract_lif_module(module: norse.LIFBoxCell) -> Optional[nir.NIRNode]:
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),
)
_my_module_dict[norse.LIFBoxcell] =_extract_lif_module
# Declare a user-facing function to output NIR graphs
def to_nir(module: torch.nn.Module) -> nir.NIRNode:
return nirtorch.torch_to_nir(
module, # The Norse/Torch module to parse
_my_module_dict # The dictionary that maps Norse module to NIR
)
```
Read the full documentation at [neuroir.org/docs](https://neuroir.org/docs/dev_pytorch.html).
## Acknowledgements
If you use NIRTorch in your work, please cite the [our work in Nature Communications](https://www.nature.com/articles/s41467-024-52259-9)
```
article{NIR2024,
title={Neuromorphic intermediate representation: A unified instruction set for interoperable brain-inspired computing},
author={Pedersen, Jens E. and Abreu, Steven and Jobst, Matthias and Lenz, Gregor and Fra, Vittorio and Bauer, Felix Christian and Muir, Dylan Richard and Zhou, Peng and Vogginger, Bernhard and Heckel, Kade and Urgese, Gianvito and Shankar, Sadasivan and Stewart, Terrence C. and Sheik, Sadique and Eshraghian, Jason K.},
rights={2024 The Author(s)},
DOI={10.1038/s41467-024-52259-9},
number={1},
journal={Nature Communications},
volume={15},
year={2024},
month=sep,
pages={8122},
}
```
## For developers
If you want to make sure that your code is linted correctly on your local machine, use [pre-commit](https://pre-commit.com/) to automatically perform checks before every git commit. To use it, first install the package in your environment
```
pip install pre-commit
```
and then install the pre-commit hooks that are listed in the root of this repository
```
pre-commit install
```
Next time you commit some changes, all the checks will be run!
Raw data
{
"_id": null,
"home_page": null,
"name": "nirtorch",
"maintainer": null,
"docs_url": null,
"requires_python": null,
"maintainer_email": null,
"keywords": "neuromorphic, intermediate, representation",
"author": null,
"author_email": "Steven Abreu <s.abreu@rug.nl>, Felix Bauer <felix.bauer@synsen.ai>, Jason Eshraghian <jeshragh@ucsc.edu>, Matthias Jobst <matthias.jobst2@tu-dresden.de>, Gregor Lenz <mail@lenzgregor.com>, Jens Egholm Pedersen <jens@jepedersen.dk>, Sadique Sheik <sadique.sheik@synsense.ai>",
"download_url": "https://files.pythonhosted.org/packages/c1/a9/b3f7d5fef556dea6208008ed7b67c720f62a44699e21f86f884537759a90/nirtorch-2.0.5.tar.gz",
"platform": null,
"description": "<picture>\n<source media=\"(prefers-color-scheme: dark)\" srcset=\"https://github.com/neuromorphs/NIR/raw/main/docs/logo_dark.png\">\n<img alt=\"NIR Logo\" src=\"https://github.com/neuromorphs/NIR/raw/main/docs/logo_light.png\">\n</picture>\n\n# NIRTorch - Torch for the Neuromorphic Intermediate Representation\n\n[](https://doi.org/10.1038/s41467-024-52259-9)\n[](https://pypi.org/project/nirtorch/)\n[](https://github.com/neuromorphs/NIRTorch/releases)\n[](https://discord.gg/JRMRGP9h3c)\n\n\nPyTorch helpers for the [Neuromorphic Intermediate Representation (NIR)](https://github.com/neuromorphs/nir).\nThis is a no frills python package to enable torch based libraries to translate to NIR.\n\n## Installation\n```shell\npip install nirtorch\n```\n\n## Usage\n\n> [!TIP]\n> Read the full documentation at [neuroir.org/docs](https://neuroir.org/docs/dev_pytorch.html).\n\n\nNIRTorch is typically only interfaced by library/hardwarae developers.\nNIRTorch provides two functions `nir_to_torch` and `torch_to_nir` that maps NIR graphs to Torch and vice versa.\n\n### Translating from NIR to Torch\nNIRTorch automatically creates the Torch graph and module behind the scenes.\nThe only thing you have to do as a developer is to provide a mapping between NIR nodes to PyTorch modules.\nThat is, a dictionary of nodes with associated functions that produces Torch modules, given a NIR Node (`Dict[nir.NIRNode, Callable[[nir.NIRNode], torch.nn.Module]]`), where a NIR node is an element in the NIR compute graph, corresponding to neuromorphic ODEs.\nHere is one example taken from the [Norse](https://github.com/norse/norse) library:\n\n```python\n# Define NIR -> Norse mapping\nmy_node_dictionary = {}\ndef _map_linear(linear: nir.Linear) -> torch.nn.Module:\n output_shape, input_shape = linear.weight.shape[-2:]\n l = torch.nn.Linear(input_shape, output_shape, bias=False)\n l.weight.data = linear.weight\n return l\n_my_node_dictionary[nir.Linear] = _map_linear\n... # And so on for other modules\n\n# Declare a user-facing function to output Norse modules\ndef from_nir(module: nir.NIRGraph) -> torch.nn.Module:\n return nirtorch.nir_to_torch(\n module, # The NIR Graph to map\n _my_node_dictionary # The dictionary used to map modules to Norse\n )\n```\n\n### Translating from Torch to NIR\nNIRTorch helps construct a NIR graph by tracing through the PyTorch graph module by module.\nOur job as a developer is now to provide a mapping that helps translate PyTorch modules into NIR nodes, where a NIR node is an element in the NIR compute graph.\nThat is, a `Dict[torch.nn.Module, Callable[[torch.nn.Module], nir.NIRNode]]`.\nNote that the output node may be a subgraph.\n\ndef torch_to_nir(\n module: torch.nn.Module,\n module_map: Dict[torch.nn.Module, Callable[[torch.nn.Module], nir.NIRNode]],\n default_dict: Optional[\n Dict[torch.nn.Module, Callable[[torch.nn.Module], nir.NIRNode]]\n ] = None,\n) -> nir.NIRGraph:\n\n```python\n# Define Norse -> NIR mapping\n_my_module_dict = {}\ndef _extract_lif_module(module: norse.LIFBoxCell) -> Optional[nir.NIRNode]:\n return nir.LIF(\n tau=module.p.tau_mem_inv,\n v_th=module.p.v_th,\n v_leak=module.p.v_leak,\n r=torch.ones_like(module.p.v_leak),\n )\n_my_module_dict[norse.LIFBoxcell] =_extract_lif_module\n\n# Declare a user-facing function to output NIR graphs\ndef to_nir(module: torch.nn.Module) -> nir.NIRNode:\n return nirtorch.torch_to_nir(\n module, # The Norse/Torch module to parse\n _my_module_dict # The dictionary that maps Norse module to NIR\n )\n```\nRead the full documentation at [neuroir.org/docs](https://neuroir.org/docs/dev_pytorch.html).\n\n\n## Acknowledgements\nIf you use NIRTorch in your work, please cite the [our work in Nature Communications](https://www.nature.com/articles/s41467-024-52259-9)\n\n```\narticle{NIR2024, \n title={Neuromorphic intermediate representation: A unified instruction set for interoperable brain-inspired computing}, \n author={Pedersen, Jens E. and Abreu, Steven and Jobst, Matthias and Lenz, Gregor and Fra, Vittorio and Bauer, Felix Christian and Muir, Dylan Richard and Zhou, Peng and Vogginger, Bernhard and Heckel, Kade and Urgese, Gianvito and Shankar, Sadasivan and Stewart, Terrence C. and Sheik, Sadique and Eshraghian, Jason K.}, \n rights={2024 The Author(s)},\n DOI={10.1038/s41467-024-52259-9}, \n number={1},\n journal={Nature Communications}, \n volume={15},\n year={2024}, \n month=sep, \n pages={8122},\n}\n```\n\n## For developers\nIf you want to make sure that your code is linted correctly on your local machine, use [pre-commit](https://pre-commit.com/) to automatically perform checks before every git commit. To use it, first install the package in your environment\n```\npip install pre-commit\n```\nand then install the pre-commit hooks that are listed in the root of this repository\n```\npre-commit install\n```\nNext time you commit some changes, all the checks will be run!\n",
"bugtrack_url": null,
"license": null,
"summary": "Neuromorphic Intermediate Representation",
"version": "2.0.5",
"project_urls": null,
"split_keywords": [
"neuromorphic",
" intermediate",
" representation"
],
"urls": [
{
"comment_text": null,
"digests": {
"blake2b_256": "3af190466c875ddb662b59a56a48e8ce609beb80a1e276b672919c6d578b5e95",
"md5": "b5db9837c5c94f59ef1f583b71bf32eb",
"sha256": "32c57c45af7057bebf14c89e9b2b0bc8b4da3210312472e38494d785591ee76a"
},
"downloads": -1,
"filename": "nirtorch-2.0.5-py3-none-any.whl",
"has_sig": false,
"md5_digest": "b5db9837c5c94f59ef1f583b71bf32eb",
"packagetype": "bdist_wheel",
"python_version": "py3",
"requires_python": null,
"size": 27329,
"upload_time": "2025-07-10T15:15:37",
"upload_time_iso_8601": "2025-07-10T15:15:37.891008Z",
"url": "https://files.pythonhosted.org/packages/3a/f1/90466c875ddb662b59a56a48e8ce609beb80a1e276b672919c6d578b5e95/nirtorch-2.0.5-py3-none-any.whl",
"yanked": false,
"yanked_reason": null
},
{
"comment_text": null,
"digests": {
"blake2b_256": "c1a9b3f7d5fef556dea6208008ed7b67c720f62a44699e21f86f884537759a90",
"md5": "0953546aeccf940f09963032eb354bae",
"sha256": "445a41e14547fd5f45714218591b321622679a92ac17dc4d3e42b252c93cc44e"
},
"downloads": -1,
"filename": "nirtorch-2.0.5.tar.gz",
"has_sig": false,
"md5_digest": "0953546aeccf940f09963032eb354bae",
"packagetype": "sdist",
"python_version": "source",
"requires_python": null,
"size": 52878,
"upload_time": "2025-07-10T15:15:39",
"upload_time_iso_8601": "2025-07-10T15:15:39.211399Z",
"url": "https://files.pythonhosted.org/packages/c1/a9/b3f7d5fef556dea6208008ed7b67c720f62a44699e21f86f884537759a90/nirtorch-2.0.5.tar.gz",
"yanked": false,
"yanked_reason": null
}
],
"upload_time": "2025-07-10 15:15:39",
"github": false,
"gitlab": false,
"bitbucket": false,
"codeberg": false,
"lcname": "nirtorch"
}
JSON