Refactor Base to BaseArray and update usages

Renamed the Base class in ndarray.py to BaseArray and updated all references in related modules, including Array and Variable. Also fixed the while_loop wrapper in controls.py to correctly unpack operands for the body and condition functions. Uncommented bm.clear_buffer_memory() calls in test_aligns.py to ensure buffer memory is cleared after tests.

Author: Routhleck

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/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), 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: