Synaptic delays for Dense connections

src/lava/proc/dense/models.py

@@ -14,7 +14,7 @@
 from lava.magma.core.resources import CPU
 from lava.magma.core.decorator import implements, requires, tag
 from lava.magma.core.model.py.model import PyLoihiProcessModel
-from lava.proc.dense.process import Dense, LearningDense
+from lava.proc.dense.process import Dense, LearningDense, DelayDense
 from lava.utils.weightutils import SignMode, determine_sign_mode,\
     truncate_weights, clip_weights
 
@@ -188,3 +188,152 @@ def run_spk(self):
         )
 
         self.recv_traces(s_in)
+
+
+class AbstractPyDelayDenseModel(PyLoihiProcessModel):
+    """Abstract Conn Process with Dense synaptic connections which incorporates
+    delays into the Conn Process.
+    """
+
+    @staticmethod
+    def get_del_wgts(weights, delays) -> np.ndarray:
+        """
+        Use self.weights and self.delays to create a matrix where the
+        weights are separated by delay. Returns 2D matrix of form
+        (num_flat_output_neurons * max_delay + 1, num_flat_input_neurons) where
+        del_wgts[
+            k * num_flat_output_neurons : (k + 1) * num_flat_output_neurons, :
+        ]
+        contains the weights for all connections with a delay equal to k.
+        This allows for the updating of the activation buffer and updating
+        weights.
+        """
+        return np.vstack([
+            np.where(delays == k, weights, 0)
+            for k in range(np.max(delays) + 1)
+        ])
+
+    def calc_act(self, s_in) -> np.ndarray:
+        """
+        Calculate the activations by performing del_wgts * s_in. This matrix
+        is then summed across each row to get the activations to the output
+        neurons for different delays. This activation vector is reshaped to a
+        matrix of the form
+        (n_flat_output_neurons * (max_delay + 1), n_flat_output_neurons)
+        which is then transposed to get the activation matrix.
+        """
+        return np.reshape(
+            np.sum(self.get_del_wgts(self.weights,
+                                     self.delays) * s_in, axis=1),
+            (np.max(self.delays) + 1, self.weights.shape[0])).T
+
+    def update_act(self, s_in):
+        """
+        Updates the activations for the connection.
+        Clears first column of a_buff and rolls them to the last column.
+        Finally, calculates the activations for the current time step and adds
+        them to a_buff.
+        This order of operations ensures that delays of 0 correspond to
+        the next time step.
+        """
+        self.a_buff[:, 0] = 0
+        self.a_buff = np.roll(self.a_buff, -1)
+        self.a_buff += self.calc_act(s_in)
+
+
+@implements(proc=DelayDense, protocol=LoihiProtocol)
+@requires(CPU)
+@tag("floating_pt")
+class PyDelayDenseModelFloat(AbstractPyDelayDenseModel):
+    """Implementation of Conn Process with Dense synaptic connections in
+    floating point precision. This short and simple ProcessModel can be used
+    for quick algorithmic prototyping, without engaging with the nuances of a
+    fixed point implementation. DelayDense incorporates delays into the Conn
+    Process.
+    """
+    s_in: PyInPort = LavaPyType(PyInPort.VEC_DENSE, bool, precision=1)
+    a_out: PyOutPort = LavaPyType(PyOutPort.VEC_DENSE, float)
+    a_buff: np.ndarray = LavaPyType(np.ndarray, float)
+    # weights is a 2D matrix of form (num_flat_output_neurons,
+    # num_flat_input_neurons) in C-order (row major).
+    weights: np.ndarray = LavaPyType(np.ndarray, float)
+    # delays is a 2D matrix of form (num_flat_output_neurons,
+    # num_flat_input_neurons) in C-order (row major).
+    delays: np.ndarray = LavaPyType(np.ndarray, int)
+    num_message_bits: np.ndarray = LavaPyType(np.ndarray, int, precision=5)
+
+    def run_spk(self):
+        # The a_out sent on a each timestep is a buffered value from dendritic
+        # accumulation at timestep t-1. This prevents deadlocking in
+        # networks with recurrent connectivity structures.
+        self.a_out.send(self.a_buff[:, 0])
+        if self.num_message_bits.item() > 0:
+            s_in = self.s_in.recv()
+        else:
+            s_in = self.s_in.recv().astype(bool)
+        self.update_act(s_in)
+
+
+@implements(proc=DelayDense, protocol=LoihiProtocol)
+@requires(CPU)
+@tag("bit_accurate_loihi", "fixed_pt")
+class PyDelayDenseModelBitAcc(AbstractPyDelayDenseModel):
+    """Implementation of Conn Process with Dense synaptic connections that is
+    bit-accurate with Loihi's hardware implementation of Dense, which means,
+    it mimics Loihi behaviour bit-by-bit. DelayDense incorporates delays into
+    the Conn Process. Loihi 2 has a maximum of 6 bits for delays, meaning a
+    spike can be delayed by 0 to 63 time steps."""
+
+    s_in: PyInPort = LavaPyType(PyInPort.VEC_DENSE, bool, precision=1)
+    a_out: PyOutPort = LavaPyType(PyOutPort.VEC_DENSE, np.int32, precision=16)
+    a_buff: np.ndarray = LavaPyType(np.ndarray, np.int32, precision=16)
+    # weights is a 2D matrix of form (num_flat_output_neurons,
+    # num_flat_input_neurons) in C-order (row major).
+    weights: np.ndarray = LavaPyType(np.ndarray, np.int32, precision=8)
+    delays: np.ndarray = LavaPyType(np.ndarray, np.int32, precision=6)
+    num_message_bits: np.ndarray = LavaPyType(np.ndarray, int, precision=5)
+
+    def __init__(self, proc_params):
+        super().__init__(proc_params)
+        # Flag to determine whether weights have already been scaled.
+        self.weights_set = False
+
+    def run_spk(self):
+        self.weight_exp: int = self.proc_params.get("weight_exp", 0)
+
+        # Since this Process has no learning, weights are assumed to be static
+        # and only require scaling on the first timestep of run_spk().
+        if not self.weights_set:
+            num_weight_bits: int = self.proc_params.get("num_weight_bits", 8)
+            sign_mode: SignMode = self.proc_params.get("sign_mode") \
+                or determine_sign_mode(self.weights)
+
+            self.weights = clip_weights(self.weights, sign_mode, num_bits=8)
+            self.weights = truncate_weights(self.weights,
+                                            sign_mode,
+                                            num_weight_bits)
+            self.weights_set = True
+
+            # Check if delays are within Loihi 2 constraints
+            if np.max(self.delays) > 63:
+                raise ValueError("DelayDense Process 'delays' expects values "
+                                 f"between 0 and 63 for Loihi, got "
+                                 f"{self.delays}.")
+
+        # The a_out sent at each timestep is a buffered value from dendritic
+        # accumulation at timestep t-1. This prevents deadlocking in
+        # networks with recurrent connectivity structures.
+        self.a_out.send(self.a_buff[:, 0])
+        if self.num_message_bits.item() > 0:
+            s_in = self.s_in.recv()
+        else:
+            s_in = self.s_in.recv().astype(bool)
+
+        a_accum = self.calc_act(s_in)
+        self.a_buff[:, 0] = 0
+        self.a_buff = np.roll(self.a_buff, -1)
+        self.a_buff += (
+            np.left_shift(a_accum, self.weight_exp)
+            if self.weight_exp > 0
+            else np.right_shift(a_accum, -self.weight_exp)
+        )

