Merge branch 'update' of https://github.com/brainpy/BrainPy into update

# Conflicts:
#	brainpy/_src/math/object_transform/controls.py

Author: chaoming0625

brainpy/_src/dyn/projections/tests/test_aligns.py

@@ -74,7 +74,7 @@ def update(self):
   bp.visualize.raster_plot(indices * bm.dt, spks, show=True)
 
   plt.close()
-  # bm.clear_buffer_memory()
+  bm.clear_buffer_memory()
 
 
 def test_ProjAlignPostMg2():
@@ -150,7 +150,7 @@ def update(self):
   bp.visualize.raster_plot(indices * bm.dt, spks, show=True)
 
   plt.close()
-  # bm.clear_buffer_memory()
+  bm.clear_buffer_memory()
 
 
 def test_ProjAlignPost1():
@@ -185,7 +185,7 @@ def update(self, input):
   indices = bm.arange(400)
   spks = bm.for_loop(lambda i: model.step_run(i, 20.), indices)
   bp.visualize.raster_plot(indices, spks, show=True)
-  # bm.clear_buffer_memory()
+  bm.clear_buffer_memory()
   plt.close()
 
 
@@ -244,7 +244,7 @@ def update(self, inp):
   spks = bm.for_loop(lambda i: model.step_run(i, 20.), indices)
   bp.visualize.raster_plot(indices, spks, show=True)
 
-  # bm.clear_buffer_memory()
+  bm.clear_buffer_memory()
   plt.close()
 
 
@@ -280,7 +280,7 @@ def update(self, input):
   indices = bm.arange(400)
   spks = bm.for_loop(lambda i: model.step_run(i, 20.), indices)
   bp.visualize.raster_plot(indices, spks, show=True)
-  # bm.clear_buffer_memory()
+  bm.clear_buffer_memory()
   plt.close()
 
 
@@ -338,7 +338,7 @@ def update(self, inp):
   spks = bm.for_loop(lambda i: model.step_run(i, 20.), indices)
   bp.visualize.raster_plot(indices, spks, show=True)
 
-  # bm.clear_buffer_memory()
+  bm.clear_buffer_memory()
   plt.close()
 
 
@@ -396,7 +396,7 @@ def update(self, inp):
   spks = bm.for_loop(lambda i: model.step_run(i, 20.), indices)
   bp.visualize.raster_plot(indices, spks, show=True)
 
-  # bm.clear_buffer_memory()
+  bm.clear_buffer_memory()
   plt.close()
 
 
@@ -437,4 +437,4 @@ def update(self, input):
   spks = bm.for_loop(lambda i: model.step_run(i, 20.), indices, progress_bar=True)
   bp.visualize.raster_plot(indices, spks, show=True)
   plt.close()
-  # bm.clear_buffer_memory()
+  bm.clear_buffer_memory()

brainpy/_src/math/delayvars.py

@@ -14,7 +14,7 @@
 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
+from .ndarray import ndarray, BaseArray
 from .object_transform.base import BrainPyObject
 from .object_transform.controls import cond
 from .object_transform.variables import Variable
@@ -29,7 +29,7 @@
 
 
 def _as_jax_array(arr):
-  return arr.value if isinstance(arr, Array) else arr
+  return arr.value if isinstance(arr, BaseArray) else arr
 
 
 class AbstractDelay(BrainPyObject):
@@ -129,8 +129,8 @@ def __init__(
     super(TimeDelay, self).__init__(name=name)
 
     # shape
-    if not isinstance(delay_target, (jnp.ndarray, Array)):
-      raise ValueError(f'Must be an instance of Array or jax.numpy.ndarray. But we got {type(delay_target)}')
+    if not isinstance(delay_target, (jnp.ndarray, BaseArray)):
+      raise ValueError(f'Must be an instance of BaseArray or jax.numpy.ndarray. But we got {type(delay_target)}')
 
     # delay_len
     self.t0 = t0
@@ -453,7 +453,7 @@ def retrieve(self, delay_len, *indices):
     # the delay data
     return self.data[indices]
 
-  def update(self, value: Union[numbers.Number, Array, jax.Array] = None):
+  def update(self, value: Union[numbers.Number, BaseArray, jax.Array] = None):
     """Update delay variable with the new data.
 
     Parameters

brainpy/_src/math/ndarray.py

@@ -63,7 +63,7 @@ def _get_dtype(v):
     dtype = canonicalize_dtype(type(v))
   return dtype
 
-class Base:
+class BaseArray:
   @property
   def sharding(self):
     return self._value.sharding
@@ -1488,7 +1488,7 @@ def double(self): return jnp.asarray(self.value, dtype=jnp.float64)
 
 
 @register_pytree_node_class
-class Array(Base):
+class Array(BaseArray):
   """Multiple-dimensional array in BrainPy.
 
   Compared to ``jax.Array``, :py:class:`~.Array` has the following advantages:

brainpy/_src/math/object_transform/controls.py

@@ -381,4 +381,4 @@ def while_loop(
     """
     if not isinstance(operands, (tuple, list)):
         operands = (operands,)
-    return brainstate.transform.while_loop(body_fun, cond_fun, *operands)
+    return brainstate.transform.while_loop(lambda x: body_fun(*x), lambda x: cond_fun(*x), operands)

brainpy/_src/math/object_transform/variables.py

@@ -6,7 +6,7 @@
 from jax.dtypes import canonicalize_dtype
 from jax.tree_util import register_pytree_node_class
 
-from brainpy._src.math.ndarray import Array, Base
+from brainpy._src.math.ndarray import Array, BaseArray
 from brainpy._src.math.sharding import BATCH_AXIS
 from brainpy.errors import MathError
 
@@ -224,7 +224,7 @@ def __add__(self, other: dict):
 
 
 @register_pytree_node_class
-class Variable(brainstate.State, Base):
+class Variable(brainstate.State, BaseArray):
   """The pointer to specify the dynamical variable.
 
   Initializing an instance of ``Variable`` by two ways:

brainpy/_src/math/random.py

@@ -546,8 +546,20 @@ def seed(self, seed_or_key=None, seed=None):
       key = seed_or_key
     self._value = key
 
+
+  def _ensure_value_exists(self):
+    """Ensure that the random state has a valid value, regenerate if needed."""
+    if not isinstance(self._value, np.ndarray):
+      with jax.ensure_compile_time_eval():
+        if not isinstance(self._value, jax.core.Tracer):
+          if self._value.is_deleted():
+            seed_or_key = np.random.randint(0, 100000, 2, dtype=np.uint32)
+            self._value = seed_or_key
+
+
   @property
   def value(self):
+    self._ensure_value_exists()
     record_state_value_read(self)
     return self._read_value()