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(n_steps, int], t_max, …)

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) → jaxsnn.event.topology.Topology

Convert NIRGraph to jaxsnn Topology

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(topology: jaxsnn.event.topology.Topology, params: dict, output: Optional[List[str], None] = None) → nir.ir.graph.NIRGraph

Convert a jaxsnn model to a NIRGraph. Note that the output node is not explicitly represented in the jaxsnn topology. Therefore, if an output node is desired in the NIRGraph, it can be specified explicitly via the output argument after which node of the jaxsnn topology the output node should be added.

Parameters

  • topology – jaxsnn topology to convert.

  • params – Parameters for the model.

  • output – Keys of the output nodes in the jaxsnn topology.

Returns

Converted NIRGraph.

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”,)