change data format specification from time_major to data_first_axis

Author: chaoming0625

brainpy/dyn/runners.py

@@ -31,6 +31,7 @@
 def _is_brainpy_array(x):
   return isinstance(x, bm.Array)
 
+
 def check_and_format_inputs(host, inputs):
   """Check inputs and get the formatted inputs for the given population.
 
@@ -292,10 +293,10 @@ class DSRunner(Runner):
   numpy_mon_after_run : bool
     When finishing the network running, transform the JAX arrays into numpy ndarray or not?
 
-  time_major: bool
+  data_first_axis: str
     Set the default data dimension arrangement.
-    To indicate whether the first axis is the batch size (``time_major=False``) or the
-    time length (``time_major=True``).
+    To indicate whether the first axis is the batch size (``data_first_axis='B'``) or the
+    time length (``data_first_axis='T'``).
     In order to be compatible with previous API, default is set to be ``False``.
 
     .. versionadded:: 2.3.1
@@ -311,22 +312,22 @@ def __init__(
       inputs: Union[Sequence, Callable] = (),
 
       # monitors
-      monitors: Union[Sequence, Dict] = None,
+      monitors: Optional[Union[Sequence, Dict]] = None,
       numpy_mon_after_run: bool = True,
 
       # jit
       jit: Union[bool, Dict[str, bool]] = True,
-      dyn_vars: Union[bm.Variable, Sequence[bm.Variable], Dict[str, bm.Variable]] = None,
+      dyn_vars: Optional[Union[bm.Variable, Sequence[bm.Variable], Dict[str, bm.Variable]]] = None,
 
       # extra info
-      dt: float = None,
+      dt: Optional[float] = None,
       t0: Union[float, int] = 0.,
       progress_bar: bool = True,
-      time_major: bool = False,
+      data_first_axis: Optional[str] = None,
 
       # deprecated
-      fun_inputs: Callable = None,
-      fun_monitors: Dict[str, Callable] = None,
+      fun_inputs: Optional[Callable] = None,
+      fun_monitors: Optional[Dict[str, Callable]] = None,
   ):
     if not isinstance(target, DynamicalSystem):
       raise RunningError(f'"target" must be an instance of {DynamicalSystem.__name__}, '
@@ -344,7 +345,10 @@ def __init__(
     self._t0 = t0
     self.i0 = bm.Variable(bm.asarray([1], dtype=bm.int_))
     self.t0 = bm.Variable(bm.asarray([t0], dtype=bm.float_))
-    self.time_major = time_major
+    if data_first_axis is None:
+      data_first_axis = 'B' if isinstance(self.target, bm.BatchingMode) else 'T'
+    assert data_first_axis in ['B', 'T']
+    self.data_first_axis = data_first_axis
 
     # parameters
     dt = bm.get_dt() if dt is None else dt
@@ -372,7 +376,7 @@ def __repr__(self):
     return (f'{name}(target={tools.repr_context(str(self.target), indent2)}, \n'
             f'{indent}jit={self.jit},\n'
             f'{indent}dt={self.dt},\n'
-            f'{indent}time_major={self.time_major})')
+            f'{indent}data_first_axis={self.data_first_axis})')
 
   def reset_state(self):
     """Reset state of the ``DSRunner``."""
@@ -407,8 +411,8 @@ def predict(
 
       - If the mode of ``target`` is instance of :py:class:`~.BatchingMode`,
         ``inputs`` must be a PyTree of data with two dimensions:
-        ``(batch, time, ...)`` when ``time_major=False``,
-        or ``(time, batch, ...)`` when ``time_major=True``.
+        ``(batch, time, ...)`` when ``data_first_axis='B'``,
+        or ``(time, batch, ...)`` when ``data_first_axis='T'``.
       - If the mode of ``target`` is instance of :py:class:`~.NonBatchingMode`,
         the ``inputs`` should be a PyTree of data with one dimension:
         ``(time, ...)``.
@@ -462,7 +466,7 @@ def predict(
     shared['i'] += self.i0
     shared['t'] += self.t0
 
-    if isinstance(self.target.mode, bm.BatchingMode) and not self.time_major:
+    if isinstance(self.target.mode, bm.BatchingMode) and self.data_first_axis == 'B':
       inputs = tree_map(lambda x: bm.moveaxis(x, 0, 1),
                         inputs,
                         is_leaf=lambda x: isinstance(x, bm.Array))
@@ -530,7 +534,7 @@ def _predict(
     """
     _predict_func = self._get_f_predict(shared_args)
     outs_and_mons = _predict_func(xs)
