jaxsnn.base.topology.AbstractTopology
-
class
jaxsnn.base.topology.AbstractTopology Bases:
abc.ABC,Generic[jaxsnn.base.topology.ModelInitFnT,jaxsnn.base.topology.ModelApplyFnT]-
__init__() Initialize self. See help(type(self)) for accurate signature.
Methods
add(modules)Add modules to the topology graph.
connect(edges)Add directed edges between nodes in the topology graph.
done()Finalize the topology and return (init_fn, apply_fn).
-
abstract
add(modules: Dict[str, ModuleT]) → None Add modules to the topology graph.
- Parameters
modules – Dictionary mapping module names to modules.
-
abstract
connect(edges: List[Tuple[str, str]]) → None Add directed edges between nodes in the topology graph.
- Parameters
edges – List of (source, target) tuples.
-
abstract
done() → Tuple[ModelInitFnT, ModelApplyFnT] Finalize the topology and return (init_fn, apply_fn).
- Returns
Tuple of initialization and application functions.
-