jaxsnn.discrete.topology.BaseTopology
-
class
jaxsnn.discrete.topology.BaseTopology 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.
-
__init__() → None Initialize an empty topology.
Methods
__init__()Initialize an empty topology.
add(modules)Add a dictionary of nodes to the topology.
Attach functional closures (init/state/step) to each graph node.
connect(edges)Add directed edges between nodes in the topology.
done()Finalize the graph topology into one init/apply pair.
generate_apply_fn(sccs_ordered)Generate the model application (forward pass) function for the topology.
generate_init_fn(sccs_ordered)Generate the model initialization function for the topology.
pre_nodes(node)pre_population_nodes(node)-
add(modules: Dict[str, ModuleT]) → None Add a dictionary of nodes to the topology.
- Parameters
nodes – A dictionary mapping node names to layer definitions.
-
abstract
attach_layer_fns() → None Attach functional closures (init/state/step) to each graph node.
-
connect(edges: List[Tuple[str, str]]) → None Add directed edges between nodes in the topology.
- Parameters
edges – List of (source, target) node name tuples.
-
done() → Tuple[ModelInitFnT, ModelApplyFnT] Finalize the graph topology into one init/apply pair.
Constructs initialization and forward pass functions based on the topology and user-defined node generators. This includes automatic SCC decomposition, and handling of the forward and backward dynamics.
- Returns
A tuple containing: - init: function to initialize weights - apply: function to apply the topology to inputs
-
abstract
generate_apply_fn(sccs_ordered: List[Tuple[str, …]]) → ModelApplyFnT Generate the model application (forward pass) function for the topology.
- Parameters
sccs_ordered – List of strongly connected components in topological order.
- Returns
Model application function.
-
generate_init_fn(sccs_ordered: List[Tuple[str, …]]) → ModelInitFnT Generate the model initialization function for the topology.
- Parameters
sccs_ordered – List of strongly connected components in topological order.
- Returns
Model initialization function.
-
pre_nodes(node: str) → List[str]
-
pre_population_nodes(node: str) → List[str]
-
pre_process() → List[Tuple[str, …]]
-