update docs and apis

Author: chaoming0625

.github/workflows/Linux_CI.yml

@@ -28,7 +28,7 @@ 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 https://github.com/google/jax/archive/refs/tags/jax-v0.3.14.tar.gz
         if [ -f requirements-dev.txt ]; then pip install -r requirements-dev.txt; fi
         python setup.py install
     - name: Lint with flake8

.github/workflows/MacOS_CI.yml

@@ -28,7 +28,7 @@ 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 https://github.com/google/jax/archive/refs/tags/jax-v0.3.14.tar.gz
         if [ -f requirements-dev.txt ]; then pip install -r requirements-dev.txt; fi
         python setup.py install
     - name: Lint with flake8

brainpy/dyn/layers/linear.py

@@ -75,6 +75,12 @@ def __init__(
       self.W = bm.TrainVar(self.W)
       self.b = None if (self.b is None) else bm.TrainVar(self.b)
 
+  def __repr__(self):
+    return (f'{self.__class__.__name__}(name={self.name}, '
+            f'num_in={self.num_in}, '
+            f'num_out={self.num_out}, '
+            f'mode={self.mode})')
+
   def reset_state(self, batch_size=None):
     pass
 
@@ -173,7 +179,7 @@ def offline_fit(self,
       xs = bm.concatenate([bm.ones(xs.shape[:2] + (1,)), xs], axis=-1)  # (..., 1 + num_ff_input)
 
     # solve weights by offline training methods
-    weights = self.offline_fit_by(target, xs, ys)
+    weights = self.offline_fit_by(self.name, target, xs, ys)
 
     # assign trained weights
     if self.b is None:

brainpy/math/jaxarray.py

@@ -76,6 +76,15 @@ def update(self, value):
     """
     if self._outside_global_jit and _global_jit_mode:
       raise MathError(msg)
+    if isinstance(value, JaxArray):
+      value = value.value
+    elif isinstance(value, np.ndarray):
+      value = jnp.asarray(value)
+    elif isinstance(value, jnp.ndarray):
+      pass
+    else:
+      value = jnp.asarray(value)
+    # check
     if value.shape != self._value.shape:
       raise MathError(f"The shape of the original data is {self._value.shape}, "
                       f"while we got {value.shape}.")

docs/index.rst

@@ -10,10 +10,10 @@ high-performance Brain Dynamics Programming (BDP). Among its key ingredients, Br
   delay differential equations (DDEs),
   fractional differential equations (FDEs), etc.
 - **Dynamics building** with the modular and composable programming interface.
-- **Dynamics simulation** tools for various brain objects.
-- **Dynamics training** tools with various machine learning algorithms,
+- **Dynamics simulation** for various brain objects with parallel supports.
+- **Dynamics training** with various machine learning algorithms,
   like FORCE learning, ridge regression, back-propagation, etc.
-- **Dynamics analysis** tools for differential equations, including
+- **Dynamics analysis** for low- and high-dimensional systems, including
   phase plane analysis, bifurcation analysis, linearization analysis,
   and fixed/slow point finding.
 - And more others ......