jaxsnn.event.modules.lif.lif

Classes

LIFParameters(tau_syn, tau_mem, v_th, …)

LIFState(V, I)

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

Event-driven population module

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

StepState(neuron_state, time)

partial

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

Functions

jaxsnn.event.modules.lif.lif.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.modules.lif.lif.add_grads(tau_syn: Union[float, jax.Array], pre_pop_node: str, pre_param_node: str, grads: Dict[str, jax.Array], index_for_layer: int, adjoint_states: jaxsnn.event.types.StepState, weights: jax.Array, adjoint_spikes: Dict[str, jaxsnn.event.types.Spike], queue_index: int) → Tuple[Dict[str, jax.Array], Dict[str, jaxsnn.event.types.Spike]]

Add gradients with respect to synaptic weights and update adjoint spikes time for backpropagating the gradient.

Parameters

  • tau_syn – Synaptic time constant.

  • pre_pop_node – Name of the pre-synaptic population node.

  • pre_param_node – Name of the parameter node associated with the connection.

  • grads – Dictionary containing gradients with respect to the parameters.

  • index_for_layer – Index of the pre-synaptic neuron.

  • adjoint_states – Adjoint states of the post-synaptic layer.

  • weights – Dictionary containing the weights.

  • adjoint_spikes – Adjoint spikes structure used for backpropagation.

  • queue_index – Index in the spike queue for the input event.

Returns

Updated grads and adjoint_spikes.

jaxsnn.event.modules.lif.lif.adjoint_lif_exponential_flow(tau_syn: Union[jax.Array, float], tau_mem: Union[jax.Array, float]) → Callable[[jaxsnn.event.states.LIFState, jax.Array], jaxsnn.event.states.LIFState]
jaxsnn.event.modules.lif.lif.adjoint_step(adjoint_dynamics: Callable[[Optional[jaxsnn.base.types.BaseState, None], jax.Array], Optional[jaxsnn.base.types.BaseState, None]], adjoint_tr_dynamics: Callable[[jaxsnn.event.types.StepState, jaxsnn.event.types.Spike, jaxsnn.event.types.Spike, jax.Array, Dict[str, jax.Array], Dict[str, jax.Array], Dict[str, jaxsnn.event.types.Spike]], Tuple[jaxsnn.event.types.StepState, Dict[str, jax.Array], Dict[str, jaxsnn.event.types.Spike]]], t_max: float, adjoint_step_input: 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]

Perform one backward adjoint step integrating dynamics and transitions.

Parameters

  • adjoint_dynamics – Function integrating adjoint neuron dynamics.

  • adjoint_tr_dynamics – Function applying adjoint transition.

  • t_max – Maximum simulation time.

  • adjoint_step_input –

    Tuple containing (weights, spike, queue_index, adjoint_spike, adjoint_state, grads,

    adjoint_spikes).

Returns

Tuple with updated (weights, adjoint_spikes, adjoint_state, grads).

jaxsnn.event.modules.lif.lif.adjoint_transition(v_threshold: Union[jax.Array, float], input_layers: List[int], add_grads_fns: List[Callable[[Dict[str, jax.Array], int, jaxsnn.event.types.StepState, jax.Array, Dict[str, jaxsnn.event.types.Spike], int], Tuple[Dict[str, jax.Array], Dict[str, jaxsnn.event.types.Spike]]]], adjoint_states: jaxsnn.event.types.StepState, spike: jaxsnn.event.types.Spike, adjoint_spike: jaxsnn.event.types.Spike, queue_index: jax.Array, grads: Dict[str, jax.Array], weights: Dict[str, jax.Array], adjoint_spikes: Dict[str, jaxsnn.event.types.Spike]) → Tuple[jaxsnn.event.types.StepState, Dict[str, jax.Array], Dict[str, jaxsnn.event.types.Spike]]

Perform the adjoint transition step for a neuron or input spike.

Parameters

  • v_threshold – Membrane threshold potential.

  • input_layers – List of input layer indices.

  • add_grads_fns – List of functions to add gradients for weights.

  • adjoint_states – Current adjoint neuron states.

  • spike – Current spike event.

  • adjoint_spike – Current Adjoint spike event.

  • queue_index – Index of the spike in the queue.

  • grads – Gradients for weights.

  • weights – Weights of the current layer.

  • adjoint_spikes – Adjoint spikes through which the gradient flows.

Returns

Updated adjoint_states, grads, and adjoint_spikes.

jaxsnn.event.modules.lif.lif.build_adjoint_step_function(pre_layer_pop_nodes: List[str], pre_layer_param_nodes: List[str], node_index_mapping: Dict[str, int], lif_params: jaxsnn.event.modules.lif.parameters.LIFParameters, t_max: float) → 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]]

Build the adjoint step function used for backprop in spiking neurons. Constructs the adjoint step function by composing the adjoint neuron dynamics, adjoint transition dynamics, and gradient accumulation functions for each input connection. This function is used in the backward pass during event-based backpropagation.

Parameters

  • pre_layer_pop_nodes – List of names of input population nodes.

  • pre_layer_param_nodes – List of names of input parameter nodes.

  • node_index_mapping – Mapping from node names to indices.

  • lif_params – Parameters of the LIF neuron model.

  • t_max – Maximum simulation time.

