jaxsnn.event.utils.from_nir

Conversion of a NIR graph to jaxsnn-model

Classes

ConversionConfig(t_max, n_steps, int], …)

Configuration for the conversion from NIR to jaxsnn.

LIFParameters(tau_syn, tau_mem, v_th, …)

Parameter

alias of jax.Array

Population(generator, List[str], Dict[str, …)

Event-driven population module

Projection(generator, int], …)

Event-driven projection module

SourcePopulation(generator, parameters, …)

Event-driven source population module

Topology(t_max, backprop_method, mock, …)

Represents a spiking neural network (SNN) topology as a directed graph of layers.

Functions

jaxsnn.event.utils.from_nir.LIF(size: int, n_steps: int, params: jaxsnn.event.modules.lif.parameters.LIFParameters) → jaxsnn.event.types.Population

Creates a LIF layer for event-based simulation and backpropagation. Returns a generator function that builds the layer based on input connections, delays, and chosen backpropagation strategy. Supports both forward-only simulation and custom backward passes (eventprop).

Parameters

  • size – Number of neurons in the layer.

  • n_steps – Number of event steps in the simulation.

  • params – Parameters for the LIF neuron dynamics.

Returns

A Population object containing the generator and parameters.

jaxsnn.event.utils.from_nir.Linear(mean: float = 0.5, std: float = 2.0, min_delay: float = 0.0, pre_weights: Optional[jax.Array, None] = None) → jaxsnn.event.types.Projection

Creates a Linear projection layer

Either:

  • initialize weights from a Gaussian (mean, std), or

  • provide a concrete weight array.

Parameters

  • mean – Mean of the Gaussian distribution for weight initialization.

  • std – Standard deviation of the Gaussian distribution.

  • min_delay – Minimum delay associated with this projection.

  • pre_weights – Optional weight array. If provided, mean and std are ignored.

Returns

A Projection object containing the generator and parameters.

jaxsnn.event.utils.from_nir.Source(size: int) → jaxsnn.event.types.SourcePopulation

Creates a Source population layer representing external input.

Parameters

size – Number of neurons/channels in the source layer.

Returns

A SourcePopulation object containing the generator and parameters.

jaxsnn.event.utils.from_nir.dataclass(cls=None, /, *, init=True, repr=True, eq=True, order=False, unsafe_hash=False, frozen=False, match_args=True, kw_only=False, slots=False, weakref_slot=False)

Add dunder methods based on the fields defined in the class.

Examines PEP 526 __annotations__ to determine fields.

If init is true, an __init__() method is added to the class. If repr is true, a __repr__() method is added. If order is true, rich comparison dunder methods are added. If unsafe_hash is true, a __hash__() method is added. If frozen is true, fields may not be assigned to after instance creation. If match_args is true, the __match_args__ tuple is added. If kw_only is true, then by default all fields are keyword-only. If slots is true, a new class with a __slots__ attribute is returned.

jaxsnn.event.utils.from_nir.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.utils.from_nir.node_from_nir(node_key: str, node: nir.ir.node.NIRNode, config: jaxsnn.event.utils.from_nir.ConversionConfig) → Union[jaxsnn.event.types.Population, jaxsnn.event.types.Projection, jaxsnn.event.types.SourcePopulation]

Convert a NIR node to a jaxsnn module.

Parameters

  • node_key – Key of the node in the NIR graph.

  • node – NIR node to convert.

  • config – Conversion configuration.

Returns

Converted jaxsnn module.