update brainpy package

Author: chaoming0625

brainpy/__init__.py

@@ -58,14 +58,14 @@
   DynamicalSystem as DynamicalSystem,
   DynSysGroup as DynSysGroup,  # collectors
   Sequential as Sequential,
-  Network as Network,
   Dynamic as Dynamic,  # category
   Projection as Projection,
 )
 DynamicalSystemNS = DynamicalSystem
+Network = DynSysGroup
 # delays
 from brainpy._src.delay import (
-  VariableDelay as VariableDelay,
+  VariDelay as VariDelay,
 )
 
 # building blocks

brainpy/_src/delay.py

@@ -21,8 +21,9 @@
 
 __all__ = [
   'Delay',
-  'VariableDelay',
+  'VariDelay',
   'DataDelay',
+  'DelayAccess',
 ]
 
 
@@ -431,8 +432,8 @@ def _check_target_sharding(sharding, ndim, mode: bm.Mode):
   return sharding
 
 
-class VariableDelay(Delay):
-  """Delay variable which has a fixed delay length.
+class VariDelay(Delay):
+  """Generate Delays for the given :py:class:`~.Variable` instance.
 
   The data in this delay variable is arranged as::
 
@@ -517,8 +518,8 @@ def __init__(
 
     # other info
     if entries is not None:
-      for entry, value in entries.items():
-        self.register_entry(entry, value)
+      for entry, delay_time in entries.items():
+        self.register_entry(entry, delay_time)
 
   def register_entry(
       self,
@@ -572,11 +573,17 @@ def at(self, entry: str, *indices) -> bm.Array:
       raise KeyError(f'Does not find delay entry "{entry}".')
     delay_step = self._registered_entries[entry]
     if delay_step is None or delay_step == 0.:
-      return self.target.value
+      if len(indices):
+        return self.target[indices]
+      else:
+        return self.target.value
     else:
       assert self.data is not None
       if delay_step == 0:
-        return self.target.value
+        if len(indices):
+          return self.target[indices]
+        else:
+          return self.target.value
       else:
         return self.retrieve(delay_step, *indices)
 
@@ -683,16 +690,15 @@ def _init_data(self, length: int, batch_size: int = None):
       self.data[:] = self._init((length,) + self.target.shape, dtype=self.target.dtype)
 
 
-class DataDelay(VariableDelay):
-  
+class DataDelay(VariDelay):
   not_desc_params = ('time', 'entries')
 
   def __init__(
       self,
 
       # delay target
-      target: bm.Variable,
-      target_init: Callable,
+      data: bm.Variable,
+      data_init: Union[Callable, bm.Array, jax.Array],
 
       # delay time
       time: Optional[Union[int, float]] = None,
@@ -710,8 +716,8 @@ def __init__(
       name: Optional[str] = None,
       mode: Optional[bm.Mode] = None,
   ):
-    self.target_init = target_init
-    super().__init__(target=target,
+    self.target_init = data_init
+    super().__init__(target=data,
                      time=time,
                      init=init,
                      entries=entries,
@@ -736,3 +742,20 @@ def update(
     super().update(latest_value)
 
 
+class DelayAccess(DynamicalSystem):
+  def __init__(
+      self,
+      delay: Delay,
+      time: Union[None, int, float],
+      *indices
+  ):
+    super().__init__(mode=delay.mode)
+    self.delay = delay
+    assert isinstance(delay, Delay)
+    delay.register_entry(self.name, time)
+    self.indices = indices
+
+  def update(self):
+    return self.delay.at(self.name, *self.indices)
+
+

brainpy/_src/dnn/linear.py

@@ -4,6 +4,7 @@
 from typing import Dict, Optional, Union, Callable
 
 import jax
+import numpy as np
 import jax.numpy as jnp
 
 from brainpy import math as bm
@@ -63,8 +64,8 @@ def __init__(
       num_out: int,
       W_initializer: Union[Initializer, Callable, ArrayType] = XavierNormal(),
       b_initializer: Optional[Union[Initializer, Callable, ArrayType]] = ZeroInit(),
-      mode: bm.Mode = None,
-      name: str = None,
+      mode: Optional[bm.Mode] = None,
+      name: Optional[str] = None,
   ):
     super(Dense, self).__init__(mode=mode, name=name)
 
@@ -642,7 +643,7 @@ def __init__(
       num_out: int,
       prob: float,
       weight: float,
-      seed: int,
+      seed: Optional[int] = None,
       sharding: Optional[Sharding] = None,
       mode: Optional[bm.Mode] = None,
       name: Optional[str] = None,
@@ -654,7 +655,7 @@ def __init__(
     self.prob = prob
     self.sharding = sharding
     self.transpose = transpose
-    self.seed = seed
+    self.seed = np.random.randint(0, 100000) if seed is None else seed
     self.atomic = atomic
     self.num_in = num_in
     self.num_out = num_out
@@ -723,7 +724,7 @@ def __init__(
       prob: float,
       w_low: float,
       w_high: float,
-      seed: int,
+      seed: Optional[int] = None,
       sharding: Optional[Sharding] = None,
       mode: Optional[bm.Mode] = None,
       name: Optional[str] = None,
@@ -735,7 +736,7 @@ def __init__(
     self.prob = prob
     self.sharding = sharding
     self.transpose = transpose
-    self.seed = seed
+    self.seed = np.random.randint(0, 100000) if seed is None else seed
     self.atomic = atomic
     self.num_in = num_in
     self.num_out = num_out
@@ -803,7 +804,7 @@ def __init__(
       prob: float,
       w_mu: float,
       w_sigma: float,
-      seed: int,
+      seed: Optional[int] = None,
       sharding: Optional[Sharding] = None,
       transpose: bool = False,
       atomic: bool = False,
@@ -815,7 +816,7 @@ def __init__(
     self.prob = prob
     self.sharding = sharding
     self.transpose = transpose
-    self.seed = seed
+    self.seed = np.random.randint(0, 100000) if seed is None else seed
     self.atomic = atomic
     self.num_in = num_in
     self.num_out = num_out
@@ -881,7 +882,7 @@ def __init__(
       num_out: int,
       prob: float,
       weight: float,
-      seed: int,
+      seed: Optional[int] = None,
       sharding: Optional[Sharding] = None,
       mode: Optional[bm.Mode] = None,
       name: Optional[str] = None,
@@ -893,7 +894,7 @@ def __init__(
     self.prob = prob
     self.sharding = sharding
     self.transpose = transpose
-    self.seed = seed
+    self.seed = np.random.randint(0, 1000000) if seed is None else seed
     self.atomic = atomic
     self.num_in = num_in
     self.num_out = num_out
@@ -962,7 +963,7 @@ def __init__(
       prob: float,
       w_low: float,
       w_high: float,
-      seed: int,
+      seed: Optional[int] = None,
       sharding: Optional[Sharding] = None,
       mode: Optional[bm.Mode] = None,
       name: Optional[str] = None,
@@ -974,7 +975,7 @@ def __init__(
     self.prob = prob
     self.sharding = sharding
     self.transpose = transpose
-    self.seed = seed
+    self.seed = np.random.randint(0, 100000) if seed is None else seed
     self.atomic = atomic
     self.num_in = num_in
     self.num_out = num_out
@@ -1042,7 +1043,7 @@ def __init__(
       prob: float,
       w_mu: float,
       w_sigma: float,
-      seed: int,
+      seed: Optional[int] = None,
       sharding: Optional[Sharding] = None,
       transpose: bool = False,
       atomic: bool = False,
@@ -1054,7 +1055,7 @@ def __init__(
     self.prob = prob
     self.sharding = sharding
     self.transpose = transpose
-    self.seed = seed
+    self.seed = np.random.randint(0, 100000) if seed is None else seed
     self.atomic = atomic
     self.num_in = num_in
     self.num_out = num_out

brainpy/_src/dyn/others/input.py

@@ -40,11 +40,11 @@ def __init__(
       mode: Optional[bm.Mode] = None,
       name: Optional[str] = None,
   ):
-    super(InputGroup, self).__init__(name=name,
-                                     sharding=sharding,
-                                     size=size,
-                                     keep_size=keep_size,
-                                     mode=mode)
+    super().__init__(name=name,
+                     sharding=sharding,
+                     size=size,
+                     keep_size=keep_size,
+                     mode=mode)
 
   def update(self, x):
     return x
@@ -74,11 +74,11 @@ def __init__(
       mode: Optional[bm.Mode] = None,
       name: Optional[str] = None,
   ):
-    super(OutputGroup, self).__init__(name=name,
-                                      sharding=sharding,
-                                      size=size,
-                                      keep_size=keep_size,
-                                      mode=mode)
+    super().__init__(name=name,
+                     sharding=sharding,
+                     size=size,
+                     keep_size=keep_size,
+                     mode=mode)
 
   def update(self, x):
     return x
@@ -130,11 +130,11 @@ def __init__(
       mode: Optional[bm.Mode] = None,
       need_sort: bool = True,
   ):
-    super(SpikeTimeGroup, self).__init__(size=size,
-                                         sharding=sharding,
-                                         name=name,
-                                         keep_size=keep_size,
-                                         mode=mode)
+    super().__init__(size=size,
+                     sharding=sharding,
+                     name=name,
+                     keep_size=keep_size,
+                     mode=mode)
 
     # parameters
     if keep_size:
@@ -202,11 +202,11 @@ def __init__(
       mode: Optional[bm.Mode] = None,
       seed=None,
   ):
-    super(PoissonGroup, self).__init__(size=size,
-                                       sharding=sharding,
-                                       name=name,
-                                       keep_size=keep_size,
-                                       mode=mode)
+    super().__init__(size=size,
+                     sharding=sharding,
+                     name=name,
+                     keep_size=keep_size,
+                     mode=mode)
 
     if seed is not None:
       warnings.warn('')

brainpy/_src/dyn/synapses/abstract_models.py

@@ -334,7 +334,8 @@ def add_current(self, inp):
     self.g_decay += inp
 
   def return_info(self):
-    return ReturnInfo(self.varshape, self.sharding, self.mode, bm.zeros)
+    return ReturnInfo(self.varshape, self.sharding, self.mode,
+                      lambda shape: self.coeff * (self.g_decay - self.g_rise))
 
 
 DualExponV2.__doc__ = DualExponV2.__doc__ % (pneu_doc,)
@@ -677,22 +678,21 @@ def update(self, pre_spike):
     t = share.load('t')
     dt = share.load('dt')
     u, x = self.integral(self.u.value, self.x.value, t, dt)
-    if pre_spike.dtype == jax.numpy.bool_:
-      u = bm.where(pre_spike, u + self.U * (1 - self.u), u)
-      x = bm.where(pre_spike, x - u * self.x, x)
-    else:
-      u = pre_spike * (u + self.U * (1 - self.u)) + (1 - pre_spike) * u
-      x = pre_spike * (x - u * self.x) + (1 - pre_spike) * x
+    # if pre_spike.dtype == jax.numpy.bool_:
+    #   u = bm.where(pre_spike, u + self.U * (1 - self.u), u)
+    #   x = bm.where(pre_spike, x - u * self.x, x)
+    # else:
+    # u = pre_spike * (u + self.U * (1 - self.u)) + (1 - pre_spike) * u
+    # x = pre_spike * (x - u * self.x) + (1 - pre_spike) * x
+    u = pre_spike * self.U * (1 - self.u) + u
+    x = pre_spike * -u * self.x + x
     self.x.value = x
     self.u.value = u
     return u * x
 
   def return_info(self):
-    return ReturnInfo(size=self.varshape,
-                      batch_or_mode=self.mode,
-                      axis_names=self.sharding,
-                      init=Constant(self.U))
+    return ReturnInfo(self.varshape, self.sharding, self.mode,
+                      lambda shape: self.u * self.x)
 
 
 STP.__doc__ = STP.__doc__ % (pneu_doc,)
-

brainpy/_src/math/object_transform/jit.py

@@ -405,13 +405,14 @@ def _make_jit_fun(
 
   @wraps(fun)
   def call_fun(self, *args, **kwargs):
-    fun2 = partial(fun, self)
     if jax.config.jax_disable_jit:
-      return fun2(*args, **kwargs)
+      return fun(self, *args, **kwargs)
 
     hash_v = hash(fun) + hash(self)
     cache = get_stack_cache(hash_v)  # TODO: better cache mechanism
     if cache is None:
+      fun2 = partial(fun, self)
+      
       with jax.ensure_compile_time_eval():
         if len(static_argnums) or len(static_argnames):
           fun3, args_, kwargs_ = _partial_fun(fun2, args, kwargs, static_argnums, static_argnames)

brainpy/_src/math/object_transform/variables.py

@@ -71,7 +71,7 @@ def dict_data(self) -> dict:
     """Get all data in the collected variables with a python dict structure."""
     new_dict = dict()
     for id_, elem in tuple(self.items()):
-      new_dict[id_] = elem.value if isinstance(elem, Array) else elem
+      new_dict[id_] = elem._value if isinstance(elem, Array) else elem
     return new_dict
 
   def list_data(self) -> list: