Persistent delay buffer to correctly run slayer blocks one timestep at a time

src/lava/lib/dl/slayer/block/base.py

@@ -12,6 +12,33 @@
 from lava.lib.dl.slayer.utils import recurrent
 
 
+def step_delay(module, x):
+    """Step delay computation. This simulates the 1 timestep delay needed
+    for communication between layers.
+
+    Parameters
+    ----------
+    module: module
+        python module instance
+    x : torch.tensor
+        Tensor data to be delayed.
+    """
+    if hasattr(module, 'delay_buffer') is False:
+        module.delay_buffer = None
+    persistent_state = hasattr(module, 'neuron') \
+        and module.neuron.persistent_state is True
+    if module.delay_buffer is not None:
+        if module.delay_buffer.shape[0] != x.shape[0]:  # batch mismatch
+            module.delay_buffer = None
+    if persistent_state:
+        delay_buffer = 0 if module.delay_buffer is None else module.delay_buffer
+        module.delay_buffer = x[..., -1]
+    x = delay(x, 1)
+    if persistent_state:
+        x[..., 0] = delay_buffer
+    return x
+
+
 class AbstractInput(torch.nn.Module):
     """Abstract input block class. This should never be instantiated on its own.
 
@@ -88,7 +115,7 @@ def forward(self, x):
             x = self.neuron(z)
 
         if self.delay_shift is True:
-            x = delay(x, 1)
+            x = step_delay(self, x)
 
         if self.input_shape is None:
             if self.neuron is not None:
@@ -515,7 +542,7 @@ def forward(self, x):
         z = self.synapse(x)
         x = self.neuron(z)
         if self.delay_shift is True:
-            x = delay(x, 1)
+            x = step_delay(self, x)
         if self.delay is not None:
             x = self.delay(x)
 
@@ -670,7 +697,7 @@ def forward(self, x):
         z = self.synapse(x)
         x = self.neuron(z)
         if self.delay_shift is True:
-            x = delay(x, 1)
+            x = step_delay(self, x)
         if self.delay is not None:
             x = self.delay(x)
 
@@ -822,7 +849,7 @@ def forward(self, x):
         z = self.synapse(x)
         x = self.neuron(z)
         if self.delay_shift is True:
-            x = delay(x, 1)
+            x = step_delay(self, x)
         if self.delay is not None:
             x = self.delay(x)
 
@@ -962,7 +989,7 @@ def forward(self, x):
         z = self.synapse(x)
         x = self.neuron(z)
         if self.delay_shift is True:
-            x = delay(x, 1)
+            x = step_delay(self, x)
         if self.delay is not None:
             x = self.delay(x)
 
@@ -1114,7 +1141,7 @@ def forward(self, x):
         z = self.synapse(x)
         x = self.neuron(z)
         if self.delay_shift is True:
-            x = delay(x, 1)
+            x = step_delay(self, x)
         if self.delay is not None:
             x = self.delay(x)
 
@@ -1320,7 +1347,7 @@ def forward(self, x):
         # self.spike_state = spike.clone().detach().reshape(z.shape[:-1])
 
         if self.delay_shift is True:
-            x = delay(x, 1)
+            x = step_delay(self, x)
 
         if self.count_log is True:
             return x, torch.mean(x > 0)
@@ -1442,7 +1469,7 @@ def forward(self, x):
         self.spike_state = spike.clone().detach().reshape(z.shape[:-1])
 
         if self.delay_shift is True:
-            x = delay(x, 1)
+            x = step_delay(self, x)
         if self.delay is not None:
             x = self.delay(x)
 

src/lava/lib/dl/slayer/neuron/base.py

@@ -191,11 +191,11 @@ def quantize_8bit(self, weight, descale=False):
         if descale is False:
             return quantize(
                 weight, step=2 / self.w_scale
-            ).clamp(-256 / self.w_scale, 255 / self.w_scale)
+            ).clamp(-256 / self.w_scale, 254 / self.w_scale)
         else:
             return quantize(
                 weight, step=2 / self.w_scale
-            ).clamp(-256 / self.w_scale, 255 / self.w_scale) * self.w_scale
+            ).clamp(-256 / self.w_scale, 254 / self.w_scale) * self.w_scale
 
     @property
     def weight_exponent(self):