jaxsnn.event.functional

Modules

jaxsnn.event.functional.trajectory(…)

Simulate over multiple time steps for recurrent sub-networks.

Functions

jaxsnn.event.functional.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).