updates

Author: chaoming0625

brainpy/_src/dynsys.py

@@ -411,7 +411,7 @@ class Network(DynSysGroup):
   pass
 
 
-class Sequential(DynamicalSystem, AutoDelaySupp):
+class Sequential(DynamicalSystem, AutoDelaySupp, Container):
   """A sequential `input-output` module.
 
   Modules will be added to it in the order they are passed in the
@@ -468,22 +468,12 @@ def __init__(
       **modules_as_dict
   ):
     super().__init__(name=name, mode=mode)
-    self._dyn_modules = bm.NodeDict()
-    self._static_modules = dict()
-    i = 0
-    for m in modules_as_tuple + tuple(modules_as_dict.values()):
-      key = self.__format_key(i)
-      if isinstance(m, bm.BrainPyObject):
-        self._dyn_modules[key] = m
-      else:
-        self._static_modules[key] = m
-      i += 1
-    self._num = i
+    self.children = bm.node_dict(self.format_elements(object, *modules_as_tuple, **modules_as_dict))
 
   def update(self, x):
     """Update function of a sequential model.
     """
-    for m in self.__all_nodes():
+    for m in self.children.values():
       x = m(x)
     return x
 
@@ -494,15 +484,6 @@ def return_info(self):
                              f'not instance of {AutoDelaySupp.__name__}')
     return last.return_info()
 
-  def append(self, module: Callable):
-    assert isinstance(module, Callable)
-    key = self.__format_key(self._num)
-    if isinstance(module, bm.BrainPyObject):
-      self._dyn_modules[key] = module
-    else:
-      self._static_modules[key] = module
-    self._num += 1
-
   def __format_key(self, i):
     return f'l-{i}'
 
@@ -518,19 +499,17 @@ def __all_nodes(self):
 
   def __getitem__(self, key: Union[int, slice, str]):
     if isinstance(key, str):
-      if key in self._dyn_modules:
-        return self._dyn_modules[key]
-      elif key in self._static_modules:
-        return self._static_modules[key]
+      if key in self.children:
+        return self.children[key]
       else:
         raise KeyError(f'Does not find a component named {key} in\n {str(self)}')
     elif isinstance(key, slice):
-      return Sequential(*(self.__all_nodes()[key]))
+      return Sequential(**dict(tuple(self.children.items())[key]))
     elif isinstance(key, int):
-      return self.__all_nodes()[key]
+      return tuple(self.children.values())[key]
     elif isinstance(key, (tuple, list)):
-      _all_nodes = self.__all_nodes()
-      return Sequential(*[_all_nodes[k] for k in key])
+      _all_nodes = tuple(self.children.items())
+      return Sequential(**dict(_all_nodes[k] for k in key))
     else:
       raise KeyError(f'Unknown type of key: {type(key)}')
 
@@ -653,7 +632,7 @@ def init_variable(self, var_data, batch_or_mode, shape=None, sharding=None):
                      batch_axis_name=bm.sharding.BATCH_AXIS)
 
   def __repr__(self):
-    return f'{self.__class__.__name__}(name={self.name}, mode={self.mode}, size={self.size})'
+    return f'{self.name}(mode={self.mode}, size={self.size})'
 
   def __getitem__(self, item):
     return DynView(target=self, index=item)

brainpy/_src/math/object_transform/base.py

@@ -649,8 +649,8 @@ def __init__(self, seq=()):
     self.extend(seq)
 
   def append(self, element) -> 'NodeList':
-    if not isinstance(element, BrainPyObject):
-      raise TypeError(f'element must be an instance of {BrainPyObject.__name__}.')
+    # if not isinstance(element, BrainPyObject):
+    #   raise TypeError(f'element must be an instance of {BrainPyObject.__name__}.')
     super().append(element)
     return self
 
@@ -668,10 +668,10 @@ class NodeDict(dict):
   :py:func:`.vars()` operation in a :py:class:`~.BrainPyObject`.
   """
 
-  def _check_elem(self, elem):
-    if not isinstance(elem, BrainPyObject):
-      raise TypeError(f'Element should be {BrainPyObject.__name__}, but got {type(elem)}.')
-    return elem
+  # def _check_elem(self, elem):
+  #   if not isinstance(elem, BrainPyObject):
+  #     raise TypeError(f'Element should be {BrainPyObject.__name__}, but got {type(elem)}.')
+  #   return elem
 
   def __init__(self, *args, **kwargs):
     super().__init__()
@@ -690,7 +690,7 @@ def update(self, *args, **kwargs) -> 'VarDict':
     return self
 
   def __setitem__(self, key, value) -> 'VarDict':
-    super().__setitem__(key, self._check_elem(value))
+    super().__setitem__(key, value)
     return self
 
 

brainpy/_src/math/object_transform/naming.py

@@ -32,7 +32,7 @@ def check_name_uniqueness(name, obj):
     _name2id[name] = id(obj)
 
 
-def get_unique_name(type_):
+def get_unique_name(type_: str):
   """Get the unique name for the given object type."""
   if type_ not in _typed_names:
     _typed_names[type_] = 0

brainpy/_src/mixin.py

@@ -9,6 +9,7 @@
 import numpy as np
 
 from brainpy import math as bm, tools
+from brainpy._src.math.object_transform.naming import get_unique_name
 from brainpy._src.initialize import parameter
 from brainpy.types import ArrayType
 
@@ -192,19 +193,25 @@ def __repr__(self):
     string = ", \n".join(child_str)
     return f'{cls_name}({string})'
 
+  def __get_elem_name(self, elem):
+    if isinstance(elem, bm.BrainPyObject):
+      return elem.name
+    else:
+      return get_unique_name('ContainerElem')
+
   def format_elements(self, child_type: type, *children_as_tuple, **children_as_dict):
     res = dict()
 
     # add tuple-typed components
     for module in children_as_tuple:
       if isinstance(module, child_type):
-        res[module.name] = module
+        res[self.__get_elem_name(module)] = module
       elif isinstance(module, (list, tuple)):
         for m in module:
           if not isinstance(m, child_type):
             raise ValueError(f'Should be instance of {child_type.__name__}. '
                              f'But we got {type(m)}')
-          res[m.name] = m
+          res[self.__get_elem_name(m)] = m
       elif isinstance(module, dict):
         for k, v in module.items():
           if not isinstance(v, child_type):
@@ -226,12 +233,12 @@ def add_elem(self, **elements):
     """Add new elements.
 
     >>> obj = Container()
-    >>> obj.add_elem(1.)
+    >>> obj.add_elem(a=1.)
 
     Args:
       elements: children objects.
     """
-    self.check_hierarchies(type(self), **elements)
+    # self.check_hierarchies(type(self), **elements)
     self.children.update(self.format_elements(object, **elements))
 
 

brainpy/_src/runners.py

@@ -379,7 +379,6 @@ def __repr__(self):
   def reset_state(self):
     """Reset state of the ``DSRunner``."""
     self.i0 = 0
-    self.t0 = self.t0
 
   def predict(
       self,