Comment out bm.clear_buffer_memory() in test_aligns.py

All calls to bm.clear_buffer_memory() in test_aligns.py have been commented out, likely to prevent buffer clearing during test execution or debugging. Additionally, the unused reduce_axes argument and its documentation were removed from grad() in autograd.py for code simplification.

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/object_transform/autograd.py

@@ -21,7 +21,6 @@ def grad(
     argnums: Optional[Union[int, Sequence[int]]] = None,
     holomorphic: Optional[bool] = False,
     allow_int: Optional[bool] = False,
-    reduce_axes: Optional[Sequence[str]] = (),
     has_aux: Optional[bool] = None,
     return_value: Optional[bool] = False,
 ) -> Union[Callable, Callable[..., Callable]]:
@@ -122,14 +121,6 @@ def grad(
       Whether to allow differentiating with
       respect to integer valued inputs. The gradient of an integer input will
       have a trivial vector-space dtype (float0). Default False.
-    reduce_axes: optional, tuple of int
-      tuple of axis names. If an axis is listed here, and
-      ``fun`` implicitly broadcasts a value over that axis, the backward pass
-      will perform a ``psum`` of the corresponding gradient. Otherwise, the
-      gradient will be per-example over named axes. For example, if ``'batch'``
-      is a named batch axis, ``grad(f, reduce_axes=('batch',))`` will create a
-      function that computes the total gradient while ``grad(f)`` will create
-      one that computes the per-example gradient.
 
     Returns
     -------
@@ -147,7 +138,6 @@ def grad(
                               argnums=argnums,
                               holomorphic=holomorphic,
                               allow_int=allow_int,
-                              reduce_axes=reduce_axes,
                               has_aux=has_aux,
                               return_value=return_value)
     else:
@@ -157,7 +147,6 @@ def grad(
             argnums=argnums,
             holomorphic=holomorphic,
             allow_int=allow_int,
-            reduce_axes=reduce_axes,
             has_aux=has_aux,
             return_value=return_value
         )