updates

Author: chaoming0625

brainpy/__init__.py

@@ -65,7 +65,7 @@
 Network = DynSysGroup
 # delays
 from brainpy._src.delay import (
-  VariDelay as VariDelay,
+  VarDelay as VarDelay,
 )
 
 # building blocks
@@ -129,12 +129,16 @@
 from brainpy._add_deprecations import deprecation_getattr2
 
 __deprecations = {
+  'Module': ('brainpy.Module', 'brainpy.DynamicalSystem', DynamicalSystem),
+  'Channel': ('brainpy.Channel', 'brainpy.dyn.IonChannel', dyn.IonChannel),
+  'NeuGroup': ('brainpy.NeuGroup', 'brainpy.dyn.NeuDyn', dyn.NeuDyn),
+  'SynConn': ('brainpy.SynConn', 'brainpy.dyn.SynConn', dyn.SynConn),
   'Container': ('brainpy.Container', 'brainpy.DynSysGroup', DynSysGroup),
+
   'optimizers': ('brainpy.optimizers', 'brainpy.optim', optim),
   'TensorCollector': ('brainpy.TensorCollector', 'brainpy.ArrayCollector', ArrayCollector),
   'SynSTP': ('brainpy.SynSTP', 'brainpy.synapses.SynSTP', synapses.SynSTP),
   'SynOut': ('brainpy.SynOut', 'brainpy.synapses.SynOut', synapses.SynOut),
-  'SynConn': ('brainpy.SynConn', 'brainpy.dyn.SynConn', dyn.SynConn),
   'TwoEndConn': ('brainpy.TwoEndConn', 'brainpy.synapses.TwoEndConn', synapses.TwoEndConn),
   'CondNeuGroup': ('brainpy.CondNeuGroup', 'brainpy.syn.CondNeuGroup', dyn.CondNeuGroup),
 }

brainpy/_src/delay.py

@@ -21,7 +21,7 @@
 
 __all__ = [
   'Delay',
-  'VariDelay',
+  'VarDelay',
   'DataDelay',
   'DelayAccess',
 ]
@@ -432,7 +432,7 @@ def _check_target_sharding(sharding, ndim, mode: bm.Mode):
   return sharding
 
 
