[documentation] update documentation to brainpy>=2.4.0 (#361)

[documentation] Update documentation to brainpy>=2.4.0

Author: chaoming0625

brainpy/_src/analysis/highdim/slow_points.py

@@ -234,6 +234,8 @@ def __init__(
     if f_loss_batch is not None:
       raise UnsupportedError('"f_loss_batch" is no longer supported, please '
                              'use "f_loss" instead.')
+    if fun_inputs is not None:
+      raise UnsupportedError('"fun_inputs" is no longer supported.')
     if f_loss is None:
       f_loss = losses.mean_squared_error if f_type == constants.DISCRETE else losses.mean_square
     self.f_loss = f_loss

brainpy/_src/analysis/lowdim/lowdim_bifurcation.py

@@ -2,6 +2,7 @@
 
 from functools import partial
 
+import jax
 import jax.numpy as jnp
 from jax import vmap
 import numpy as np

brainpy/_src/analysis/lowdim/lowdim_phase_plane.py

@@ -1,5 +1,6 @@
 # -*- coding: utf-8 -*-
 
+import jax
 import jax.numpy as jnp
 import numpy as np
 from jax import vmap

brainpy/_src/checkpoints/io.py

@@ -8,7 +8,8 @@
 
 from brainpy import errors
 import brainpy.math as bm
-from brainpy._src.math.object_transform.base import BrainPyObject, ArrayCollector
+from brainpy._src.math.object_transform.base import BrainPyObject
+from brainpy._src.math.object_transform.collectors import ArrayCollector
 
 
 logger = logging.getLogger('brainpy.brainpy_object.io')

brainpy/_src/checkpoints/serialization.py

@@ -951,7 +951,7 @@ def _record_saved_duration(checkpoint_start_time: float):
   # Note: for the very first checkpoint, this is the interval between program
   # init and the current checkpoint start time.
   duration_since_last_checkpoint = checkpoint_start_time - _LAST_CHECKPOINT_WRITE_TIME
-  if jax.version.__version_info__ > (0, 3, 25):
+  if monitoring is not None:
     monitoring.record_event_duration_secs(
       '/jax/checkpoint/write/duration_since_last_checkpoint_secs',
       duration_since_last_checkpoint)
@@ -1151,7 +1151,7 @@ def save_main_ckpt_task():
   else:
     save_main_ckpt_task()
   end_time = time.time()
-  if jax.version.__version_info__ > (0, 3, 25):
+  if monitoring is not None:
     monitoring.record_event_duration_secs(_WRITE_CHECKPOINT_EVENT,
                                           end_time - start_time)
   return ckpt_path
@@ -1281,7 +1281,7 @@ def save_main_ckpt_task():
   else:
     save_main_ckpt_task()
   end_time = time.time()
-  if jax.version.__version_info__ > (0, 3, 25):
+  if monitoring is not None:
     monitoring.record_event_duration_secs(_WRITE_CHECKPOINT_EVENT,
                                           end_time - start_time)
 
@@ -1390,7 +1390,7 @@ def save_main_ckpt_task():
                keep, overwrite, keep_every_n_steps, start_time, async_manager)
 
   end_time = time.time()
-  if jax.version.__version_info__ > (0, 3, 25):
+  if monitoring is not None:
     monitoring.record_event_duration_secs(_WRITE_CHECKPOINT_EVENT,
                                           end_time - start_time)
   return ckpt_path
@@ -1553,7 +1553,7 @@ def read_chunk(i):
     restored_checkpoint = from_state_dict(target, state_dict)
 
   end_time = time.time()
-  if jax.version.__version_info__ > (0, 3, 25):
+  if monitoring is not None:
     monitoring.record_event_duration_secs(_READ_CHECKPOINT_EVENT, end_time - start_time)
 
   return restored_checkpoint
@@ -1616,7 +1616,7 @@ def read_chunk(i):
 
   state_dict = msgpack_restore(checkpoint_contents)
   end_time = time.time()
-  if jax.version.__version_info__ > (0, 3, 25):
+  if monitoring is not None:
     monitoring.record_event_duration_secs(_READ_CHECKPOINT_EVENT, end_time - start_time)
 
   return state_dict

brainpy/_src/dyn/base.py

@@ -1475,16 +1475,10 @@ def __getitem__(self, key: Union[int, slice, str]):
     elif isinstance(key, slice):
       return Sequential(*(self.__all_nodes()[key]))
     elif isinstance(key, int):
-      key = self.__format_key(key)
-      return self._static_modules[key] if (key not in self._dyn_modules) else self._dyn_modules[key]
+      return self.__all_nodes()[key]
     elif isinstance(key, (tuple, list)):
-      nodes = []
-      for i in key:
-        if isinstance(i, int):
-          i = self.__format_key(i)
-        assert isinstance(i, str)
-        nodes.append(self._static_modules[i] if (i not in self._dyn_modules) else self._dyn_modules[i])
-      return Sequential(*nodes)
+      _all_nodes = self.__all_nodes()
+      return Sequential(*[_all_nodes[k] for k in key])
     else:
       raise KeyError(f'Unknown type of key: {type(key)}')
 

brainpy/_src/dyn/context.py

@@ -65,6 +65,14 @@ def save(self, *args, **kwargs) -> None:
     for identifier, data in kwargs.items():
       self._arguments[identifier] = data
 
+  def __setitem__(self, key, value):
+    """Enable setting the shared item by ``bp.share[key] = value``."""
+    self.save(key, value)
+
+  def __getitem__(self, item):
+    """Enable loading the shared parameter by ``bp.share[key]``."""
+    return self.load(item)
+
   def get_shargs(self) -> DotDict:
     """Get all shared arguments in the global context."""
     return self._arguments.copy()

brainpy/_src/dyn/neurons/biological_models.py

@@ -456,8 +456,8 @@ def __init__(
     self.input_var = input_var
 
     # initializers
-    self._W_initializer = check.is_initializer(V_initializer, allow_none=False)
-    self._V_initializer = check.is_initializer(W_initializer, allow_none=False)
+    self._W_initializer = check.is_initializer(W_initializer, allow_none=False)
+    self._V_initializer = check.is_initializer(V_initializer, allow_none=False)
 
     # variables
     self.reset_state(self.mode)
@@ -491,7 +491,7 @@ def dW(self, W, t, V):
 
   @property
   def derivative(self):
-    return JointEq([self.dV, self.dW])
+    return JointEq(self.dV, self.dW)
 
   def update(self, x=None):
     t = share.load('t')

brainpy/_src/dyn/neurons/input_groups.py

@@ -131,10 +131,10 @@ def __init__(
     self.num_times = len(times)
 
     # data about times and indices
-    self.times = jnp.asarray(times)
-    self.indices = jnp.asarray(indices, dtype=bm.int_)
+    self.times = bm.asarray(times)
+    self.indices = bm.asarray(indices, dtype=bm.int_)
     if need_sort:
-      sort_idx = jnp.argsort(self.times)
+      sort_idx = bm.argsort(self.times)
       self.indices.value = self.indices[sort_idx]
       self.times.value = self.times[sort_idx]
 
@@ -144,7 +144,7 @@ def __init__(
     # functions
     def cond_fun(t):
       i = self.i.value
-      return jnp.logical_and(i < self.num_times, t >= self.times[i])
+      return bm.logical_and(i < self.num_times, t >= self.times[i])
 
     def body_fun(t):
       i = self.i.value

brainpy/_src/dyn/runners.py

@@ -640,7 +640,6 @@ def _step_func_predict(self, shared_args, t, i, x):
     shared = tools.DotDict(t=t, i=i, dt=self.dt)
     shared.update(shared_args)
     share.save(**shared)
-    self.target.clear_input()
     self._step_func_input(shared)
 
     # dynamics update step
@@ -655,6 +654,7 @@ def _step_func_predict(self, shared_args, t, i, x):
     if self.progress_bar:
       id_tap(lambda *arg: self._pbar.update(), ())
     share.clear_shargs()
+    self.target.clear_input()
 
     if self._memory_efficient:
       id_tap(self._step_mon_on_cpu, mon)