fix bugs in delay vars

Author: chaoming0625

brainpy/_src/dyn/base.py

@@ -134,7 +134,7 @@ def __init__(
     self.local_delay_vars: Dict[str, bm.LengthDelay] = Collector()
 
     # super initialization
-    super(DynamicalSystem, self).__init__(name=name)
+    BrainPyObject.__init__(self, name=name)
 
   @property
   def mode(self) -> bm.Mode:
@@ -155,7 +155,8 @@ def __call__(self, *args, **kwargs):
     """The shortcut to call ``update`` methods."""
     if hasattr(self.update, '_new_style') and getattr(self.update, '_new_style'):
       if len(args) and isinstance(args[0], dict):
-        bm.share.save(**args[0])
+        for k, v in args[0].items():
+          bm.share.save(k, v)
         return self.update(*args[1:], **kwargs)
       else:
         return self.update(*args, **kwargs)

brainpy/_src/dyn/runners.py

@@ -615,7 +615,8 @@ def _step_func_predict(self, shared_args, t, i, x):
     # input step
     shared = tools.DotDict(t=t, i=i, dt=self.dt)
     shared.update(shared_args)
-    bm.share.save(**shared)
+    for k, v in shared.items():
+      bm.share.save(k, v)
     self.target.clear_input()
     self._step_func_input(shared)
 
@@ -630,7 +631,6 @@ def _step_func_predict(self, shared_args, t, i, x):
     # finally
     if self.progress_bar:
       id_tap(lambda *arg: self._pbar.update(), ())
-    bm.share.remove_shargs()
     return out, mon
 
   def _get_f_predict(self, shared_args: Dict = None):

brainpy/_src/experimental/delay.py

@@ -21,18 +21,18 @@ def __init__(
       length: int = 0,
       before_t0: Union[float, int, bool, bm.Array, jax.Array, Callable] = None,
       entries: Optional[Dict] = None,
-      method: str = None,
+      method: str = ROTATE_UPDATE,
       mode: bm.Mode = None,
       name: str = None,
   ):
+    DynamicalSystem.__init__(self, mode=mode)
     if method is None:
       if self.mode.is_a(bm.NonBatchingMode):
         method = ROTATE_UPDATE
       elif self.mode.is_parent_of(bm.TrainingMode):
         method = CONCAT_UPDATE
       else:
         method = ROTATE_UPDATE
-    DynamicalSystem.__init__(self, mode=mode)
     DelayVariable.__init__(self,
                            target=target,
                            length=length,

brainpy/_src/math/context.py

@@ -64,24 +64,24 @@ def save(self, identifier: str, data: Any) -> None:
     """Save shared arguments in the global context."""
     assert isinstance(identifier, str)
 
-    if identifier in self._identifiers:
-      raise ValueError(f'{identifier} has been used. Please assign another name.')
     if isinstance(data, DelayVariable):
+      if identifier in self._identifiers:
+        raise ValueError(f'{identifier} has been used. Please assign another name.')
       self._delays[identifier] = data
-    elif isinstance(data, DelayEntry):
-      if isinstance(data.target, DelayVariable):
-        delay_key = f'delay{id(data)}'
-        self.save(delay_key, data.target)
-        delay = data.target
-      elif isinstance(data.target, str):
-        if data.target not in self._delays:
-          raise ValueError(f'Delay target {data.target} has not been registered.')
-        delay = self._delays[data.target]
-        delay_key = data.target
-      else:
-        raise ValueError(f'Unknown delay target. {type(data.target)}')
-      delay.register_entry(identifier, delay_time=data.time, delay_step=data.step)
-      self._delay_entries[identifier] = delay_key
+    # elif isinstance(data, DelayEntry):
+    #   if isinstance(data.target, DelayVariable):
+    #     delay_key = f'delay{id(data)}'
+    #     self.save(delay_key, data.target)
+    #     delay = data.target
+    #   elif isinstance(data.target, str):
+    #     if data.target not in self._delays:
+    #       raise ValueError(f'Delay target {data.target} has not been registered.')
+    #     delay = self._delays[data.target]
+    #     delay_key = data.target
+    #   else:
+    #     raise ValueError(f'Unknown delay target. {type(data.target)}')
+    #   delay.register_entry(identifier, delay_time=data.time, delay_step=data.step)
+    #   self._delay_entries[identifier] = delay_key
     else:
       self._arguments[identifier] = data
     self._identifiers.add(identifier)

brainpy/_src/math/delayvars.py

@@ -519,7 +519,7 @@ def __init__(
       name: str = None,
       method: str = ROTATE_UPDATE,
   ):
-    super().__init__(name=name)
+    BrainPyObject.__init__(self, name=name)
     assert method in [ROTATE_UPDATE, CONCAT_UPDATE]
     self.method = method
 

brainpy/_src/train/back_propagation.py

@@ -566,7 +566,8 @@ def _step_func_fit(self, shared_args, inputs, targets):
 
   def _step_func_predict(self, shared, x=None):
     assert self.data_first_axis == 'B', f'There is no time dimension when using the trainer {self.__class__.__name__}.'
-    bm.share.save(**shared)
+    for k, v in shared.items():
+      bm.share.save(k, v)
 
     # input step
     self.target.clear_input()

brainpy/math/activations.py

@@ -24,5 +24,4 @@
   swish as swish,
   selu as selu,
   identity as identity,
-  tanh as tanh,
 )