Axonal delay as synaptic delay in NetX

src/lava/lib/dl/netx/blocks/process.py

@@ -9,6 +9,7 @@
 from lava.magma.core.process.ports.ports import InPort, OutPort
 from lava.magma.core.process.process import AbstractProcess
 from lava.proc.dense.process import Dense as DenseSynapse
+from lava.proc.dense.process import DelayDense as DelayDenseSynapse
 from lava.proc.conv.process import Conv as ConvSynapse
 
 
@@ -116,6 +117,8 @@ class Dense(AbstractBlock):
         dictionary of neuron parameters. Defaults to None.
     weight : np.ndarray
         synaptic weight.
+    delay : np.ndarray
+        synaptic delay.
     bias : np.ndarray or None
         bias of neuron. None means no bias. Defaults to None.
     has_graded_input : dict
@@ -133,15 +136,25 @@ def __init__(self, **kwargs: Union[dict, tuple, list, int, bool]) -> None:
         super().__init__(**kwargs)
 
         weight = kwargs.pop('weight')
+        delay = kwargs.pop('delay', None)
         num_weight_bits = kwargs.pop('num_weight_bits', 8)
         weight_exponent = kwargs.pop('weight_exponent', 0)
 
-        self.synapse = DenseSynapse(
-            weights=weight,
-            weight_exp=weight_exponent,
-            num_weight_bits=num_weight_bits,
-            num_message_bits=self.input_message_bits,
-        )
+        if delay is None:
+            self.synapse = DenseSynapse(
+                weights=weight,
+                weight_exp=weight_exponent,
+                num_weight_bits=num_weight_bits,
+                num_message_bits=self.input_message_bits,
+            )
+        else:
+            self.synapse = DelayDenseSynapse(
+                weights=weight,
+                delays=delay.astype(int),
+                max_delay=62,
+                num_weight_bits=num_weight_bits,
+                num_message_bits=self.input_message_bits,
+            )
 
         if self.shape != self.synapse.a_out.shape:
             raise RuntimeError(

src/lava/lib/dl/netx/hdf5.py

@@ -54,6 +54,7 @@ class Network(AbstractProcess):
     reset_offset: int
         determines the phase shift of network reset if enabled. Defaults to 0.
     """
+
     def __init__(self,
                  net_config: str,
                  num_layers: Optional[int] = None,
@@ -326,12 +327,12 @@ def create_dense(layer_config: h5py.Group,
             neuron_config = layer_config['neuron']
         else:
             # Non-leaky integrator by default
-            neuron_config = {'type' : 'CUBA',
-                             'iDecay' : 0,
-                             'vDecay' : 4096,
-                             'vThMant' : 1 << 18 - 1,
-                             'refDelay' : 1,
-                             'gradedSpike' : False}
+            neuron_config = {'type': 'CUBA',
+                             'iDecay': 0,
+                             'vDecay': 4096,
+                             'vThMant': 1 << 18 - 1,
+                             'refDelay': 1,
+                             'gradedSpike': False}
         neuron_params = Network.get_neuron_params(neuron_config,
                                                   reset_interval=reset_interval,
                                                   reset_offset=reset_offset)
@@ -381,6 +382,27 @@ def create_dense(layer_config: h5py.Group,
                       'sign_mode': sign_mode,
                       'input_message_bits': input_message_bits}
 
+            if 'delay' in layer_config.keys():
+                delay = layer_config['delay']
+
+                if np.isscalar(delay):
+                    delay = delay * np.ones_like(weight)
+                elif len(delay.shape) == 1:  # delay is one dimensional
+                    num_delays = np.prod(delay.shape)
+                    if num_delays == weight.shape[0]:  # axonal delays to syn
+                        delay = np.repeat(delay.reshape(-1, 1),
+                                          weight.shape[1], axis=1)
+                    elif num_delays == weight.shape[1]:  # input delays
+                        delay = np.repeat(delay.reshape(1, -1),
+                                          weight.shape[0], axis=0)
+                    else:
+                        raise RuntimeError(f'Found invalid delay of shape '
+                                           f'{delay.shape} to be incompatible '
+                                           f'with weight of shape '
+                                           f'{weight.shape}.')
+
+                params['delay'] = delay
+
             # optional arguments
             if 'bias' in layer_config.keys():
                 params['bias'] = layer_config['bias']
@@ -571,9 +593,10 @@ def _create(self) -> List[AbstractProcess]:
                     reset_offset += 1
                 input_message_bits = layer.output_message_bits
                 if flatten_next:
-                    layers[-2].out.transpose([2, 1, 0]).flatten().connect(
-                        layers[-1].inp
-                    )
+                    if len(layers) > 1:
+                        layers[-2].out.transpose([2, 1, 0]).flatten().connect(
+                            layers[-1].inp
+                        )
                     flatten_next = False
                 else:
                     if len(layers) > 1:

tests/lava/lib/dl/netx/gts/ntidigits/.gitattributes

@@ -0,0 +1,3 @@
+input.npy filter=lfs diff=lfs merge=lfs -text
+ntidigits.net filter=lfs diff=lfs merge=lfs -text
+output.npy filter=lfs diff=lfs merge=lfs -text

tests/lava/lib/dl/netx/test_hdf5.py

@@ -174,6 +174,54 @@ def test_pilotnet_sdnn(self) -> None:
             f'Error was {error}.'
         )
 
+    def test_axonal_delay_ntidigits(self) -> None:
+        """Tests the output of ntidigits hdf5 description. This network
+        consists of axonal delay. So this tests specifically tests for
+        correctness of axonal delay."""
+        net_config = root + '/gts/ntidigits/ntidigits.net'
+        input = np.load(root + '/gts/ntidigits/input.npy')
+        gt = np.load(root + '/gts/ntidigits/output.npy')
+        num_steps = input.shape[1]
+
+        # skipping the last average layer which is not suppprted
+        net = netx.hdf5.Network(net_config=net_config, num_layers=5)
+
+        inp_gen = io.source.RingBuffer(data=input)
+        output_logger = io.sink.RingBuffer(shape=net.out_layer.shape,
+                                           buffer=num_steps)
+
+        inp_gen.s_out.connect(net.inp)
+        net.out.connect(output_logger.a_in)
+
+        run_condition = RunSteps(num_steps=num_steps)
+        run_config = TestRunConfig(select_tag='fixed_pt')
+        net.run(condition=run_condition, run_cfg=run_config)
+        output = output_logger.data.get()
+        net.stop()
+
+        out_ev = np.argwhere(output > 0)
+        gt_ev = np.argwhere(gt > 0)
+
+        error = np.abs(output[:, -1] - gt[:, 1]).sum()
+
+        if verbose:
+            if bool(os.environ.get('DISPLAY', None)):
+                plt.figure(figsize=(10, 5))
+                plt.plot(out_ev[:, 1], out_ev[:, 0], '.',
+                         markersize=12, label='Output Spikes')
+                plt.plot(gt_ev[:, 1], gt_ev[:, 0], '.', label='GT Spikes')
+                plt.xlabel(f'time')
+                plt.ylabel('Neuron ID')
+                plt.legend()
+                plt.show()
+
+        self.assertTrue(
+            error == 0,
+            f'Output spike and ground truth do not match for NTIDIGITS network.'
+            f'Found {output[output != gt] = } and {gt[output != gt] = }. '
+            f'Error was {error}.'
+        )
+
 
 if __name__ == '__main__':
     unittest.main()