-class VariDelay(Delay):
+class VarDelay(Delay):
   """Generate Delays for the given :py:class:`~.Variable` instance.
 
   The data in this delay variable is arranged as::
@@ -690,7 +690,7 @@ def _init_data(self, length: int, batch_size: int = None):
       self.data[:] = self._init((length,) + self.target.shape, dtype=self.target.dtype)
 
 
-class DataDelay(VariDelay):
+class DataDelay(VarDelay):
   not_desc_params = ('time', 'entries')
 
   def __init__(

brainpy/_src/deprecations.py

@@ -8,6 +8,40 @@
 ]
 
 
+_update_deprecate_msg = '''
+From brainpy>=2.4.3, update() function no longer needs to receive a global shared argument.
+
+Instead of using:
+
+  def update(self, tdi, *args, **kwagrs):
+     t = tdi['t']
+     ...
+
+Please use:
+
+  def update(self, *args, **kwagrs):
+     t = bp.share['t']
+     ...
+'''
+
+
+_input_deprecate_msg = '''
+From brainpy>=2.4.3, input() function no longer needs to receive a global shared argument.
+
+Instead of using:
+
+  def input(tdi):
+     ...
+
+Please use:
+
+  def input():
+     t = bp.share['t']
+     ...
+'''
+
+
+
 def _deprecate(msg):
   warnings.simplefilter('always', DeprecationWarning)  # turn off filter
   warnings.warn(msg, category=DeprecationWarning, stacklevel=2)

brainpy/_src/dyn/projections/aligns.py

@@ -3,7 +3,7 @@
 import jax
 
 from brainpy import math as bm, check
-from brainpy._src.delay import Delay, VariDelay, DataDelay, DelayAccess
+from brainpy._src.delay import Delay, VarDelay, DataDelay, DelayAccess
 from brainpy._src.dynsys import DynamicalSystem, Projection, Dynamic, Sequential
 from brainpy._src.mixin import JointType, ParamDescInit, ReturnInfo, AutoDelaySupp, BindCondData, AlignPost
 
@@ -54,7 +54,7 @@ def update(self):
 
 def _init_delay(info: Union[bm.Variable, ReturnInfo]) -> Delay:
   if isinstance(info, bm.Variable):
-    return VariDelay(info)
+    return VarDelay(info)
   elif isinstance(info, ReturnInfo):
     if isinstance(info.batch_or_mode, int):
       shape = (info.batch_or_mode,) + tuple(info.size)
@@ -106,7 +106,7 @@ def __init__(self):
           super().__init__()
           self.N = bp.dyn.LifRef(4000, V_rest=-60., V_th=-50., V_reset=-60., tau=20., tau_ref=5.,
                                  V_initializer=bp.init.Normal(-55., 2.))
-          self.delay = bp.VariableDelay(self.N.spike, entries={'I': None})
+          self.delay = bp.VarDelay(self.N.spike, entries={'I': None})
           self.syn1 = bp.dyn.Expon(size=3200, tau=5.)
           self.syn2 = bp.dyn.Expon(size=800, tau=10.)
           self.E = bp.dyn.VanillaProj(comm=bp.dnn.EventJitFPHomoLinear(3200, 4000, prob=0.02, weight=0.6),
@@ -180,7 +180,7 @@ def __init__(self):
         super().__init__()
         self.N = bp.dyn.LifRef(4000, V_rest=-60., V_th=-50., V_reset=-60., tau=20., tau_ref=5.,
                                V_initializer=bp.init.Normal(-55., 2.))
-        self.delay = bp.VariableDelay(self.N.spike, entries={'I': None})
+        self.delay = bp.VarDelay(self.N.spike, entries={'I': None})
         self.E = bp.dyn.ProjAlignPostMg1(comm=bp.dnn.EventJitFPHomoLinear(3200, 4000, prob=0.02, weight=0.6),
                                          syn=bp.dyn.Expon.desc(size=4000, tau=5.),
                                          out=bp.dyn.COBA.desc(E=0.),
@@ -374,7 +374,7 @@ def __init__(self):
           super().__init__()
           self.N = bp.dyn.LifRef(4000, V_rest=-60., V_th=-50., V_reset=-60., tau=20., tau_ref=5.,
                                  V_initializer=bp.init.Normal(-55., 2.))
-          self.delay = bp.VariableDelay(self.N.spike, entries={'I': None})
+          self.delay = bp.VarDelay(self.N.spike, entries={'I': None})
           self.E = bp.dyn.ProjAlignPost1(comm=bp.dnn.EventJitFPHomoLinear(3200, 4000, prob=0.02, weight=0.6),
                                          syn=bp.dyn.Expon(size=4000, tau=5.),
                                          out=bp.dyn.COBA(E=0.),

brainpy/_src/dynsys.py

@@ -13,6 +13,7 @@
 from brainpy._src.mixin import AutoDelaySupp, Container, DelayRegister, global_delay_data
 from brainpy.errors import NoImplementationError, UnsupportedError
 from brainpy.types import ArrayType, Shape
+from brainpy._src.deprecations import _update_deprecate_msg
 
 share = None
 
@@ -29,21 +30,6 @@
 
 SLICE_VARS = 'slice_vars'
 
-_update_deprecate_msg = '''
-From brainpy>=2.4.3, update() function no longer needs to receive a global shared argument.
-
-Instead of using:
-          
-  def update(self, tdi, *args, **kwagrs):
-     ...
-          
-Please use:
-          
-  def update(self, *args, **kwagrs):
-     t = bp.share['t']
-     ...
-'''
-
 
 def not_pass_shared(func: Callable):
   """Label the update function as the one without passing shared arguments.
@@ -305,14 +291,14 @@ def __repr__(self):
   def __call__(self, *args, **kwargs):
     """The shortcut to call ``update`` methods."""
 
-    # update ``before_updates``
+    # ``before_updates``
     for model in self.before_updates.values():
       model()
 
     # update the model self
     ret = self.update(*args, **kwargs)
 
-    # update ``after_updates``
+    # ``after_updates``
     for model in self.after_updates.values():
       model(ret)
     return ret

brainpy/_src/math/ndarray.py

