fix bugs

Author: chaoming0625

brainpy/_src/dyn/layers/pooling.py

@@ -82,6 +82,7 @@ def __init__(
 
   def update(self, *args):
     x = args[0] if len(args) == 1 else args[1]
+    x = bm.as_jax(x)
     window_shape = self._infer_shape(x.ndim, self.kernel_size)
     stride = self._infer_shape(x.ndim, self.stride)
     padding = (self.padding
@@ -258,6 +259,7 @@ def __init__(
 
   def update(self, *args):
     x = args[0] if len(args) == 1 else args[1]
+    x = bm.as_jax(x)
     window_shape = self._infer_shape(x.ndim, self.kernel_size)
     strides = self._infer_shape(x.ndim, self.stride)
     padding = (self.padding if isinstance(self.padding, str) else
@@ -356,6 +358,7 @@ def __init__(
 
   def update(self, *args):
     x = args[0] if len(args) == 1 else args[1]
+    x = bm.as_jax(x)
     x_dim = self.pool_dim + (0 if self.channel_axis is None else 1)
     if x.ndim < x_dim:
       raise ValueError(f'Excepted input with >= {x_dim} dimensions, but got {x.ndim}.')
@@ -521,6 +524,7 @@ def __init__(
 class _AvgPoolNd(_MaxPoolNd):
   def update(self, *args):
     x = args[0] if len(args) == 1 else args[1]
+    x = bm.as_jax(x)
     x_dim = self.pool_dim + (0 if self.channel_axis is None else 1)
     if x.ndim < x_dim:
       raise ValueError(f'Excepted input with >= {x_dim} dimensions, but got {x.ndim}.')
@@ -694,6 +698,7 @@ def _adaptive_pool1d(x, target_size: int, operation: Callable):
   Returns:
     A JAX array of shape `(target_size, )`.
   """
+  x = bm.as_jax(x)
   size = jnp.size(x)
   num_head_arrays = size % target_size
   num_block = size // target_size
@@ -767,6 +772,7 @@ def update(self, *args):
       or `(..., dim_1, dim_2)`.
     """
     x = args[0] if len(args) == 1 else args[1]
+    x = bm.as_jax(x)
 
     # channel axis
     channel_axis = self.channel_axis

brainpy/_src/dyn/neurons/reduced_models.py

@@ -580,7 +580,7 @@ def __init__(
       b: Union[float, ArrayType, Initializer, Callable] = 1.,
       tau: Union[float, ArrayType, Initializer, Callable] = 10.,
       tau_w: Union[float, ArrayType, Initializer, Callable] = 30.,
-      tau_ref: Optional[Union[float, ArrayType, Initializer, Callable]] = 30.,
+      tau_ref: Optional[Union[float, ArrayType, Initializer, Callable]] = None,
       R: Union[float, ArrayType, Initializer, Callable] = 1.,
       V_initializer: Union[Initializer, Callable, ArrayType] = ZeroInit(),
       w_initializer: Union[Initializer, Callable, ArrayType] = ZeroInit(),

brainpy/_src/initialize/random_inits.py

@@ -329,7 +329,7 @@ def __call__(self, shape, dtype=None):
     n_cols = np.prod(shape) // n_rows
     matrix_shape = (n_rows, n_cols) if n_rows > n_cols else (n_cols, n_rows)
     norm_dst = self.rng.normal(size=matrix_shape)
-    q_mat, r_mat = jnp.linalg.qr(norm_dst)
+    q_mat, r_mat = jnp.linalg.qr(bm.as_jax(norm_dst))
     # Enforce Q is uniformly distributed
     q_mat *= jnp.sign(jnp.diag(r_mat))
     if n_rows < n_cols:

brainpy/_src/losses/comparison.py

@@ -11,9 +11,9 @@
 from typing import Tuple
 
 import jax.numpy as jnp
+from jax.lax import scan
 from jax.scipy.special import logsumexp
 from jax.tree_util import tree_map
-from jax.lax import scan
 
 import brainpy.math as bm
 from brainpy.types import ArrayType
@@ -106,7 +106,7 @@ def _cel(_pred, _tar):
     loss = logsumexp(bm.as_jax(_pred), axis=-1) - (_pred * _tar).sum(axis=-1)
     return _reduce(outputs=loss, reduction=reduction)
 
-  r = tree_map(_cel, predicts, targets, is_leaf=lambda x: isinstance(x, bm.Array))
+  r = tree_map(_cel, predicts, targets, is_leaf=_is_leaf)
   return _multi_return(r)
 
 
@@ -128,7 +128,7 @@ def crs(_prd, _tar):
       logits = jnp.take_along_axis(_prd, _tar, -1).squeeze(-1)
     return logsumexp(bm.as_jax(_prd), axis=-1) - logits
 
-  r = tree_map(crs, predicts, targets, is_leaf=lambda x: isinstance(x, bm.Array))
+  r = tree_map(crs, predicts, targets, is_leaf=_is_leaf)
   return _multi_return(r)
 
 
@@ -142,9 +142,14 @@ def cross_entropy_sigmoid(predicts, targets):
   Returns:
       (batch, ...) tensor of the cross-entropies for each entry.
   """
-  r = tree_map(lambda pred, tar: jnp.maximum(pred, 0) - pred * tar + jnp.log(1 + jnp.exp(-jnp.abs(pred))),
-               predicts,
-               targets)
+  r = tree_map(
+    lambda pred, tar: bm.as_jax(
+      bm.maximum(pred, 0) - pred * tar + bm.log(1 + bm.exp(-bm.abs(pred)))
+    ),
+    predicts,
+    targets,
+    is_leaf=_is_leaf
+  )
   return _multi_return(r)
 
 
@@ -201,7 +206,7 @@ def loss(pred, tar):
     norm = jnp.linalg.norm(bm.as_jax(diff), ord=1, axis=1, keepdims=False)
     return _reduce(outputs=norm, reduction=reduction)
 
-  r = tree_map(loss, logits, targets, is_leaf=lambda x: isinstance(x, bm.Array))
+  r = tree_map(loss, logits, targets, is_leaf=_is_leaf)
   return _multi_return(r)
 
 
@@ -228,7 +233,9 @@ def l2_loss(predicts, targets):
   ----------
   .. [1] Bishop, Christopher M. 2006. Pattern Recognition and Machine Learning.
   """
-  r = tree_map(lambda pred, tar: 0.5 * (pred - tar) ** 2, predicts, targets)
+  r = tree_map(lambda pred, tar: 0.5 * (pred - tar) ** 2,
+               predicts,
+               targets)
   return _multi_return(r)
 
 
@@ -243,7 +250,10 @@ def mean_absolute_error(x, y, axis=None, reduction: str = 'mean'):
   Returns:
       tensor of shape (d_i, ..., for i in keep_axis) containing the mean absolute error.
   """
-  r = tree_map(lambda a, b: _reduce(jnp.abs(a - b), reduction=reduction, axis=axis), x, y)
+  r = tree_map(lambda a, b: _reduce(bm.abs(a - b), reduction=reduction, axis=axis),
+               x,
+               y,
+               is_leaf=_is_leaf)
   return _multi_return(r)
 
 
@@ -260,7 +270,8 @@ def mean_squared_error(predicts, targets, axis=None, reduction: str = 'mean'):
   """
   r = tree_map(lambda a, b: _reduce((a - b) ** 2, reduction, axis=axis),
                predicts,
-               targets)
+               targets,
+               is_leaf=_is_leaf)
   return _multi_return(r)
 
 
@@ -276,7 +287,9 @@ def mean_squared_log_error(predicts, targets, axis=None, reduction: str = 'mean'
       tensor of shape (d_i, ..., for i in keep_axis) containing the mean squared error.
   """
   r = tree_map(lambda a, b: _reduce((jnp.log1p(a) - jnp.log1p(b)) ** 2, reduction, axis=axis),
-               predicts, targets, is_leaf=_is_leaf)
+               predicts,
+               targets,
+               is_leaf=_is_leaf)
   return _multi_return(r)
 
 
@@ -309,12 +322,13 @@ def huber_loss(predicts, targets, delta: float = 1.0):
   def _loss(y_predict, y_target):
     # 0.5 * err^2                  if |err| <= d
     # 0.5 * d^2 + d * (|err| - d)  if |err| > d
-    diff = jnp.abs(y_predict - y_target)
-    return jnp.where(diff > delta,
-                    delta * (diff - .5 * delta),
-                    0.5 * diff ** 2)
+    diff = bm.abs(y_predict - y_target)
+    r = bm.where(diff > delta,
+                 delta * (diff - .5 * delta),
+                 0.5 * diff ** 2)
+    return bm.as_jax(r)
 
-  r = tree_map(_loss, targets, predicts)
+  r = tree_map(_loss, targets, predicts, is_leaf=_is_leaf)
   return _multi_return(r)
 
 
@@ -382,7 +396,7 @@ def loss(pred, tar):
     log_not_p = bm.log_sigmoid(-pred)
     return -tar * log_p - (1. - tar) * log_not_p
 
-  r = tree_map(loss, logits, labels, is_leaf=lambda x: isinstance(x, bm.Array))
+  r = tree_map(loss, logits, labels, is_leaf=_is_leaf)
   return _multi_return(r)
 
 
@@ -433,7 +447,7 @@ def loss(pred, tar):
     errors = bm.as_jax(pred - tar)
     return jnp.logaddexp(errors, -errors) - jnp.log(2.0).astype(errors.dtype)
 
-  r = tree_map(loss, predicts, targets, is_leaf=lambda x: isinstance(x, bm.Array))
+  r = tree_map(loss, predicts, targets, is_leaf=_is_leaf)
   return _multi_return(r)
 
 

brainpy/_src/math/__init__.py

@@ -41,9 +41,9 @@
 # high-level numpy operations
 from .arraycreation import *
 from .arrayinterporate import *
-# from .arraycompatible import *
+from .arraycompatible import *
 from .others import *
-from . import random
+from . import random, linalg, fft
 
 # operators
 from .operators import *

brainpy/_src/math/_utils.py

@@ -4,7 +4,6 @@
 from typing import Callable
 
 import jax
-import numpy as np
 from jax.tree_util import tree_map
 
 from .ndarray import Array
@@ -18,8 +17,8 @@ def _as_jax_array_(obj):
   return obj.value if isinstance(obj, Array) else obj
 
 
-def _return(x):
-  return Array(x) if _return_bp_array else x
+def _return(a):
+  return Array(a) if isinstance(a, jax.Array) and a.ndim > 1 else a
 
 
 _return_bp_array = True
@@ -50,10 +49,6 @@ def wrap(op):
   return wrap
 
 
-def _as_brainpy_array(a):
-  return Array(a) if isinstance(a, (np.ndarray, jax.Array)) else a
-
-
 def _is_leaf(a):
   return isinstance(a, Array)
 
@@ -65,7 +60,7 @@ def new_fun(*args, **kwargs):
     if len(kwargs):
       kwargs = tree_map(_as_jax_array_, kwargs, is_leaf=_is_leaf)
     r = fun(*args, **kwargs)
-    return tree_map(_as_brainpy_array, r) if _return_bp_array else r
+    return tree_map(_return, r) if _return_bp_array else r
 
   new_fun.__doc__ = getattr(fun, "__doc__", None)
 

brainpy/_src/math/object_transform/autograd.py

@@ -18,9 +18,9 @@
 from jax.util import safe_map
 
 from brainpy import errors, tools, check
-from brainpy._src.math.object_transform.base import BrainPyObject
-from brainpy._src.math.object_transform.abstract import ObjectTransform
 from brainpy._src.math.ndarray import Array, Variable, add_context, del_context
+from brainpy._src.math.object_transform.abstract import ObjectTransform
+from brainpy._src.math.object_transform.base import BrainPyObject
 
 __all__ = [
   'grad',  # gradient of scalar function
@@ -75,7 +75,7 @@ def __init__(
       _argnums = tuple(a + 2 for a in _argnums)
       if len(self._grad_vars) > 0:
         _argnums = (0,) + _argnums
-    self.nonvar_argnums = argnums
+    self._nonvar_argnums = argnums
     self.return_value = return_value
     self.has_aux = has_aux
 
@@ -134,10 +134,12 @@ def __call__(self, *args, **kwargs):
     # old_dyn_vs = [v.value for v in self._dyn_vars]
     try:
       add_context(self.name)
-      grads, (outputs, new_grad_vs, new_dyn_vs) = self._call([v.value for v in self._grad_vars],
-                                                             [v.value for v in self._dyn_vars],
-                                                             *args,
-                                                             **kwargs)
+      grads, (outputs, new_grad_vs, new_dyn_vs) = self._call(
+        [v.value for v in self._grad_vars],
+        [v.value for v in self._dyn_vars],
+        *args,
+        **kwargs
+      )
       del_context(self._name)
     except UnexpectedTracerError as e:
       del_context(self._name)
@@ -155,11 +157,11 @@ def __call__(self, *args, **kwargs):
 
     # check returned grads
     if len(self._grad_vars) > 0:
-      if self.nonvar_argnums is None:
+      if self._nonvar_argnums is None:
         grads = self._grad_tree.unflatten(grads)
       else:
         var_grads = self._grad_tree.unflatten(grads[0])
-        arg_grads = grads[1] if isinstance(self.nonvar_argnums, int) else grads[1:]
+        arg_grads = grads[1] if isinstance(self._nonvar_argnums, int) else grads[1:]
         grads = (var_grads, arg_grads)
 
     # check returned value

brainpy/_src/train/offline.py

@@ -84,10 +84,6 @@ def __init__(
     for node in self.train_nodes:
       node.offline_fit_by = fit_method
 
-    # initialize the fitting method
-    for node in self.train_nodes:
-      node.offline_init()
-
   def __repr__(self):
     name = self.__class__.__name__
     prefix = ' ' * len(name)