jaxsnn.event.modules.hx.lif.neuron

Classes

BaseModule(layer_idx)

DLSGlobal

Experiment(topology, *args, …)

Hardware experiment class for executing experiment on BrainScaleS-2

LIFParameters(tau_syn, tau_mem, v_th, …)

LIFState(V, I)

Neuron(layer_idx, n_events, n_hw_spikes, …)

NeuronParameters(tau_syn, tau_mem, v_th, …)

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.hx.lif.neuron.HXLIF(size: int, n_steps: int, params: NeuronParameters, time_offset: float = 0.0, n_hw_spikes: Optional[int] = None, chip_coordinate: Optional[DLSGlobal] = None, enable_spike_recording: bool = True, enable_cadc_recording: bool = False, enable_cadc_recording_placement_in_dram: bool = False, enable_madc_recording: bool = False, record_neuron_id: Optional[int] = None, placement_constraint: Optional[List[halco.LogicalNeuronOnDLS]] = None, neuron_structure: Optional[Morphology] = None, **extra_params) → Population

Create a LIF neuron population for execution on BrainScaleS.

Parameters

  • size – Number of neurons in the population.

  • n_steps – Number of simulation steps.

  • params – Neuron parameters (hardware-compatible).

  • time_offset – Optional time offset for all neurons.

  • n_hw_spikes – Optional hardware spike limit.

  • chip_coordinate – Optional chip coordinate for placement.

  • enable_spike_recording – Enable spike recording.

  • enable_cadc_recording – Enable CADC recording.

  • enable_cadc_recording_placement_in_dram – Place CADC recording in DRAM.

  • enable_madc_recording – Enable MADC recording.

  • record_neuron_id – Restrict recording to a specific neuron.

  • placement_constraint – Placement constraints for neurons.

  • neuron_structure – Optional neuron morphology.

  • extra_params – Additional hardware-specific parameters.

Returns

Population object for hardware LIF neurons.

jaxsnn.event.modules.hx.lif.neuron.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.hx.lif.neuron.build_step_function(pre_layer_pop_nodes: List[str], pre_layer_param_nodes: List[str], pre_layer_params: Dict[str, float], lif_params: jaxsnn.event.modules.lif.parameters.LIFParameters, node_idx_mapping: Dict[str, int], 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]]

Build the step function for a hardware LIF neuron layer.

Assembles the solver, neuron dynamics, and transition functions for both recurrent and feedforward simulation contexts.

Parameters

  • pre_layer_pop_nodes – Names of presynaptic population nodes.

  • pre_layer_param_nodes – Names of presynaptic parameter nodes.

  • pre_layer_params – Parameters for presynaptic connections.

  • lif_params – LIF neuron model parameters.

  • node_idx_mapping – Mapping of node names to indices.

  • t_max – Maximum simulation time.

  • node – Name of this node.

Returns

EventStepFn for the LIF neuron layer.

jaxsnn.event.modules.hx.lif.neuron.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.hx.lif.neuron.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.

jaxsnn.event.modules.hx.lif.neuron.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.