jaxsnn.event.stepping

Modules

jaxsnn.event.stepping.step(next_input_fn, …)	Find the next spike (external or internal), and evolve state to that point.
jaxsnn.event.stepping.step_existing_events
jaxsnn.event.stepping.step_multi_layer

Functions

jaxsnn.event.stepping.min_delay_check(input_nodes: List[str], min_delays: Dict[str, float], spikes: Dict[str, jaxsnn.event.types.EventT], spike_time: jax.Array) → Tuple[jax.Array, jax.Array]

Check if a given spike time is not too far ahead in the future.

For each recurrently connected input layer, finds the latest spike time plus the minimum delay, then compares to the given spike time. This ensures that no input spike can reach the neuron before the time of the spike.

Parameters

• input_nodes – List of indices of input layers.

• min_delays – Array of minimum delays per input connection.

• spikes – List of spike records for all layers.

• spike_time – The spike time to check.

Returns

Tuple containing - boolean indicating if spike_time is safe, - the computed safe time (minimum allowed spike time).

jaxsnn.event.stepping.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.stepping.next_input(input_layers: List[str], min_delays: Dict[str, float], spikes: Dict[str, jaxsnn.event.types.Spike], queue_heads: jax.Array, time: float, t_max: float) → Tuple[jax.Array, jax.Array, jaxsnn.event.types.Spike]

Find the next input spike event across potentially multiple input layers.

jaxsnn.event.stepping.step(next_input_fn: Callable[[Dict[str, jaxsnn.event.types.Spike], jax.Array, float, float], Tuple[jax.Array, jax.Array, jaxsnn.event.types.Spike]], min_delay_check_fn: Callable[[Dict[str, jaxsnn.event.types.Spike], jax.Array], Tuple[jax.Array, jax.Array]], dynamics: Callable[[jaxsnn.base.types.BaseState, jax.Array], jaxsnn.base.types.BaseState], tr_dynamics: List[Callable], t_max: float, solver: Callable[[StateT, float], jaxsnn.event.types.Spike], step_input: jaxsnn.event.types.Step) → Tuple[jaxsnn.event.types.Spike, jaxsnn.event.types.StepState, jax.Array, jax.Array]

Find the next spike (external or internal), and evolve state to that point.

Parameters

• next_input_fn – Function to find the next input spike.

• min_delay_check_fn – Function to check if spike is allowed.

• dynamics – Function describing continuous neuron dynamics.

• tr_dynamics – List of functions describing the transition after a spike.

• t_max – Maximum simulation time for this step.

• solver – Solver that returns the next internal event.

• step_input – Tuple containing (weights, spikes, state, _, layer_idx, queue_heads, _).

Returns

Tuple containing - Spike: The new spike event. - StepState: Updated neuron state after transition. - QueueHead: Updated queue heads. - QueueIndex: Spike queue index of the spike event.

jaxsnn.event.stepping.step_existing(input_layers: List[str], dynamics: Callable, transition_fns: List[Callable], node: str, t_max: float, step_input: jaxsnn.event.types.Step) → Tuple[jaxsnn.event.types.Spike, jaxsnn.event.types.StepState, jax.Array, jax.Array]

Simulate a single event-driven step for a spiking neuron layer using existing (from hardware) events.

Finds the next relevant spike (known hardware spike or input event), advances the neuron state to that event, applies the appropriate transition, and updates all event queues and indices accordingly.

Parameters

• input_layers – List of input layer names.

• valid_input_layer_indices – Indices of valid input layers for event matching.

• dynamics – Function describing the continuous neuron dynamics.

• transition_fns – List of transition functions for each possible input.

• node – Name of the current layer/node.

• t_max – Maximum simulation time for this step.

• step_input – Step object containing all state, parameters, queues, and indices.

Returns

Tuple containing: - Spike: The spike event at this step (or empty if none). - StepState: Updated neuron state after the event. - QueueHead: Updated queue head indices for all inputs. - QueueIndex: Index of the input queue used for the event.