Docs update and bugs fixed

.github/workflows/MacOS_CI.yml

@@ -28,7 +28,8 @@ jobs:
       run: |
         python -m pip install --upgrade pip
         python -m pip install flake8 pytest
-        # python -m pip install https://github.com/google/jax/archive/refs/tags/jax-v0.3.14.tar.gz
+        python -m pip install jax==0.3.14
+        python -m pip install jaxlib==0.3.14
         if [ -f requirements-dev.txt ]; then pip install -r requirements-dev.txt; fi
         python setup.py install
     - name: Lint with flake8

brainpy/algorithms/offline.py

@@ -7,7 +7,7 @@
 
 import brainpy.math as bm
 from brainpy.base import Base
-from brainpy.types import Tensor
+from brainpy.types import Array
 from .utils import (Sigmoid,
                     Regularization, L1Regularization, L1L2Regularization, L2Regularization,
                     polynomial_features, normalize)
@@ -60,7 +60,7 @@ def __call__(self, identifier, target, input, output):
     """
     return self.call(identifier, target, input, output)
 
-  def call(self, identifier, targets, inputs, outputs) -> Tensor:
+  def call(self, identifier, targets, inputs, outputs) -> Array:
     """The training procedure.
 
     Parameters
@@ -355,7 +355,7 @@ def __init__(
     self.gradient_descent = gradient_descent
     self.sigmoid = Sigmoid()
 
-  def call(self, identifier, targets, inputs, outputs=None) -> Tensor:
+  def call(self, identifier, targets, inputs, outputs=None) -> Array:
     # prepare data
     inputs = _check_data_2d_atls(bm.asarray(inputs))
     targets = _check_data_2d_atls(bm.asarray(targets))

brainpy/analysis/highdim/slow_points.py

@@ -18,7 +18,7 @@
 from brainpy.dyn.runners import build_inputs, check_and_format_inputs
 from brainpy.errors import AnalyzerError, UnsupportedError
 from brainpy.tools.others.dicts import DotDict
-from brainpy.types import Tensor
+from brainpy.types import Array
 
 __all__ = [
   'SlowPointFinder',
@@ -295,7 +295,7 @@ def selected_ids(self, val):
 
   def find_fps_with_gd_method(
       self,
-      candidates: Union[Tensor, Dict[str, Tensor]],
+      candidates: Union[Array, Dict[str, Array]],
       tolerance: Union[float, Dict[str, float]] = 1e-5,
       num_batch: int = 100,
       num_opt: int = 10000,
@@ -305,7 +305,7 @@ def find_fps_with_gd_method(
 
     Parameters
     ----------
-    candidates : Tensor, dict
+    candidates : Array, dict
       The array with the shape of (batch size, state dim) of hidden states
       of RNN to start training for fixed points.
 
@@ -402,14 +402,14 @@ def batch_train(start_i, n_batch):
 
   def find_fps_with_opt_solver(
       self,
-      candidates: Union[Tensor, Dict[str, Tensor]],
+      candidates: Union[Array, Dict[str, Array]],
       opt_solver: str = 'BFGS'
   ):
     """Optimize fixed points with nonlinear optimization solvers.
 
     Parameters
     ----------
-    candidates: Tensor, dict
+    candidates: Array, dict
       The candidate (initial) fixed points.
     opt_solver: str
       The solver of the optimization.
@@ -536,7 +536,7 @@ def exclude_outliers(self, tolerance: float = 1e0):
 
   def compute_jacobians(
       self,
-      points: Union[Tensor, Dict[str, Tensor]],
+      points: Union[Array, Dict[str, Array]],
       stack_dict_var: bool = True,
       plot: bool = False,
       num_col: int = 4,

brainpy/dyn/base.py

@@ -18,7 +18,7 @@
 from brainpy.integrators import odeint, sdeint
 from brainpy.modes import Mode, TrainingMode, BatchingMode, normal, training
 from brainpy.tools.others import to_size, size2num
-from brainpy.types import Tensor, Shape
+from brainpy.types import Array, Shape
 
 __all__ = [
   # general class
@@ -70,17 +70,17 @@ def __init__(
       name: str = None,
       mode: Optional[Mode] = None,
   ):
-    super(DynamicalSystem, self).__init__(name=name)
-
-    # local delay variables
-    self.local_delay_vars: Dict[str, bm.LengthDelay] = Collector()
-
     # mode setting
     if mode is None: mode = normal
     if not isinstance(mode, Mode):
       raise ValueError(f'Should be instance of {Mode.__name__}, but we got {type(Mode)}: {Mode}')
     self._mode = mode
 
+    super(DynamicalSystem, self).__init__(name=name)
+
+    # local delay variables
+    self.local_delay_vars: Dict[str, bm.LengthDelay] = Collector()
+
     # fitting parameters
     self.online_fit_by = None
     self.offline_fit_by = None
@@ -106,9 +106,9 @@ def __call__(self, *args, **kwargs):
   def register_delay(
       self,
       identifier: str,
-      delay_step: Optional[Union[int, Tensor, Callable, Initializer]],
+      delay_step: Optional[Union[int, Array, Callable, Initializer]],
       delay_target: bm.Variable,
-      initial_delay_data: Union[Initializer, Callable, Tensor, float, int, bool] = None,
+      initial_delay_data: Union[Initializer, Callable, Array, float, int, bool] = None,
   ):
     """Register delay variable.
 
@@ -317,14 +317,14 @@ def offline_init(self):
 
   @tools.not_customized
   def online_fit(self,
-                 target: Tensor,
-                 fit_record: Dict[str, Tensor]):
+                 target: Array,
+                 fit_record: Dict[str, Array]):
     raise NoImplementationError('Subclass must implement online_fit() function when using OnlineTrainer.')
 
   @tools.not_customized
   def offline_fit(self,
-                  target: Tensor,
-                  fit_record: Dict[str, Tensor]):
+                  target: Array,
+                  fit_record: Dict[str, Array]):
     raise NoImplementationError('Subclass must implement offline_fit() function when using OfflineTrainer.')
 
   def clear_input(self):
@@ -482,7 +482,7 @@ def f(x):
     entries = '\n'.join(f'  [{i}] {f(x)}' for i, x in enumerate(self))
     return f'{self.__class__.__name__}(\n{entries}\n)'
 
-  def update(self, sha: dict, x: Any) -> Tensor:
+  def update(self, sha: dict, x: Any) -> Array:
     """Update function of a sequential model.
 
     Parameters
@@ -494,7 +494,7 @@ def update(self, sha: dict, x: Any) -> Tensor:
 
     Returns
     -------
-    y: Tensor
+    y: Array
       The output tensor.
     """
     for node in self.implicit_nodes.values():
@@ -686,7 +686,7 @@ def __init__(
       self,
       pre: NeuGroup,
       post: NeuGroup,
-      conn: Union[TwoEndConnector, Tensor, Dict[str, Tensor]] = None,
+      conn: Union[TwoEndConnector, Array, Dict[str, Array]] = None,
       name: str = None,
       mode: Mode = normal,
   ):
@@ -904,7 +904,7 @@ def __init__(
       self,
       pre: NeuGroup,
       post: NeuGroup,
-      conn: Union[TwoEndConnector, Tensor, Dict[str, Tensor]] = None,
+      conn: Union[TwoEndConnector, Array, Dict[str, Array]] = None,
       output: SynOut = NullSynOut(),
       stp: SynSTP = NullSynSTP(),
       ltp: SynLTP = NullSynLTP(),
@@ -946,10 +946,10 @@ def __init__(
 
   def init_weights(
       self,
-      weight: Union[float, Tensor, Initializer, Callable],
+      weight: Union[float, Array, Initializer, Callable],
       comp_method: str,
       sparse_data: str = 'csr'
-  ) -> Union[float, Tensor]:
+  ) -> Union[float, Array]:
     if comp_method not in ['sparse', 'dense']:
       raise ValueError(f'"comp_method" must be in "sparse" and "dense", but we got {comp_method}')
     if sparse_data not in ['csr', 'ij']:
@@ -1061,11 +1061,11 @@ def __init__(
       self,
       size: Shape,
       keep_size: bool = False,
-      C: Union[float, Tensor, Initializer, Callable] = 1.,
-      A: Union[float, Tensor, Initializer, Callable] = 1e-3,
-      V_th: Union[float, Tensor, Initializer, Callable] = 0.,
-      V_initializer: Union[Initializer, Callable, Tensor] = Uniform(-70, -60.),
-      noise: Union[float, Tensor, Initializer, Callable] = None,
+      C: Union[float, Array, Initializer, Callable] = 1.,
+      A: Union[float, Array, Initializer, Callable] = 1e-3,
+      V_th: Union[float, Array, Initializer, Callable] = 0.,
+      V_initializer: Union[Initializer, Callable, Array] = Uniform(-70, -60.),
+      noise: Union[float, Array, Initializer, Callable] = None,
       method: str = 'exp_auto',
       name: str = None,
       mode: Mode = normal,