Returns

A callable adjoint step function for use in backward scans.

jaxsnn.event.modules.lif.lif.build_apply_function(non_recurrent_step_fn: Callable, n_steps: int, layer_idx: int) → Callable

Construct the apply function for a LIF layer. Builds the forward pass function for a layer, scanning over time using the provided non-recurrent step function. The apply function handles integration over simulation time steps and manages the propagation of spikes, states, and queue indices. The result is a function that applies the spiking neuron dynamics to given weights, input spikes, and internal state.

Parameters

  • non_recurrent_step_fn – Step function to apply at each time step for the layer’s dynamics.

  • n_steps – Number of simulation time steps.

  • layer_idx – Index of the current layer to apply the function to.

Returns

A callable that executes the forward pass over n_steps for the specified layer.

jaxsnn.event.modules.lif.lif.build_custom_vjp_apply_function(apply_fn: Callable, wrapped_adjoint_step_fn: Callable, n_steps: int) → Callable

Build a custom VJP-enabled apply function for the LIF layer. Wraps the forward apply function to support event-based backpropagation by defining a custom vector-Jacobian product (VJP). Uses the provided wrapped adjoint step function to scan backwards through time steps.

Parameters

  • apply_fn – The forward apply function (without VJP).

  • wrapped_adjoint_step_fn – The adjoint step function wrapped with layer index, used for the backward scan.

  • n_steps – Number of simulation time steps.

Returns

A callable apply function with custom VJP defined.

jaxsnn.event.modules.lif.lif.build_step_functions(pre_layer_pop_nodes: List[str], pre_layer_param_nodes: List[str], pre_layer_params: Dict[str, float], scc_mask: List[str], lif_params: jaxsnn.event.modules.lif.parameters.LIFParameters, t_max: float, node: str) → Callable[[Tuple[Dict[str, jax.Array], Dict[str, jaxsnn.event.types.Spike], jaxsnn.event.types.StepState, Any, int, jax.Array, Any]], Tuple[jaxsnn.event.types.Spike, jaxsnn.event.types.StepState, jax.Array, jax.Array]]

Construct the recurrent and non-recurrent step functions for a LIF layer. This function builds the two core step functions required to simulate a spiking neuron layer in both recurrent and feedforward cases. It assembles the solver for finding next events, sets up the neuron dynamics, builds the per-connection transition functions, and incorporates minimum delay checks as defined by the graph structure and neuron parameters. The recurrent step function is used when a layer is part of a strongly connected component (SCC), and the non-recurrent function wraps it for use in feedforward contexts.

Parameters

  • pre_layer_pop_nodes – List of names of input population nodes.

  • pre_layer_param_nodes – List of names of input parameter nodes.

  • pre_layer_params – Dictionary of synaptic parameters (delays).

  • scc_mask – List of nodes indicating which inputs belong to SCC.

  • lif_params – Parameters for the LIF neuron model.

  • t_max – Maximum simulation time.

  • node – Name of the current layer/node.

Returns

The step function for the layer.

jaxsnn.event.modules.lif.lif.lif_exponential_flow(tau_syn: Union[jax.Array, float], tau_mem: Union[jax.Array, float]) → Callable[[jaxsnn.event.states.LIFState, jax.Array], jaxsnn.event.states.LIFState]
jaxsnn.event.modules.lif.lif.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.modules.lif.lif.next_event(solver: Callable[[StateT, float], jaxsnn.event.types.Spike], neuron_state: jaxsnn.base.types.BaseState, time: jax.Array, t_max: float) → jaxsnn.event.types.Event

Wraps a root solver to provide a cleaner API for returning next event.

Parameters

  • solver – The actual root solver function.

  • neuron_state – The state of the neurons.

  • time – Current simulation time.

  • t_max – Maximum time of the simulation.

Returns

Event object representing the spike which will occur next.

jaxsnn.event.modules.lif.lif.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.modules.lif.lif.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.modules.lif.lif.transition(v_reset: Union[float, jax.Array], input_node: str, state: jaxsnn.event.types.StepState, weights: jax.Array, spike_idx: int, spike_in_layer: bool) → jaxsnn.event.types.StepState

Apply state transition based on internal/input spike.

Parameters

  • v_reset – Reset potential for the neuron.

  • input_node – Name of the input node/layer.

  • state – Current step state containing neuron states.

  • weights – Synaptic weight matrix.

  • spike_idx – Index of the neuron that spiked.

  • spike_in_layer – Boolean indicating if spike is internal/input.

Returns

Updated step state after transition.

jaxsnn.event.modules.lif.lif.ttfs_solver(tau_mem: float, tau_syn: float, v_th: float, state: jaxsnn.event.states.LIFState, t_max: float) → jax.Array

Find the next spike time for special cases $ au_mem = au_syn$ and $ au_mem = 2 * au_syn$

Args:

tau_mem (float): Membrane time constant tau_syn (float): Synaptic time constant v_th (float): Threshold Voltage state (LIFState): State of the neuron (voltage, current) t_max (float): Maximum time which is to be searched

Returns:

float: Time of next threshold crossing or t_max if no crossing