-    if isinstance(self.target.mode, bm.BatchingMode) and not self.time_major:
+    if isinstance(self.target.mode, bm.BatchingMode) and self.data_first_axis == 'B':
       outs_and_mons = tree_map(lambda x: bm.moveaxis(x, 0, 1),
                                outs_and_mons,
                                is_leaf=lambda x: isinstance(x, bm.Array))
@@ -573,9 +577,9 @@ def _get_input_batch_size(self, xs=None) -> Optional[int]:
     if isinstance(self.target.mode, bm.NonBatchingMode):
       return None
     if isinstance(xs, (bm.Array, jax.Array, np.ndarray)):
-      return xs.shape[1] if self.time_major else xs.shape[0]
+      return xs.shape[1] if self.data_first_axis == 'T' else xs.shape[0]
     leaves, _ = tree_flatten(xs, is_leaf=_is_brainpy_array)
-    if self.time_major:
+    if self.data_first_axis == 'T':
       num_batch_sizes = [x.shape[1] for x in leaves]
     else:
       num_batch_sizes = [x.shape[0] for x in leaves]
@@ -590,19 +594,13 @@ def _get_input_time_step(self, duration=None, xs=None) -> int:
       return int(duration / self.dt)
     if xs is not None:
       if isinstance(xs, (bm.Array, jnp.ndarray)):
-        if isinstance(self.target.mode, bm.BatchingMode):
-          return xs.shape[0] if self.time_major else xs.shape[1]
-        else:
-          return xs.shape[0]
+        return xs.shape[0] if self.data_first_axis == 'T' else xs.shape[1]
       else:
         leaves, _ = tree_flatten(xs, is_leaf=lambda x: isinstance(x, bm.Array))
-        if isinstance(self.target.mode, bm.BatchingMode):
-          if self.time_major:
-            num_steps = [x.shape[0] for x in leaves]
-          else:
-            num_steps = [x.shape[1] for x in leaves]
-        else:
+        if self.data_first_axis == 'T':
           num_steps = [x.shape[0] for x in leaves]
+        else:
+          num_steps = [x.shape[1] for x in leaves]
         if len(set(num_steps)) != 1:
           raise ValueError(f'Number of time step is different across arrays in '
                            f'the provided "xs". We got {set(num_steps)}.')

brainpy/dyn/transform.py

@@ -71,13 +71,13 @@ class LoopOverTime(DynSysToBPObj):
   >>> over_time.reset_state(n_batch)
   (30, 128, 2)
   >>>
-  >>> hist_l3 = over_time(bm.random.rand(n_time, n_batch, n_in), time_major=True)
+  >>> hist_l3 = over_time(bm.random.rand(n_time, n_batch, n_in), data_first_axis='T')
   >>> print(hist_l3.shape)
   >>>
   >>> # monitor the "l1" layer state
   >>> over_time = bp.LoopOverTime(model, out_vars=model['l1'].state)
   >>> over_time.reset_state(n_batch)
-  >>> hist_l3, hist_l1 = over_time(bm.random.rand(n_time, n_batch, n_in), time_major=True)
+  >>> hist_l3, hist_l1 = over_time(bm.random.rand(n_time, n_batch, n_in), data_first_axis='T')
   >>> print(hist_l3.shape)
   (30, 128, 2)
   >>> print(hist_l1.shape)
