jaxsnn.event.utils.from_nir_data.Topology

class jaxsnn.event.utils.from_nir_data.Topology(t_max: float, backprop_method: str = 'analytical', mock: bool = True, inter_batch_entry_wait: int = 0, has_external_events: bool = False)

Bases: jaxsnn.base.topology.AbstractTopology[jaxsnn.base.topology.ModelInitFnT, jaxsnn.base.topology.ModelApplyFnT]

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

Stores layer generators with associated parameters (nodes) and the connections (edges) between them. Handles construction of cleaned graphs, strongly connected components (SCCs), and simulation trajectories for efficient event-based computation. The backward pass can be configured to use EventProp or analytical gradients.

Params t_max

Maximum time for the simulation.

Params backprop_method

Either “eventprop” for the EventProp algorithm or “analytical” for the time-to-first-spike solver.

Params mock

If True, the topology will not be connected to a BSS2 experiment and will only run in software simulation mode.

Params inter_batch_entry_wait

Wait time between batch entries in FPGA cycles.

Params has_external_events

If True, the recordings of an externally performed forward pass are expected to be provided.

__init__(t_max: float, backprop_method: str = 'analytical', mock: bool = True, inter_batch_entry_wait: int = 0, has_external_events: bool = False)

Initialize an empty topology.

Methods

__init__(t_max[, backprop_method, mock, …])	Initialize an empty topology.
attach_layer_fns()	Generate and attach functional closures to each graph node.
extract_bss2_data(node)	Extract BSS2 data from the given node’s HX module.
generate_apply_fn(sccs_ordered)	Generate the main apply function for the SNN model.

attach_layer_fns() → None

Generate and attach functional closures to each graph node.

Iterates through the graph nodes and calls the generator function of each module to create its specific initialization, state, and step functions. These functions are then attached to the module instance stored in the graph node.

extract_bss2_data(node)

Extract BSS2 data from the given node’s HX module.

generate_apply_fn(sccs_ordered: List[Tuple[str, …]]) → Callable[[Dict[str, jaxsnn.event.types.Spike], Dict[str, jax.Array]], Tuple[Optional[Dict[str, StateT], None], Dict[str, jaxsnn.event.types.Spike]]]

Generate the main apply function for the SNN model.

Constructs a function that executes the network simulation by iterating through the ordered strongly connected components (SCCs). It handles both software simulation (via JAX scan) and hardware execution (via custom VJP and callback).

Parameters

sccs_ordered – List of SCCs in topological order.

Returns

A callable ModelApplyFn that takes inputs and parameters and returns the network output.