sinabs 3.0.4.dev2__py3-none-any.whl → 3.1.0__py3-none-any.whl

sinabs/backend/dynapcnn/config_builder.py

@@ -1,9 +1,9 @@
 from abc import ABC, abstractmethod
-from typing import List
+from typing import Dict, List
 
 import samna
 
-from .dvs_layer import DVSLayer
+from .dynapcnn_layer import DynapcnnLayer
 from .mapping import LayerConstraints, get_valid_mapping
 
 
@@ -13,34 +13,35 @@ class ConfigBuilder(ABC):
     def get_samna_module(self):
         """Get the samna parent module that hosts all the appropriate sub-modules and classes.
 
-        Returns
-        -------
-        samna module
+        Returns:
+            samna module
         """
 
     @classmethod
     @abstractmethod
     def get_default_config(cls):
         """
-        Returns
-        -------
-        Returns the default configuration for the device type
+        Returns:
+            Default configuration for the device type
         """
 
     @classmethod
     @abstractmethod
-    def build_config(cls, model: "DynapcnnNetwork", chip_layers: List[int]):
+    def build_config(
+        cls,
+        layers: Dict[int, DynapcnnLayer],
+        layer2core_map: Dict[int, int],
+        destination_map: Dict[int, List[int]],
+    ):
         """Build the configuration given a model.
 
-        Parameters
-        ----------
-        model:
-            The target model
-        chip_layers:
-            Chip layers where the given model layers are to be mapped.
+        Args:
+            layers (Dict): Keys are layer indices, values are DynapcnnLayer instances.
+            layer2core_map (Dict): Keys are layer indices, values are corresponding
+                cores on hardware. Needed to map the destinations.
+            destination_map (Dict): Indices of destination layers for `layer`.
 
-        Returns
-        -------
+        Returns:
             Samna Configuration object
         """
 
@@ -49,8 +50,7 @@ class ConfigBuilder(ABC):
     def get_constraints(cls) -> List[LayerConstraints]:
         """Returns the layer constraints of a the given device.
 
-        Returns
-        -------
+        Returns:
             List[LayerConstraints]
         """
 
@@ -60,43 +60,27 @@ class ConfigBuilder(ABC):
         """Enable the monitor for a given set of layers in the config object."""
 
     @classmethod
-    def get_valid_mapping(cls, model: "DynapcnnNetwork") -> List[int]:
-        """Find a valid set of layers for a given model.
-
-        Parameters
-        ----------
-        model (DynapcnnNetwork):
-            A model
-
-        Returns
-        -------
-        List of core indices corresponding to each layer of the model:
-        The index of the core on chip to which the i-th layer in the
-        model is mapped is the value of the i-th entry in the list.
+    def map_layers_to_cores(cls, layers: Dict[int, DynapcnnLayer]) -> Dict[int, int]:
+        """Find a mapping from DynapcnnLayers onto on-chip cores
+
+        Args:
+            layers: Dict with layer indices as keys and DynapcnnLayer instances as values.
+
+        Returns:
+            Dict mapping layer indices (keys) to assigned core IDs (values).
         """
-        mapping = get_valid_mapping(model, cls.get_constraints())
-        # turn the mapping into a dict
-        mapping = {m[0]: m[1] for m in mapping}
-        # Check if there is a dvs layer in the model
-        num_dynapcnn_cores = len(model.sequence)
-        if isinstance(model.sequence[0], DVSLayer):
-            num_dynapcnn_cores -= 1
-        # apply the mapping
-        chip_layers_ordering = [mapping[i] for i in range(num_dynapcnn_cores)]
-        return chip_layers_ordering
+
+        return get_valid_mapping(layers, cls.get_constraints())
 
     @classmethod
     def validate_configuration(cls, config) -> bool:
         """Check if a given configuration is valid.
 
-        Parameters
-        ----------
-        config:
-            Configuration object
+        Args:
+            config: Configuration object.
 
