rockpool.nn.modules.to_nir

rockpool.nn.modules.to_nir(module: TorchModule, sample_data: Tensor | None = None, model_name: str = 'rockpool') → NIRNode[source]

Convert a Rockpool module into a NIR graph for export

Parameters:

  • module (TorchModule) – A Rockpool TorchModule to export

  • sample_data (Optional[Tensor]) – If needed, a Tensor containg dummy data of the shape expected by module. Default: Will be generated automatically from module.size_in.

  • model_name (str) – An optional string naming the model. Default: "rockpool".

Returns:

A NIR graph for export

Return type:

nir.NIRNode