@@ -748,7 +748,7 @@ def split(self, indices_or_sections, axis=0):
     sub-arrays : list of ndarrays
       A list of sub-arrays as views into `ary`.
     """
-    return [_return(a) for a in self.value.split(indices_or_sections, axis=axis)]
+    return [_return(a) for a in jnp.split(self.value, indices_or_sections, axis=axis)]
 
   def take(self, indices, axis=None, mode=None):
     """Return an array formed from the elements of a at the given indices."""

brainpy/_src/math/object_transform/controls.py

@@ -722,7 +722,7 @@ def _get_for_loop_transform(
     progress_bar: bool,
     remat: bool,
     reverse: bool,
-    unroll: int
+    unroll: int,
 ):
   def fun2scan(carry, x):
     for k in dyn_vars.keys():
@@ -753,6 +753,7 @@ def for_loop(
     remat: bool = False,
     jit: Optional[bool] = None,
     progress_bar: bool = False,
+    unroll_kwargs: Optional[Dict] = None,
 
     # deprecated
     dyn_vars: Union[Variable, Sequence[Variable], Dict[str, Variable]] = None,
@@ -845,6 +846,8 @@ def for_loop(
     .. deprecated:: 2.4.0
        No longer need to provide ``child_objs``. This function is capable of automatically
        collecting the children objects used in the target ``func``.
+  unroll_kwargs: dict
+    The keyword arguments without unrolling.
 
   Returns
   -------
@@ -855,6 +858,9 @@ def for_loop(
   dynvar_deprecation(dyn_vars)
   node_deprecation(child_objs)
 
+  if unroll_kwargs is None:
+    unroll_kwargs = dict()
+
   if not isinstance(operands, (list, tuple)):
     operands = (operands,)
 
@@ -885,7 +891,9 @@ def for_loop(
     dyn_vars = VariableStack()
 
   # TODO: cache mechanism?
-  transform = _get_for_loop_transform(body_fun, dyn_vars, bar, progress_bar, remat, reverse, unroll)
+  transform = _get_for_loop_transform(body_fun, dyn_vars, bar,
+                                      progress_bar, remat, reverse,
+                                      unroll)
   if jit:
     dyn_vals, out_vals = transform(operands)
   else:

brainpy/dyn/channels.py

@@ -1,29 +1,29 @@
 from brainpy._src.dyn.channels.base import (
-  IonChannel,
+  IonChannel as IonChannel,
 )
 
 from brainpy._src.dyn.channels.calcium import (
-  CalciumChannel,
-  ICaN_IS2008,
-  ICaT_HM1992,
-  ICaT_HP1992,
-  ICaHT_HM1992,
-  ICaHT_Re1993,
-  ICaL_IS2008,
+  CalciumChannel as CalciumChannel,
+  ICaN_IS2008 as ICaN_IS2008,
+  ICaT_HM1992 as ICaT_HM1992,
+  ICaT_HP1992 as ICaT_HP1992,
+  ICaHT_HM1992 as ICaHT_HM1992,
+  ICaHT_Re1993 as ICaHT_Re1993,
+  ICaL_IS2008 as ICaL_IS2008,
 )
 
 
 from brainpy._src.dyn.channels.potassium import (
-  PotassiumChannel,
-  IKDR_Ba2002v2,
-  IK_TM1991v2,
-  IK_HH1952v2,
-  IKA1_HM1992v2,
-  IKA2_HM1992v2,
-  IKK2A_HM1992v2,
-  IKK2B_HM1992v2,
-  IKNI_Ya1989v2,
-  IK_Leak,
+  PotassiumChannel as PotassiumChannel,
+  IKDR_Ba2002v2 as IKDR_Ba2002v2,
+  IK_TM1991v2 as IK_TM1991v2,
+  IK_HH1952v2 as IK_HH1952v2,
+  IKA1_HM1992v2 as IKA1_HM1992v2,
+  IKA2_HM1992v2 as IKA2_HM1992v2,
+  IKK2A_HM1992v2 as IKK2A_HM1992v2,
+  IKK2B_HM1992v2 as IKK2B_HM1992v2,
+  IKNI_Ya1989v2 as IKNI_Ya1989v2,
+  IK_Leak as IK_Leak,
 )
 from brainpy._src.dyn.channels.potassium_compatible import (
   IKDR_Ba2002,