@@ -148,7 +148,7 @@ def __call__(
       t0: float = 0.,
       dt: Optional[float] = None,
       shared_arg: Optional[Dict] = None,
-      time_major: bool = True
+      data_first_axis: str = 'T'
   ):
     """Forward propagation along the time or inputs.
 
@@ -164,7 +164,7 @@ def __call__(
     shared_arg: dict
       The shared arguments across the nodes.
       For instance, `shared_arg={'fit': False}` for the prediction phase.
-    time_major: bool
+    data_first_axis: str
       Denote whether the input data is time major.
       If so, we treat the data as `(time, batch, ...)` when the `target` is in Batching mode.
       Default is True.
@@ -174,6 +174,8 @@ def __call__(
     out: PyTree
       The accumulated outputs over time.
     """
+    assert data_first_axis in ['B', 'T']
+
     is_float(t0, 't0')
     is_float(dt, 'dt', allow_none=True)
     dt = bm.get_dt() if dt is None else dt
@@ -194,11 +196,11 @@ def __call__(
     else:
       inp_err_msg = ('\n'
                      'Input should be a Array PyTree with the shape '
-                     'of (B, T, ...) or (T, B, ...) with `time_major=True`, '
+                     'of (B, T, ...) or (T, B, ...) with `data_first_axis="T"`, '
                      'where B the batch size and T the time length.')
       xs, tree = tree_flatten(duration_or_xs, lambda a: isinstance(a, bm.Array))
       if isinstance(self.target.mode, bm.BatchingMode):
-        b_idx, t_idx = (1, 0) if time_major else (0, 1)
+        b_idx, t_idx = (1, 0) if data_first_axis == 'T' else (0, 1)
 
         try:
           batch = tuple(set([x.shape[b_idx] for x in xs]))
@@ -224,10 +226,10 @@ def __call__(
         if self.no_state:
           xs = [jnp.reshape(x, (length[0] * batch[0],) + x.shape[2:]) for x in xs]
         else:
-          if not time_major:
+          if data_first_axis == 'B':
             xs = [jnp.moveaxis(x, 0, 1) for x in xs]
         xs = tree_unflatten(tree, xs)
-        origin_shape = (length[0], batch[0]) if time_major else (batch[0], length[0])
+        origin_shape = (length[0], batch[0]) if data_first_axis == 'T' else (batch[0], length[0])
 
       else:
 

brainpy/train/back_propagation.py

@@ -491,17 +491,17 @@ def loss_fun(predicts, targets):
     Make the monitored results as NumPy arrays.
   logger: Any
     A file-like object (stream). Used to output the running results. Default is the current `sys.stdout`.
-  time_major: bool
-    To indicate whether the first axis is the batch size (``time_major=False``) or the
-    time length (``time_major=True``).
+  data_first_axis: str
+    To indicate whether the first axis is the batch size (``data_first_axis='B'``) or the
+    time length (``data_first_axis='T'``).
   """
 
   def _step_func_loss(self, shared_args, inputs, targets):
     num_step = self._get_input_time_step(xs=inputs)
     indices = jnp.arange(num_step, dtype=bm.int_)
     times = indices * self.dt + self.t0
     indices = indices + self.i0
-    if isinstance(self.target.mode, bm.BatchingMode) and not self.time_major:
+    if isinstance(self.target.mode, bm.BatchingMode) and self.data_first_axis == 'B':
       inputs = tree_map(lambda x: bm.moveaxis(x, 0, 1), inputs, is_leaf=lambda x: isinstance(x, bm.Array))
     inputs = (times, indices, inputs)
     outs, mons = self._predict(xs=inputs, shared_args=shared_args)
@@ -535,7 +535,7 @@ def _step_func_train(self, shared_args, inputs, targets):
     return res[1:]
 
   def _step_func_predict(self, shared, x=None):
-    assert not self.time_major, f'There is no time dimension when using the trainer {self.__class__.__name__}.'
+    assert self.data_first_axis == 'B', f'There is no time dimension when using the trainer {self.__class__.__name__}.'
 
     # input step
     self.target.clear_input()

brainpy/train/online.py

@@ -171,7 +171,7 @@ def fit(
     shared['t'] += self.t0
     shared['i'] += self.i0
 
-    if not self.time_major:
+    if self.data_first_axis == 'B':
       xs = tree_map(lambda x: bm.moveaxis(x, 0, 1),
                     xs,
                     is_leaf=lambda x: isinstance(x, bm.Array))