updates

Author: chaoming0625

brainpy/__init__.py

@@ -129,7 +129,7 @@ def build_csr(self):
   def build_mat(self):
     pre_state = self._jaxrand.uniform(size=(self.pre_num, 1)) < self.pre_ratio
     mat = (self._jaxrand.uniform(size=(self.pre_num, self.post_num)) < self.prob) * pre_state
-    mat = jnp.asarray(mat)
+    mat = bm.asarray(mat)
     if not self.include_self:
       mat = bm.fill_diagonal(mat, False)
     return mat.astype(MAT_DTYPE)

brainpy/_src/connect/random_conn.py

@@ -81,7 +81,11 @@ def __repr__(self):
             f'num_out={self.num_out}, '
             f'mode={self.mode})')
 
-  def update(self, sha, x):
+  def update(self, *args):
+    if len(args) == 1:
+      sha, x = dict(), bm.as_jax(args[0])
+    else:
+      sha, x = args[0], bm.as_jax(args[1])
     res = x @ self.W
     if self.b is not None:
       res += self.b
@@ -102,7 +106,7 @@ def online_init(self):
       num_input = self.num_in
     else:
       num_input = self.num_in + 1
-    self.online_fit_by.initialize(feature_in=num_input, feature_out=self.num_out, identifier=self.name)
+    self.online_fit_by.register_target(feature_in=num_input, identifier=self.name)
 
   def online_fit(self,
                  target: ArrayType,
@@ -139,13 +143,6 @@ def online_fit(self,
       self.b += db[0]
       self.W += dW
 
-  def offline_init(self):
-    if self.b is None:
-      num_input = self.num_in + 1
-    else:
-      num_input = self.num_in
-    self.offline_fit_by.initialize(feature_in=num_input, feature_out=self.num_out, identifier=self.name)
-
   def offline_fit(self,
                   target: ArrayType,
                   fit_record: Dict[str, ArrayType]):
@@ -176,7 +173,7 @@ def offline_fit(self,
       xs = jnp.concatenate([jnp.ones(xs.shape[:2] + (1,)), xs], axis=-1)  # (..., 1 + num_ff_input)
 
     # solve weights by offline training methods
-    weights = self.offline_fit_by(self.name, target, xs, ys)
+    weights = self.offline_fit_by(target, xs, ys)
 
     # assign trained weights
     if self.b is None:

brainpy/_src/dyn/layers/linear.py

@@ -6,7 +6,7 @@
 
 import brainpy.math as bm
 from brainpy._src.initialize import Normal, ZeroInit, Initializer, parameter, variable
-from brainpy.check import is_float, is_initializer, is_string
+from brainpy import check
 from brainpy.tools import to_size
 from brainpy.types import ArrayType
 from .base import Layer
@@ -36,8 +36,9 @@ class Reservoir(Layer):
     A float between 0 and 1.
   activation : str, callable, optional
     Reservoir activation function.
+
     - If a str, should be a :py:mod:`brainpy.math.activations` function name.
-    - If a callable, should be an element-wise operator on tensor.
+    - If a callable, should be an element-wise operator.
   activation_type : str
     - If "internal" (default), then leaky integration happens on states transformed
       by the activation function:
@@ -66,9 +67,12 @@ class Reservoir(Layer):
     neurons connected to other reservoir neurons, including themselves.
     Must be in [0, 1], by default 0.1
   comp_type: str
-    The connectivity type, can be "dense" or "sparse".
+    The connectivity type, can be "dense" or "sparse", "jit".
+
+    - ``"dense"`` means the connectivity matrix is a dense matrix.
+    - ``"sparse"`` means the connectivity matrix is a CSR sparse matrix.
   spectral_radius : float, optional
-    Spectral radius of recurrent weight matrix, by default None
+    Spectral radius of recurrent weight matrix, by default None.
   noise_rec : float, optional
     Gain of noise applied to reservoir internal states, by default 0.0
   noise_in : float, optional
@@ -118,37 +122,38 @@ def __init__(
     self.num_unit = num_out
     assert num_out > 0, f'Must be a positive integer, but we got {num_out}'
     self.leaky_rate = leaky_rate
-    is_float(leaky_rate, 'leaky_rate', 0., 1.)
-    self.activation = getattr(bm.activations, activation)
+    check.is_float(leaky_rate, 'leaky_rate', 0., 1.)
+    self.activation = getattr(bm.activations, activation) if isinstance(activation, str) else activation
+    check.is_callable(self.activation, allow_none=False)
     self.activation_type = activation_type
-    is_string(activation_type, 'activation_type', ['internal', 'external'])
+    check.is_string(activation_type, 'activation_type', ['internal', 'external'])
     self.rng = bm.random.default_rng(seed)
-    is_float(spectral_radius, 'spectral_radius', allow_none=True)
+    check.is_float(spectral_radius, 'spectral_radius', allow_none=True)
     self.spectral_radius = spectral_radius
 
     # initializations
-    is_initializer(Win_initializer, 'ff_initializer', allow_none=False)
-    is_initializer(Wrec_initializer, 'rec_initializer', allow_none=False)
-    is_initializer(b_initializer, 'bias_initializer', allow_none=True)
+    check.is_initializer(Win_initializer, 'ff_initializer', allow_none=False)
+    check.is_initializer(Wrec_initializer, 'rec_initializer', allow_none=False)
+    check.is_initializer(b_initializer, 'bias_initializer', allow_none=True)
     self._Win_initializer = Win_initializer
     self._Wrec_initializer = Wrec_initializer
     self._b_initializer = b_initializer
 
     # connectivity
-    is_float(in_connectivity, 'ff_connectivity', 0., 1.)
-    is_float(rec_connectivity, 'rec_connectivity', 0., 1.)
+    check.is_float(in_connectivity, 'ff_connectivity', 0., 1.)
+    check.is_float(rec_connectivity, 'rec_connectivity', 0., 1.)
     self.ff_connectivity = in_connectivity
     self.rec_connectivity = rec_connectivity
-    is_string(comp_type, 'conn_type', ['dense', 'sparse'])
+    check.is_string(comp_type, 'conn_type', ['dense', 'sparse', 'jit'])
     self.comp_type = comp_type
 
     # noises
-    is_float(noise_in, 'noise_ff')
-    is_float(noise_rec, 'noise_rec')
+    check.is_float(noise_in, 'noise_ff')
+    check.is_float(noise_rec, 'noise_rec')
     self.noise_ff = noise_in
     self.noise_rec = noise_rec
     self.noise_type = noise_type
-    is_string(noise_type, 'noise_type', ['normal', 'uniform'])
+    check.is_string(noise_type, 'noise_type', ['normal', 'uniform'])
 
     # initialize feedforward weights
     weight_shape = (input_shape[-1], self.num_unit)
@@ -170,7 +175,7 @@ def __init__(
       conn_mat = self.rng.random(recurrent_shape) > self.rec_connectivity
       self.Wrec[conn_mat] = 0.
     if self.spectral_radius is not None:
-      current_sr = max(abs(jnp.linalg.eig(self.Wrec)[0]))
+      current_sr = max(abs(jnp.linalg.eig(bm.as_jax(self.Wrec))[0]))
       self.Wrec *= self.spectral_radius / current_sr
     if self.comp_type == 'sparse' and self.rec_connectivity < 1.:
       self.rec_pres, self.rec_posts = jnp.where(jnp.logical_not(bm.as_jax(conn_mat)))
@@ -186,11 +191,13 @@ def __init__(
   def reset_state(self, batch_size=None):
     self.state.value = variable(jnp.zeros, batch_size, self.output_shape)
 
-  def update(self, sha, x):
+  def update(self, *args):
     """Feedforward output."""
     # inputs
-    x = jnp.concatenate(x, axis=-1)
-    if self.noise_ff > 0: x += self.noise_ff * self.rng.uniform(-1, 1, x.shape)
+    x = args[0] if len(args) == 1 else args[1]
+    x = bm.as_jax(x)
+    if self.noise_ff > 0:
+      x += self.noise_ff * self.rng.uniform(-1, 1, x.shape)
     if self.comp_type == 'sparse' and self.ff_connectivity < 1.:
       sparse = {'data': self.Win,
                 'index': (self.ff_pres, self.ff_posts),

brainpy/_src/dyn/layers/reservoir.py

@@ -114,7 +114,7 @@ def __init__(self, mean=0., scale=1., seed=None):
   def __call__(self, *shape, dtype=None):
     shape = _format_shape(shape)
     weights = self.rng.normal(size=shape, loc=self.mean, scale=self.scale)
-    return jnp.asarray(weights, dtype=dtype)
+    return bm.as_jax(weights, dtype=dtype)
 
   def __repr__(self):
     return f'{self.__class__.__name__}(scale={self.scale}, rng={self.rng})'
@@ -140,7 +140,7 @@ def __init__(self, min_val: float = 0., max_val: float = 1., seed=None):
   def __call__(self, shape, dtype=None):
     shape = _format_shape(shape)
     r = self.rng.uniform(low=self.min_val, high=self.max_val, size=shape)
-    return jnp.asarray(r, dtype=dtype)
+    return bm.as_jax(r, dtype=dtype)
 
   def __repr__(self):
     return (f'{self.__class__.__name__}(min_val={self.min_val}, '
@@ -180,14 +180,14 @@ def __call__(self, shape, dtype=None):
     variance = (self.scale / denominator).astype(dtype)
     if self.distribution == "truncated_normal":
       stddev = (jnp.sqrt(variance) / .87962566103423978).astype(dtype)
-      return self.rng.truncated_normal(-2, 2, shape, dtype) * stddev
+      res = self.rng.truncated_normal(-2, 2, shape, dtype) * stddev
     elif self.distribution == "normal":
       res = self.rng.randn(*shape) * jnp.sqrt(variance).astype(dtype)
     elif self.distribution == "uniform":
       res = self.rng.uniform(low=-1, high=1, size=shape) * jnp.sqrt(3 * variance).astype(dtype)
     else:
       raise ValueError("invalid distribution for variance scaling initializer")
-    return jnp.asarray(res, dtype=dtype)
+    return bm.as_jax(res, dtype=dtype)
 
   def __repr__(self):
     name = self.__class__.__name__
@@ -336,7 +336,7 @@ def __call__(self, shape, dtype=None):
       q_mat = q_mat.T
     q_mat = jnp.reshape(q_mat, (n_rows,) + tuple(np.delete(shape, self.axis)))
     q_mat = jnp.moveaxis(q_mat, 0, self.axis)
-    return self.scale * jnp.asarray(q_mat, dtype=dtype)
+    return self.scale * bm.as_jax(q_mat, dtype=dtype)
 
   def __repr__(self):
     return f'{self.__class__.__name__}(scale={self.scale}, axis={self.axis}, rng={self.rng})'

brainpy/_src/initialize/random_inits.py

@@ -30,70 +30,70 @@ def test2(self, last_epoch):
       self.assertTrue(lr1 == lr2)
 
 
-class TestStepLR(parameterized.TestCase):
-
-  @parameterized.named_parameters(
-    {'testcase_name': f'last_epoch={last_epoch}',
-     'last_epoch': last_epoch}
-    for last_epoch in [-1, 0, 5, 10]
-  )
-  def test1(self, last_epoch):
-    scheduler1 = scheduler.StepLR(0.1, 10, gamma=0.1, last_epoch=last_epoch)
-    scheduler2 = scheduler.StepLR(0.1, 10, gamma=0.1, last_epoch=last_epoch)
-
-    for i in range(1, 25):
-      lr1 = scheduler1(i + last_epoch)
-      lr2 = scheduler2()
-      scheduler2.step_epoch()
-      print(f'{scheduler2.last_epoch}, {lr1:.4f}, {lr2:.4f}')
-      self.assertTrue(lr1 == lr2)
-
-
-class TestCosineAnnealingLR(unittest.TestCase):
-  def test1(self):
-    max_epoch = 50
-    iters = 200
-    sch = scheduler.CosineAnnealingLR(0.1, T_max=5, eta_min=0, last_epoch=-1)
-    all_lr1 = [[], []]
-    all_lr2 = [[], []]
-    for epoch in range(max_epoch):
-      for batch in range(iters):
-        all_lr1[0].append(epoch + batch / iters)
-        all_lr1[1].append(sch())
-        sch.step_epoch()
-      all_lr2[0].append(epoch)
-      all_lr2[1].append(sch())
-      sch.step_epoch()
-    plt.subplot(211)
-    plt.plot(jax.numpy.asarray(all_lr1[0]), jax.numpy.asarray(all_lr1[1]))
-    plt.subplot(212)
-    plt.plot(jax.numpy.asarray(all_lr2[0]), jax.numpy.asarray(all_lr2[1]))
-    plt.show()
-    plt.close()
-
-
-class TestCosineAnnealingWarmRestarts(unittest.TestCase):
-  def test1(self):
-    max_epoch = 50
-    iters = 200
-    sch = scheduler.CosineAnnealingWarmRestarts(0.1,
-                                                iters,
-                                                T_0=5,
-                                                T_mult=1,
-                                                last_call=-1)
-    all_lr1 = []
-    all_lr2 = []
-    for epoch in range(max_epoch):
-      for batch in range(iters):
-        all_lr1.append(sch())
-        sch.step_call()
-      all_lr2.append(sch())
-      sch.step_epoch()
-    plt.subplot(211)
-    plt.plot(jax.numpy.asarray(all_lr1))
-    plt.subplot(212)
-    plt.plot(jax.numpy.asarray(all_lr2))
-    plt.show()
-    plt.close()
+# class TestStepLR(parameterized.TestCase):
+#
+#   @parameterized.named_parameters(
+#     {'testcase_name': f'last_epoch={last_epoch}',
+#      'last_epoch': last_epoch}
+#     for last_epoch in [-1, 0, 5, 10]
+#   )
+#   def test1(self, last_epoch):
+#     scheduler1 = scheduler.StepLR(0.1, 10, gamma=0.1, last_epoch=last_epoch)
+#     scheduler2 = scheduler.StepLR(0.1, 10, gamma=0.1, last_epoch=last_epoch)
+#
+#     for i in range(1, 25):
+#       lr1 = scheduler1(i + last_epoch)
+#       lr2 = scheduler2()
+#       scheduler2.step_epoch()
+#       print(f'{scheduler2.last_epoch}, {lr1:.4f}, {lr2:.4f}')
+#       self.assertTrue(lr1 == lr2)
+#
+#
+# class TestCosineAnnealingLR(unittest.TestCase):
+#   def test1(self):
+#     max_epoch = 50
+#     iters = 200
+#     sch = scheduler.CosineAnnealingLR(0.1, T_max=5, eta_min=0, last_epoch=-1)
+#     all_lr1 = [[], []]
+#     all_lr2 = [[], []]
+#     for epoch in range(max_epoch):
+#       for batch in range(iters):
+#         all_lr1[0].append(epoch + batch / iters)
+#         all_lr1[1].append(sch())
+#         sch.step_epoch()
+#       all_lr2[0].append(epoch)
+#       all_lr2[1].append(sch())
+#       sch.step_epoch()
+#     plt.subplot(211)
+#     plt.plot(jax.numpy.asarray(all_lr1[0]), jax.numpy.asarray(all_lr1[1]))
+#     plt.subplot(212)
+#     plt.plot(jax.numpy.asarray(all_lr2[0]), jax.numpy.asarray(all_lr2[1]))
+#     plt.show()
+#     plt.close()
+#
+#
+# class TestCosineAnnealingWarmRestarts(unittest.TestCase):
+#   def test1(self):
+#     max_epoch = 50
+#     iters = 200
+#     sch = scheduler.CosineAnnealingWarmRestarts(0.1,
+#                                                 iters,
+#                                                 T_0=5,
+#                                                 T_mult=1,
+#                                                 last_call=-1)
+#     all_lr1 = []
+#     all_lr2 = []
+#     for epoch in range(max_epoch):
+#       for batch in range(iters):
+#         all_lr1.append(sch())
+#         sch.step_call()
+#       all_lr2.append(sch())
+#       sch.step_epoch()
+#     plt.subplot(211)
+#     plt.plot(jax.numpy.asarray(all_lr1))
+#     plt.subplot(212)
+#     plt.plot(jax.numpy.asarray(all_lr2))
+#     plt.show()
+#     plt.close()