-        Returns
-        -------
-        True if the configuration is valid, else false
+        Returns:
+            True if the configuration is valid, else false
         """
         is_valid, message = cls.get_samna_module().validate_configuration(config)
         if not is_valid:
@@ -127,22 +111,18 @@ class ConfigBuilder(ABC):
     def reset_states(cls, config, randomize=False):
         """Randomize or reset the neuron states.
 
-        Parameters
-        ----------
-            randomize (bool):
-                If true, the states will be set to random initial values. Else they will be set to zero
+        Args:
+            randomize (bool): If true, the states will be set to random initial values.
+                Else they will be set to zero
         """
 
     @classmethod
     def set_all_v_mem_to_zeros(cls, samna_device, layer_id: int) -> None:
         """Reset all memory states to zeros.
 
-        Parameters
-        ----------
-        samna_device:
-            samna device object to erase vmem memory.
-        layer_id:
-            layer index
+        Args:
+            samna_device: samna device object to erase vmem memory.
+            layer_id: layer index
         """
         mod = cls.get_samna_module()
         layer_constraint: LayerConstraints = cls.get_constraints()[layer_id]

sinabs/backend/dynapcnn/connectivity_specs.py

@@ -0,0 +1,48 @@
+"""
+functionality : list device-independent supported connections between layers on chip
+"""
+
+from typing import Union
+
+import torch.nn as nn
+
+import sinabs.layers as sl
+
+from .dvs_layer import DVSLayer
+
+Pooling = (sl.SumPool2d, nn.AvgPool2d)
+Weight = (nn.Conv2d, nn.Linear)
+Neuron = (sl.IAFSqueeze,)
+DVS = (DVSLayer,)
+SupportedNodeTypes = (*Pooling, *Weight, *Neuron, *DVS)
+
+VALID_SINABS_EDGE_TYPES_ABSTRACT = {
+    # convoluion is always followed by a neuron layer.
+    (Weight, Neuron): "weight-neuron",
+    # Neuron layer can be followed by pooling
+    (Neuron, Pooling): "neuron-pooling",
+    # Pooling can be followed by another pooling (will be consolidated)
+    (Pooling, Pooling): "pooling-pooling",
+    # Neuron layer can be followed by weight layer of next core
+    (Neuron, Weight): "neuron-weight",
+    # Pooling can be followed by weight layer of next core
+    (Pooling, Weight): "pooling-weight",
+    # Dvs can be followed by weight layer of next core
+    (DVS, Weight): "dvs-weight",
+    # Dvs can be followed by pooling layer
+    (DVS, Pooling): "dvs-pooling",
+}
+
+# Unpack dict
+VALID_SINABS_EDGE_TYPES = {
+    (source_type, target_type): name
+    for types, name in VALID_SINABS_EDGE_TYPES_ABSTRACT.items()
+    for source_type in types[0]
+    for target_type in types[1]
+}
+
+# Only `Merge` layers are allowed to join multiple inputs
+LAYER_TYPES_WITH_MULTIPLE_INPUTS = (sl.Merge,)
+
+# Neuron and pooling layers can have their output sent to multiple cores
+LAYER_TYPES_WITH_MULTIPLE_OUTPUTS = (*Neuron, *Pooling, *DVS)

sinabs/backend/dynapcnn/discretize.py

@@ -19,21 +19,14 @@ def discretize_conv_spike(
     This function takes a 2D convolutional and a spiking layer and returns a
     copy of each, with discretized weights, bias and threshold.
 
-    Parameters
-    ----------
-    conv_lyr: nn.Conv2d
-        Convolutional layer
-    spike_lyr: sl.IAF
-        Spiking layer
-    to_int: bool
-        Use integer types for discretized parameter
-
-    Returns
-    -------
-    nn.Conv2d
-        Discretized copy of convolutional layer
-    sl.IAF
-        Discretized copy of spiking layer
+    Args:
+        conv_lyr (nn.Conv2d):  Convolutional layer.
+        spike_lyr (sl.IAF):  Spiking layer.
+        to_int (bool): Use integer types for discretized parameter.
+
+    Returns:
+        Tuple containing a discretized copy of convolutional layer and
+        a discretized copy of spiking layer.
     """
     conv_lyr_copy = deepcopy(conv_lyr)
     spike_lyr_copy = deepcopy(spike_lyr)
@@ -48,21 +41,14 @@ def discretize_conv_spike_(
     This function takes a 2D convolutional and a spiking layer and discretizes
     weights, bias and threshold in-place.
 
-    Parameters
-    ----------
-    conv_lyr: nn.Conv2d
-        Convolutional layer
-    spike_lyr: sl.IAF
-        Spiking layer
-    to_int: bool
-        Use integer types for discretized parameter
-
-    Returns
-    -------
-    nn.Conv2d
-        Discretized convolutional layer
-    sl.IAF
-        Discretized spiking layer
+    Args:
+        conv_lyr (nn.Conv2d): Convolutional layer.
+        spike_lyr (sl.IAF): Spiking layer.
+        to_int (bool): Use integer types for discretized parameter.
+
+    Returns:
+        A tuple containing a discretized convolutional layer and a
+        discretized spiking layer.
     """
 
     return _discretize_conv_spk_(conv_lyr, spike_lyr, to_int=to_int)
@@ -74,28 +60,20 @@ def discretize_conv(
     spk_thr_low: float,
     spk_state: Optional[torch.Tensor] = None,
     to_int: bool = True,
-):
+) -> nn.Conv2d:
     """Discretize convolutional layer.
 
     This function takes a 2D convolutional layer and parameters of a subsequent
     spiking layer to return a discretized copy of the convolutional layer.
 
-    Parameters
-    ----------
-    layer: nn.Conv2d
-        Convolutional layer
-    spk_thr: float
-        Upper threshold of subsequent spiking layer
-    spk_thr_low: float
-        Lower threshold of subsequent spiking layer
-    spk_state: torch.Tensor or None
-        State of spiking layer.
-    to_int: bool
-        Use integer types for discretized parameter
-
-    Returns
-    -------
-    nn.Conv2d
+    Args:
+        layer (nn.Conv2d): Convolutional layer.
+        spk_thr (float): Upper threshold of subsequent spiking layer.
+        spk_thr_low (float): Lower threshold of subsequent spiking layer.
+        spk_state (torch.Tensor): State of spiking layer. Defaults to None.
+        to_int (bool): Use integer types for discretized parameter. Defaults to True.
+
+    Returns:
         Discretized copy of convolutional layer
     """
     lyr_copy = deepcopy(layer)
@@ -115,28 +93,20 @@ def discretize_conv_(
     spk_thr_low: float,
     spk_state: Optional[torch.Tensor] = None,
     to_int: bool = True,
-):
+) -> nn.Conv2d:
     """Discretize convolutional layer, in-place.
 
     This function discretizes a 2D convolutional layer in-place, based on
     parameters of a subsequent spiking layer.
 
-    Parameters
-    ----------
-    layer: nn.Conv2d
-        Convolutional layer
-    spk_thr: float
-        Upper threshold of subsequent spiking layer
-    spk_thr_low: float
-        Lower threshold of subsequent spiking layer
-    spk_state: torch.Tensor or None
-        State of spiking layer.
-    to_int: bool
-        Use integer types for discretized parameter
-
-    Returns
-    -------
-    nn.Conv2d
+    Args:
+        layer (nn.Conv2d): Convolutional layer.
+        spk_thr (float):  Upper threshold of subsequent spiking layer.
+        spk_thr_low (float): Lower threshold of subsequent spiking layer.
+        spk_state (torch.Tensor): State of spiking layer. Defaults to None.
+        to_int (bool): Use integer types for discretized parameter. Defaults to True.
+
+    Returns:
         Discretized convolutional layer
     """
     layer_discr, __ = _discretize_conv_spk_(
@@ -154,26 +124,20 @@ def discretize_spk(
     conv_weight: torch.Tensor,
     conv_bias: Optional[torch.Tensor] = None,
     to_int: bool = True,
-):
+) -> sl.IAF:
     """Discretize spiking layer.
 
     This function takes a spiking layer and parameters of a preceding
     convolutional layer to return a discretized copy of the spiking layer.
 
-    Parameters
-    ----------
-    layer: sl.IAF
-        Spiking layer
-    conv_weight: torch.Tensor
-        Weight tensor of preceding convolutional layer
-    conv_bias: torch.Tensor or None
-        Bias of preceding convolutional layer
-    to_int: bool
-        Use integer types for discretized parameter
-
-    Returns
-    -------
-    sl.IAF
+    Args:
+    layer (sl.IAF): Spiking layer.
+    conv_weight (torch.Tensor):  Weight tensor of preceding convolutional layer.
+    conv_bias (torch.Tensor): Bias of preceding convolutional layer.
+        Optional argument, defaults to None.
+    to_int (bool): Use integer types for discretized parameter.
+
+    Returns:
         Discretized copy of spiking layer
     """
     lyr_copy = deepcopy(layer)
@@ -188,26 +152,20 @@ def discretize_spk_(
     conv_weight: torch.Tensor,
     conv_bias: Optional[torch.Tensor] = None,
     to_int: bool = True,
-):
+) -> sl.IAF:
     """Discretize spiking layer in-place.
 
     This function discretizes a spiking layer in-place, based on parameters of a
     preceding convolutional layer.
 
-    Parameters
-    ----------
-    layer: sl.IAF
-        Spiking layer
-    conv_weight: torch.Tensor
-        Weight tensor of preceding convolutional layer
-    conv_bias: torch.Tensor or None
-        Bias of preceding convolutional layer
-    to_int: bool
-        Use integer types for discretized parameter
-
-    Returns
-    -------
-    sl.IAF
+    Args:
+        layer (sl.IAF): Spiking layer.
+        conv_weight (torch.Tensor): Weight tensor of preceding convolutional layer.
+        conv_bias (torch.Tensor): Bias of preceding convolutional layer.
+            Optional argument, defaults to None.
+        to_int (bool): Use integer types for discretized parameter.
+
+    Returns:
         Discretized spiking
     """
     __, layer_discr = _discretize_conv_spk_(
@@ -225,7 +183,7 @@ def _discretize_conv_spk_(
     conv_weight: Optional[torch.Tensor] = None,
     conv_bias: Optional[torch.Tensor] = None,
     to_int: bool = True,
-):
+) -> Tuple[nn.Conv2d, sl.IAF]:
     """Discretize convolutional and spiking layer.
 
     Determine and apply a suitable scaling factor for weight and bias of
@@ -234,34 +192,27 @@ def _discretize_conv_spk_(
     providing layers, respective parameters can be provided directly. If a layer
     is not provided, `None` will be returned instead of its discrete version.
 
-    Parameters
-    ----------
-        conv_lyr: nn.Conv2d or None
-            Convolutional layer
-        spike_lyr: sl.IAF or None
-            Spiking layer
-        spk_thr: float or None
-            Upper threshold of spiking layer. Has to be provided if `spike_lyr` is `None`.
+    Args:
+        conv_lyr (nn.Conv2d): Convolutional layer. Optional argument, defaults to None.
+        spike_lyr (sl.IAF): Spiking layer. Optional argument, defaults to None.
+        spk_thr (float): Upper threshold of spiking layer. Optional argument, defaults
+            to None. Has to be provided if `spike_lyr` is `None`.
             Is ignored otherwise.
-        spk_thr_low: float or None
-            Lower threshold of spiking layer. Has to be provided if `spike_lyr` is `None`.
+        spk_thr_low (float): Lower threshold of spiking layer. Optional argument,
+            defaults to None. Has to be provided if `spike_lyr` is `None`.
             Is ignored otherwise.
-        spk_state: torch.Tensor or None
-            State of spiking layer. Ignored if `spike_lyr` is not `None`.
-        conv_weight: torch.Tensor or None
-            Weight of convolutional layer. Has to be provided if `conv_lyr` is `None`.
+        spk_state (torch.Tensor): State of spiking layer. Optional argument
+            defaults to None. Ignored if `spike_lyr` is not `None`.
+        conv_weight (torch.Tensor): Weight of convolutional layer. Optional argument,
+            defaults to None. Has to be provided if `conv_lyr` is `None`.
             Is ignored otherwise.
-        conv_bias: torch.Tensor or None
-            Bias of convolutional layer. Ignored if `conv_lyr` is not `None`.
-        to_int: bool
-            Use integer types for discretized parameters.
-
-    Returns
-    -------
-        nn.Conv2d or None
-            Discretized convolutional layer if `conv_lyr` is not `None`, else `None`
-        sl.IAF or None
-            Discretized spiking layer if `spk_lyr` is not `None`, else `None`
+        conv_bias (torch.Tensor): Bias of convolutional layer. Optional argument,
+            defaults to None. Ignored if `conv_lyr` is not `None`.
+        to_int (bool): Use integer types for discretized parameters.
+
+    Returns:
+            Discretized convolutional layer if `conv_lyr` is not `None`, else `None`.
+            and discretized spiking layer if `spk_lyr` is not `None`, else `None`
     """
 
     if conv_lyr is None:
@@ -287,7 +238,6 @@ def _discretize_conv_spk_(
             conv_bias = torch.zeros(conv_lyr.out_channels)
 
     if spike_lyr is None:
-
         discr_spk = False
 
         if spk_thr is None or spk_thr_low is None:
@@ -361,17 +311,12 @@ def determine_discretization_scale(obj: torch.Tensor, bit_precision: int) -> flo
     Determine how much the values of a torch tensor can be scaled in order to fit
     the given precision
 
-    Parameters
-    ----------
-        obj: torch.Tensor
-            Tensor that is to be scaled
-        bit_precision: int
-            The precision in bits
-
-    Returns
-    -------
-        float
-            The scaling factor
+    Args:
+        obj (torch.Tensor): Tensor that is to be scaled.
+        bit_precision (int): The precision in bits.
+
+    Returns:
+        The scaling factor
     """
 
     # Discrete range
@@ -398,19 +343,14 @@ def discretize_tensor(
 ) -> torch.Tensor:
     """Scale a torch.Tensor and cast it to discrete integer values.
 
-    Parameters
-    ----------
-        obj: torch.Tensor
-            Tensor that is to be discretized
-        scaling: float
-            Scaling factor to be applied before discretization
-        to_int: bool
-            If False, round the values, but don't cast to Int. (Default True).
-
-    Returns
-    -------
-        torch.Tensor
-            Scaled and discretized copy of `obj`.
+    Args:
+        obj (torch.Tensor): Tensor that is to be discretized.
+        scaling (float): Scaling factor to be applied before discretization.
+        to_int (bool): If False, round the values, but don't cast to Int.
+            Defaults to True.
+
+    Returns:
+        Scaled and discretized copy of `obj`.
     """
 
     # Scale the values
@@ -428,17 +368,12 @@ def discretize_tensor(
 def discretize_scalar(obj: float, scaling: float) -> int:
     """Scale a float and cast it to discrete integer values.
 
-    Parameters
-    ----------
-        obj: float
-            Value that is to be discretized
-        scaling: float
-            Scaling factor to be applied before discretization
-
-    Returns
-    -------
-        int
-            Scaled and discretized copy of `obj`.
+    Args:
+        obj (float): Value that is to be discretized.
+        scaling (float): Scaling factor to be applied before discretization.
+
+    Returns:
+        Scaled and discretized copy of `obj`.
     """
 
     # Scale the values