support brainpy.math.for_loop with the keyword unroll_kwargs

Author: chaoming0625

brainpy/_src/math/object_transform/controls.py

@@ -723,11 +723,12 @@ def _get_for_loop_transform(
     remat: bool,
     reverse: bool,
     unroll: int,
+    unroll_kwargs: tools.DotDict
 ):
   def fun2scan(carry, x):
     for k in dyn_vars.keys():
       dyn_vars[k]._value = carry[k]
-    results = body_fun(*x)
+    results = body_fun(*x, **unroll_kwargs)
     if progress_bar:
       id_tap(lambda *arg: bar.update(), ())
     return dyn_vars.dict_data(), results
@@ -860,6 +861,7 @@ def for_loop(
 
   if unroll_kwargs is None:
     unroll_kwargs = dict()
+  unroll_kwargs = tools.DotDict(unroll_kwargs)
 
   if not isinstance(operands, (list, tuple)):
     operands = (operands,)
@@ -871,19 +873,20 @@ def for_loop(
 
   if jit is None:  # jax disable jit
     jit = not jax.config.jax_disable_jit
-  dyn_vars = get_stack_cache(body_fun)
+  dyn_vars = get_stack_cache((body_fun, unroll_kwargs))
   if jit:
     if dyn_vars is None:
       # TODO: better cache mechanism?
       with new_transform('for_loop'):
         with VariableStack() as dyn_vars:
           transform = _get_for_loop_transform(body_fun, VariableStack(), bar,
-                                              progress_bar, remat, reverse, unroll)
+                                              progress_bar, remat, reverse, unroll,
+                                              unroll_kwargs)
           if current_transform_number() > 1:
             rets = transform(operands)
           else:
             rets = jax.eval_shape(transform, operands)
-      cache_stack(body_fun, dyn_vars)  # cache
+      cache_stack((body_fun, unroll_kwargs), dyn_vars)  # cache
       if current_transform_number():
         return rets[1]
       del rets
@@ -893,7 +896,7 @@ def for_loop(
   # TODO: cache mechanism?
   transform = _get_for_loop_transform(body_fun, dyn_vars, bar,
                                       progress_bar, remat, reverse,
-                                      unroll)
+                                      unroll, unroll_kwargs)
   if jit:
     dyn_vals, out_vals = transform(operands)
   else:

brainpy/_src/runners.py

@@ -466,7 +466,7 @@ def predict(
       inputs = tree_map(lambda x: jnp.moveaxis(x, 0, 1), inputs)
 
     # build monitor
-    for key in self.mon.var_names:
+    for key in self._monitors.keys():
       self.mon[key] = []  # reshape the monitor items
 
     # init progress bar
@@ -492,7 +492,7 @@ def predict(
     # post-running for monitors
     if self._memory_efficient:
       self.mon['ts'] = indices * self.dt + self.t0
-      for key in self.mon.var_names:
+      for key in self._monitors.keys():
         self.mon[key] = np.asarray(self.mon[key])
     else:
       hists['ts'] = indices * self.dt + self.t0
@@ -658,6 +658,7 @@ def _fun_predict(self, indices, *inputs, shared_args=None):
       return outs, None
 
     else:
-      return bm.for_loop(functools.partial(self._step_func_predict, shared_args=shared_args),
+      return bm.for_loop(self._step_func_predict,
                          (indices, *inputs),
-                         jit=self.jit['predict'])
+                         jit=self.jit['predict'],
+                         unroll_kwargs={'shared_args': shared_args})

brainpy/_src/running/runner.py

@@ -118,7 +118,6 @@ def __init__(
 
     # monitor for user access
     self.mon = DotDict()
-    self.mon['var_names'] = tuple(self._monitors.keys())
 
     # progress bar
     assert isinstance(progress_bar, bool), 'Must be a boolean variable.'

brainpy/_src/tools/dicts.py

@@ -42,64 +42,20 @@ class DotDict(dict):
   >>> f(d)
   TypeError: Argument 'a' of type <class 'str'> is not a valid JAX type.
 
-  At this moment, you can label this attribute `names` as not a key in the dictionary
-  by using the syntax::
-
-  >>> d.add_attr_not_key('names')
-  >>> f(d)
-  {'a': DeviceArray(10, dtype=int32, weak_type=True),
-   'b': DeviceArray(20, dtype=int32, weak_type=True),
-   'c': DeviceArray(30, dtype=int32, weak_type=True)}
-
   """
 
-  '''Used to exclude variables that '''
-  attrs_not_keys = ('attrs_not_keys', 'var_names')
-
   def __init__(self, *args, **kwargs):
     super().__init__(*args, **kwargs)
     self.__dict__ = self
-    self.var_names = ()
 
   def copy(self) -> 'DotDict':
     return type(self)(super().copy())
 
-  def keys(self):
-    """Retrieve all keys in the dict, excluding ignored keys."""
-    keys = []
-    for k in super(DotDict, self).keys():
-      if k not in self.attrs_not_keys:
-        keys.append(k)
-    return tuple(keys)
-
-  def values(self):
-    """Retrieve all values in the dict, excluding values of ignored keys."""
-    values = []
-    for k, v in super(DotDict, self).items():
-      if k not in self.attrs_not_keys:
-        values.append(v)
-    return tuple(values)
-
-  def items(self):
-    """Retrieve all items in the dict, excluding ignored items."""
-    items = []
-    for k, v in super(DotDict, self).items():
-      if k not in self.attrs_not_keys:
-        items.append((k, v))
-    return items
-
   def to_numpy(self):
     """Change all values to numpy arrays."""
     for key in tuple(self.keys()):
       self[key] = np.asarray(self[key])
 
-  def add_attr_not_key(self, *args):
-    """Add excluded attribute when retrieving dictionary keys. """
-    for arg in args:
-      if not isinstance(arg, str):
-        raise TypeError('Only support string.')
-    self.attrs_not_keys += args
-
   def update(self, *args, **kwargs):
     super().update(*args, **kwargs)
     return self
@@ -179,7 +135,7 @@ def subset(self, var_type):
 
     >>> import brainpy as bp
     >>>
-    >>> some_collector = Collector()
+    >>> some_collector = DotDict()
     >>>
     >>> # get all trainable variables
     >>> some_collector.subset(bp.math.TrainVar)

brainpy/_src/train/back_propagation.py

@@ -605,7 +605,7 @@ def predict(
       self.target.reset_state(self._get_input_batch_size(xs=inputs))
       self.reset_state()
     # init monitor
-    for key in self.mon.var_names:
+    for key in self._monitors.keys():
       self.mon[key] = []  # reshape the monitor items
     # prediction
     if not isinstance(inputs, (tuple, list)):

brainpy/_src/train/online.py

@@ -177,7 +177,7 @@ def fit(
                     is_leaf=lambda y: isinstance(y, bm.Array))
 
     # init monitor
-    for key in self.mon.var_names:
+    for key in self._monitors.keys():
       self.mon[key] = []  # reshape the monitor items
 
     # init progress bar

examples/dynamics_simulation/COBA.py

@@ -168,7 +168,6 @@ def run3():
   bp.visualize.raster_plot(runner.mon.ts, runner.mon['E.spike'], show=True)
 
 
-
 def run4():
   net = EICOBA_PostAlign(3200, 800)
   runner = bp.DSRunner(net, monitors={'E.spike': net.E.spike})

tests/simulation/test_net_COBA.py

@@ -6,17 +6,17 @@
 class ESN(bp.DynamicalSystem):
   def __init__(self, num_in, num_hidden, num_out):
     super(ESN, self).__init__()
-    self.r = bp.layers.Reservoir(num_in,
-                                 num_hidden,
-                                 Win_initializer=bp.init.Uniform(-0.1, 0.1),
-                                 Wrec_initializer=bp.init.Normal(scale=0.1),
-                                 in_connectivity=0.02,
-                                 rec_connectivity=0.02,
-                                 comp_type='dense')
-    self.o = bp.layers.Dense(num_hidden,
-                             num_out,
-                             W_initializer=bp.init.Normal(),
-                             mode=bm.training_mode)
+    self.r = bp.dnn.Reservoir(num_in,
+                              num_hidden,
+                              Win_initializer=bp.init.Uniform(-0.1, 0.1),
+                              Wrec_initializer=bp.init.Normal(scale=0.1),
+                              in_connectivity=0.02,
+                              rec_connectivity=0.02,
+                              comp_type='dense')
+    self.o = bp.dnn.Dense(num_hidden,
+                          num_out,
+                          W_initializer=bp.init.Normal(),
+                          mode=bm.training_mode)
 
   def update(self, x):
     return x >> self.r >> self.o
@@ -26,10 +26,10 @@ class NGRC(bp.DynamicalSystem):
   def __init__(self, num_in, num_out):
     super(NGRC, self).__init__()
 
-    self.r = bp.layers.NVAR(num_in, delay=2, order=2)
-    self.o = bp.layers.Dense(self.r.num_out, num_out,
-                             W_initializer=bp.init.Normal(0.1),
-                             mode=bm.training_mode)
+    self.r = bp.dnn.NVAR(num_in, delay=2, order=2)
+    self.o = bp.dnn.Dense(self.r.num_out, num_out,
+                          W_initializer=bp.init.Normal(0.1),
+                          mode=bm.training_mode)
 
   def update(self, x):
     return x >> self.r >> self.o