src/lava/proc/dense/process.py

@@ -144,3 +144,93 @@ def __init__(self,
                          log_config=log_config,
                          learning_rule=learning_rule,
                          **kwargs)
+
+
+class DelayDense(Dense):
+    def __init__(self,
+                 *,
+                 weights: np.ndarray,
+                 delays: ty.Union[np.ndarray, int],
+                 max_delay: ty.Optional[int] = 0,
+                 name: ty.Optional[str] = None,
+                 num_message_bits: ty.Optional[int] = 0,
+                 log_config: ty.Optional[LogConfig] = None,
+                 **kwargs) -> None:
+        """Dense, delayed connections between neurons. Realizes the following
+        abstract behavior: a_out = weights * s_in
+
+        Parameters
+        ----------
+        weights : numpy.ndarray
+            2D connection weight matrix of form (num_flat_output_neurons,
+            num_flat_input_neurons) in C-order (row major).
+
+        delays : numpy.ndarray, int
+            2D connection delay matrix of form (num_flat_output_neurons,
+            num_flat_input_neurons) in C-order (row major) or integer value if
+            the same delay should be used for all synapses.
+
+        max_delay: int, optional
+            Maximum expected delay. Should be set if delays change during
+            execution. Default value is 0, in this case the maximum delay
+            will be determined from the values given in 'delays'.
+
+        weight_exp : int, optional
+            Shared weight exponent of base 2 used to scale magnitude of
+            weights, if needed. Mostly for fixed point implementations.
+            Unnecessary for floating point implementations.
+            Default value is 0.
+
+        num_weight_bits : int, optional
+            Shared weight width/precision used by weight. Mostly for fixed
+            point implementations. Unnecessary for floating point
+            implementations.
+            Default is for weights to use full 8 bit precision.
+
+        sign_mode : SignMode, optional
+            Shared indicator whether synapse is of type SignMode.NULL,
+            SignMode.MIXED, SignMode.EXCITATORY, or SignMode.INHIBITORY. If
+            SignMode.MIXED, the sign of the weight is
+            included in the weight bits and the fixed point weight used for
+            inference is scaled by 2.
+            Unnecessary for floating point implementations.
+
+            In the fixed point implementation, weights are scaled according to
+            the following equations:
+            w_scale = 8 - num_weight_bits + weight_exp + isMixed()
+            weights = weights * (2 ** w_scale)
+
+        num_message_bits : int, optional
+            Determines whether the Dense Process deals with the incoming
+            spikes as binary spikes (num_message_bits = 0) or as graded
+            spikes (num_message_bits > 0). Default is 0.
+        """
+
+        super().__init__(weights=weights,
+                         num_message_bits=num_message_bits,
+                         name=name,
+                         log_config=log_config,
+                         **kwargs)
+
+        self._validate_delays(weights, delays)
+        shape = weights.shape
+        if max_delay == 0:
+            max_delay = int(np.max(delays))
+
+        # Variables
+        self.delays = Var(shape=shape, init=delays)
+        self.a_buff = Var(shape=(shape[0], max_delay + 1) , init=0)
+
+    @staticmethod
+    def _validate_delays(weights: np.ndarray, delays: np.ndarray) -> None:
+        if np.min(delays) < 0:
+            raise ValueError("DelayDense Process 'delays' expects only "
+                             f"positive values, got {delays}.")
+        if not isinstance(delays, int):
+            if np.shape(weights) != np.shape(delays):
+                raise ValueError("DelayDense Process 'delays' expects same "
+                                 f"shape than the weight matrix or int, got "
+                                 f"{delays}.")
+            if delays.dtype != int:
+                raise ValueError("DelayDense Process 'delays' expects integer "
+                                 f"value(s), got {delays}.")