fix bugs in LengthDelay and VariableView

Author: chaoming0625

brainpy/math/delayvars.py

@@ -435,7 +435,7 @@ def update(self, value: Union[float, int, bool, JaxArray, jnp.DeviceArray]):
       self.idx.value = stop_gradient((self.idx + 1) % self.num_delay_step)
 
     elif self.update_method == CONCAT_UPDATING:
-      self.data.value = bm.concatenate([self.data[1:], bm.broadcast_to(value, self.delay_target_shape)], axis=0)
+      self.data.value = bm.vstack([self.data[1:], bm.broadcast_to(value,self.data.shape[1:])])
 
     else:
       raise ValueError(f'Unknown updating method "{self.update_method}"')

brainpy/math/jaxarray.py

@@ -1673,3 +1673,23 @@ def update(self, value):
                       f"while we got {value.dtype}.")
     self._value[self.index] = value.value if isinstance(value, JaxArray) else value
 
+  @value.setter
+  def value(self, value):
+    int_shape = self.shape
+    if self.batch_axis is None:
+      ext_shape = value.shape
+    else:
+      ext_shape = value.shape[:self.batch_axis] + value.shape[self.batch_axis + 1:]
+      int_shape = int_shape[:self.batch_axis] + int_shape[self.batch_axis + 1:]
+    if ext_shape != int_shape:
+      error = f"The shape of the original data is {int_shape}, while we got {value.shape}"
+      if self.batch_axis is None:
+        error += '. Do you forget to set "batch_axis" when initialize this variable?'
+      else:
+        error += f' with batch_axis={self.batch_axis}.'
+      raise MathError(error)
+    if value.dtype != self._value.dtype:
+      raise MathError(f"The dtype of the original data is {self._value.dtype}, "
+                      f"while we got {value.dtype}.")
+    self._value[self.index] = value.value if isinstance(value, JaxArray) else value
+

brainpy/math/tests/test_jaxarray.py

@@ -58,3 +58,36 @@ def test_variable_init(self):
       not bm.array_equal(bm.Variable(bm.random.rand(10)),
                          bm.Variable(10))
     )
+
+
+class TestVariableView(unittest.TestCase):
+  def test_update(self):
+    origin = bm.Variable(bm.zeros(10))
+    view = bm.VariableView(origin, slice(0, 5, None))
+
+    view.update(bm.ones(5))
+    self.assertTrue(
+      bm.array_equal(origin, bm.concatenate([bm.ones(5), bm.zeros(5)]))
+    )
+
+    view.value = bm.arange(5.)
+    self.assertTrue(
+      bm.array_equal(origin, bm.concatenate([bm.arange(5), bm.zeros(5)]))
+    )
+
+    view += 10
+    self.assertTrue(
+      bm.array_equal(origin, bm.concatenate([bm.arange(5) + 10, bm.zeros(5)]))
+    )
+
+    bm.random.shuffle(view)
+    print(view)
+    print(origin)
+
+    view.sort()
+    self.assertTrue(
+      bm.array_equal(origin, bm.concatenate([bm.arange(5) + 10, bm.zeros(5)]))
+    )
+
+    self.assertTrue(view.sum() == bm.sum(bm.arange(5) + 10))
+