jaxsnn.event

Modules

`jaxsnn.event.custom_lax`	Implement functionality of lax for easier debugging
`jaxsnn.event.encode`
`jaxsnn.event.functional`
`jaxsnn.event.hardware`
`jaxsnn.event.loss`
`jaxsnn.event.modules`
`jaxsnn.event.solver`
`jaxsnn.event.states`
`jaxsnn.event.stepping`
`jaxsnn.event.topology`
`jaxsnn.event.training`
`jaxsnn.event.types`
`jaxsnn.event.utils`

Classes

ConversionConfig(t_max, n_steps, int], …)

Configuration for the conversion from NIR to jaxsnn.

Functions

jaxsnn.event.from_nir(graph: nir.ir.graph.NIRGraph, config: jaxsnn.event.utils.from_nir.ConversionConfig)

Convert NIRGraph to jax-snn Topology

Restrictions for NIRGraph: - Only linear feed-forward SNNs are supported - CubaLIF and Linear layers are supported - Affine layers with bias==0 are currently supported - In terms of parameters, only homogeneous layers are supported - The analytical solver is only supported for non-external inputs

Parameters

graph – NIR graph to convert
config – Conversion configuration

Example: ```python nir_graph = nir.NIRGraph(…) cfg = jaxsnn.event.ConversionConfig(…)

topology = jaxsnn.event.from_nir(nir_graph, cfg) init, apply = topology.done() ```

jaxsnn.event.from_nir_data(nir_graph_data: nir.data_ir.graph.NIRGraphData, topology: jaxsnn.event.topology.Topology, observables=('spikes')) → Dict[str, jaxsnn.event.types.Spike]

Convert NIRGraphData to a dict of EventPropSpikes (jax-snn representation)

Parameters

nir_graph_data – NIRGraphData to be converted.
topology – jaxsnn Topology object.
observables – Observables to be converted, by default (‘spikes’,)

jaxsnn.event.to_nir_data(jaxsnn_dict: Dict[str, jaxsnn.event.types.Spike], topology: jaxsnn.event.topology.Topology, observables=('spikes')) → nir.data_ir.graph.NIRGraphData

Convert a dict of Spikes (jax-snn representation) to NIRGraphData.

Parameters

jaxsnn_dict – Dictionary of Spike objects where each entry represents the spikes for a corresponding node of the jaxsnn_model. Empty events in jaxsnn are encoded by idx = -1 and time = 2 * t_max.
topology – jaxsnn Topology object.
observables – Observables to be converted, by default (‘spikes’,)