fix #626 (#631)

brainpy · Feb 22, 2024 · 5eb7cee · 5eb7cee
1 parent 4d74816
commit 5eb7cee
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/brainpy/_src/math/delayvars.py b/brainpy/_src/math/delayvars.py
@@ -11,7 +11,7 @@
 from brainpy import check
 from brainpy.check import is_float, is_integer, jit_error
 from brainpy.errors import UnsupportedError
-from .compat_numpy import vstack, broadcast_to
+from .compat_numpy import broadcast_to, expand_dims, concatenate
 from .environment import get_dt, get_float
 from .interoperability import as_jax
 from .ndarray import ndarray, Array
@@ -392,6 +392,7 @@ def reset(
                                      dtype=delay_target.dtype),
                            batch_axis=batch_axis)
     else:
+      self.data.value
       self.data._value = jnp.zeros((self.num_delay_step,) + delay_target.shape,
                                    dtype=delay_target.dtype)
 
@@ -472,7 +473,7 @@ def update(self, value: Union[numbers.Number, Array, jax.Array] = None):
 
     elif self.update_method == CONCAT_UPDATE:
       if self.num_delay_step >= 2:
-        self.data.value = vstack([broadcast_to(value, self.data.shape[1:]), self.data[1:]])
+        self.data.value = concatenate([expand_dims(value, 0), self.data[1:]], axis=0)
       else:
         self.data[:] = value