jaxsnn.event.topology

Classes

Any(*args, **kwargs)

Special type indicating an unconstrained type.

BaseTopology()

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

Experiment(topology, *args, …)

Hardware experiment class for executing experiment on BrainScaleS-2

HXInputPopulation

alias of jaxsnn.event.hardware.modules.population.InputPopulation

HXPopulation

alias of jaxsnn.event.hardware.modules.population.Population

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

Spike(time, idx, current, layer_idx, internal)

Topology(t_max, backprop_method, mock, …)

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

partial

partial(func, *args, **keywords) - new function with partial application of the given arguments and keywords.

Functions

jaxsnn.event.topology.adjoint_trajectory(multi_layer_adjoint_step_fn: Callable, n_steps: int, res: Tuple, g: Tuple) → Tuple

Run the adjoint trajectory backward through time using JAX’s scan.

Parameters

  • multi_layer_adjoint_step_fn – Function performing one adjoint step for multiple layers.

  • n_steps – Number of time steps to run backward.

  • res – Tuple containing forward pass results (spikes, states weights, queue indices).

  • g – Tuple containing adjoint initial states (adjoint spikes, adjoint state, gradients).

Returns

Tuple of (gradients, adjoint spikes, adjoint states, None).

jaxsnn.event.topology.get_logger(name: str)
jaxsnn.event.topology.multi_layer_adjoint_step(adjoint_step_fns: Dict[str, Callable[[Tuple[Dict[str, jax.Array], jaxsnn.event.types.Spike, jax.Array, jaxsnn.event.types.Spike, jaxsnn.event.types.StepState, Dict[str, jax.Array], jaxsnn.event.types.Spike]], Tuple[Dict[str, jax.Array], jaxsnn.event.types.StepState, jax.Array, jax.Array]]], nodes: List[str], carry: Tuple[Dict[str, jax.Array], Dict[str, jaxsnn.event.types.Spike], jax.Array, Dict[str, jaxsnn.event.types.Spike], Dict[str, StateT], Dict[str, jax.Array]], step_idx: int) → Tuple[Tuple, int]

Perform one adjoint backward step through all layers in reverse order.

Parameters

  • adjoint_step_fns – Dictionary of adjoint step functions per layer.

  • nodes – List of node names (layers) to process.

  • carry – Tuple containing weights, spikes, queue indices, adjoint spikes, adjoint states, and grads.

  • step_idx – Index of the current step.

Returns

Updated carry tuple and zero as dummy scan output.

jaxsnn.event.topology.multi_layer_step(step_fns: Dict[str, Callable], nodes: List[str], node_index_mapping: Dict[str, int], carry: jaxsnn.event.types.Carry, step_idx: int) → Tuple[jaxsnn.event.types.Carry, int]

Perform one simulation step across multiple neuron layers.

Iterates through all layers in the strongly connected component (SCC), applies their respective step functions, and updates spikes, states, queue heads, and queue indices accordingly.

Parameters

  • step_fns – Dictionary of step functions for each layer, keyed by node name.

  • nodes – List of node names in the SCC to process.

  • node_index_mapping – Dictionary mapping node names to their indices.

  • carry – Tuple (parameters, spikes, states, queue_heads, queue_indices).

  • step_idx – Current step index.

Returns

Updated carry tuple and an integer placeholder (always 0).

jaxsnn.event.topology.trajectory(multi_layer_step_fn: Callable, n_steps: int, parameters: Dict[str, jax.Array], spikes: Dict[str, jaxsnn.event.types.Spike], external_spikes: Optional[Dict[str, jaxsnn.event.types.Spike], None], states: Dict[str, StateT], queue_heads: Dict[str, jax.Array]) → Tuple[Dict[str, jaxsnn.event.types.Spike], Dict[str, StateT], Dict[str, jax.Array], List[jax.Array]]

Simulate over multiple time steps for recurrent sub-networks.

Applies the multi-layer step function sequentially across n_steps using JAX’s scan. Maintains the states, spike history, and queue indices.

Parameters

  • multi_layer_step_fn – Function to advance all layers one time step.

  • n_steps – Number of simulation steps.

  • parameters – Model parameters.

  • spikes – Spikes from all layers.

  • states – Initial neuron states.

  • queue_heads – Dict of queue head arrays for input queuing.

Returns

Tuple of updated (spikes, states, parameters, queue_indices).