feat: data pass of the Node is default SingleData

Author: chaoming0625

brainpy/nn/__init__.py

@@ -3,7 +3,7 @@
 """Neural Networks (nn)"""
 
 from .base import *
-from .constants import *
+from .datatypes import *
 from .graph_flow import *
 from .nodes import *
 from .graph_flow import *

brainpy/nn/base.py

@@ -28,9 +28,7 @@
                             MathError)
 from brainpy.nn.algorithms.offline import OfflineAlgorithm
 from brainpy.nn.algorithms.online import OnlineAlgorithm
-from brainpy.nn.constants import (PASS_SEQUENCE,
-                                  DATA_PASS_FUNC,
-                                  DATA_PASS_TYPES)
+from brainpy.nn.datatypes import (DataType, SingleData, MultipleData)
 from brainpy.nn.graph_flow import (find_senders_and_receivers,
                                    find_entries_and_exits,
                                    detect_cycle,
@@ -83,13 +81,13 @@ def feedback(self):
 class Node(Base):
   """Basic Node class for neural network building in BrainPy."""
 
-  '''Support multiple types of data pass, including "PASS_SEQUENCE" (by default), 
-  "PASS_NAME_DICT", "PASS_NODE_DICT" and user-customized type which registered 
-  by ``brainpy.nn.register_data_pass_type()`` function. 
+  '''Support multiple types of data pass, including "PassOnlyOne" (by default), 
+  "PassSequence", "PassNameDict", etc. and user-customized type which inherits 
+  from basic "SingleData" or "MultipleData". 
   
   This setting will change the feedforward/feedback input data which pass into 
   the "call()" function and the sizes of the feedforward/feedback input data.'''
-  data_pass_type = PASS_SEQUENCE
+  data_pass = SingleData()
 
   '''Offline fitting method.'''
   offline_fit_by: Union[Callable, OfflineAlgorithm]
@@ -115,11 +113,10 @@ def __init__(
     self._trainable = trainable
     self._state = None  # the state of the current node
     self._fb_output = None  # the feedback output of the current node
-    # data pass function
-    if self.data_pass_type not in DATA_PASS_FUNC:
-      raise ValueError(f'Unsupported data pass type {self.data_pass_type}. '
-                       f'Only support {DATA_PASS_TYPES}')
-    self.data_pass_func = DATA_PASS_FUNC[self.data_pass_type]
+    # data pass
+    if not isinstance(self.data_pass, DataType):
+      raise ValueError(f'Unsupported data pass type {type(self.data_pass)}. '
+                       f'Only support {DataType.__class__}')
 
     # super initialization
     super(Node, self).__init__(name=name)
@@ -129,11 +126,10 @@ def __init__(
       self._feedforward_shapes = {self.name: (None,) + tools.to_size(input_shape)}
 
   def __repr__(self):
-    name = type(self).__name__
-    prefix = ' ' * (len(name) + 1)
-    line1 = f"{name}(name={self.name}, forwards={self.feedforward_shapes}, \n"
-    line2 = f"{prefix}feedbacks={self.feedback_shapes}, output={self.output_shape})"
-    return line1 + line2
+    return (f"{type(self).__name__}(name={self.name}, "
+            f"forwards={self.feedforward_shapes}, "
+            f"feedbacks={self.feedback_shapes}, "
+            f"output={self.output_shape})")
 
   def __call__(self, *args, **kwargs) -> Tensor:
     """The main computation function of a Node.
@@ -298,7 +294,7 @@ def trainable(self, value: bool):
   @property
   def feedforward_shapes(self):
     """Input data size."""
-    return self.data_pass_func(self._feedforward_shapes)
+    return self.data_pass.filter(self._feedforward_shapes)
 
   @feedforward_shapes.setter
   def feedforward_shapes(self, size):
@@ -324,7 +320,7 @@ def set_feedforward_shapes(self, feedforward_shapes: Dict):
   @property
   def feedback_shapes(self):
     """Output data size."""
-    return self.data_pass_func(self._feedback_shapes)
+    return self.data_pass.filter(self._feedback_shapes)
 
   @feedback_shapes.setter
   def feedback_shapes(self, size):
@@ -530,8 +526,8 @@ def _check_inputs(self, ff, fb=None):
                              f'batch size by ".initialize(num_batch)", or change the data '
                              f'consistent with the data batch size {self.state.shape[0]}.')
     # data
-    ff = self.data_pass_func(ff)
-    fb = self.data_pass_func(fb)
+    ff = self.data_pass.filter(ff)
+    fb = self.data_pass.filter(fb)
     return ff, fb
 
   def _call(self,
@@ -747,6 +743,8 @@ def set_state(self, state):
 class Network(Node):
   """Basic Network class for neural network building in BrainPy."""
 
+  data_pass = MultipleData('sequence')
+
   def __init__(self,
                nodes: Optional[Sequence[Node]] = None,
                ff_edges: Optional[Sequence[Tuple[Node]]] = None,
@@ -1145,8 +1143,8 @@ def _check_inputs(self, ff, fb=None):
         check_shape_except_batch(size, fb[k].shape)
 
     # data transformation
-    ff = self.data_pass_func(ff)
-    fb = self.data_pass_func(fb)
+    ff = self.data_pass.filter(ff)
+    fb = self.data_pass.filter(fb)
     return ff, fb
 
   def _call(self,
@@ -1208,12 +1206,12 @@ def _call(self,
   def _call_a_node(self, node, ff, fb, monitors, forced_states,
                    parent_outputs, children_queue, ff_senders,
                    **shared_kwargs):
-    ff = node.data_pass_func(ff)
+    ff = node.data_pass.filter(ff)
     if f'{node.name}.inputs' in monitors:
       monitors[f'{node.name}.inputs'] = ff
     # get the output results
     if len(fb):
-      fb = node.data_pass_func(fb)
+      fb = node.data_pass.filter(fb)
       if f'{node.name}.feedbacks' in monitors:
         monitors[f'{node.name}.feedbacks'] = fb
       parent_outputs[node] = node.forward(ff, fb, **shared_kwargs)
@@ -1440,7 +1438,7 @@ def plot_node_graph(self,
     if len(nodes_untrainable):
       proxie.append(Line2D([], [], color='white', marker='o',
                            markerfacecolor=untrainable_color))
-      labels.append('Untrainable')
+      labels.append('Nontrainable')
     if len(ff_edges):
       proxie.append(Line2D([], [], color=ff_color, linewidth=2))
       labels.append('Feedforward')

brainpy/nn/nodes/ANN/batch_norm.py

@@ -6,17 +6,15 @@
 """
 
 
-from typing import Sequence, Optional, Dict, Callable, Union
+from typing import Union
 
 import jax.nn
 import jax.numpy as jnp
 
-import brainpy.math as bm
 import brainpy
+import brainpy.math as bm
 from brainpy.initialize import ZeroInit, OneInit, Initializer
 from brainpy.nn.base import Node
-from brainpy.nn.constants import PASS_ONLY_ONE
-
 
 __all__ = [
   'BatchNorm',
@@ -40,8 +38,6 @@ class BatchNorm(Node):
   beta_init: an initializer generating the original translation matrix
   gamma_init: an initializer generating the original scaling matrix
   """
-  data_pass_type = PASS_ONLY_ONE
-
   def __init__(self,
                axis: Union[int, tuple, list],
                epsilon: float = 1e-5,

brainpy/nn/nodes/ANN/dropout.py

@@ -2,7 +2,6 @@
 
 import brainpy.math as bm
 from brainpy.nn.base import Node
-from brainpy.nn.constants import PASS_ONLY_ONE
 
 __all__ = [
   'Dropout'
@@ -37,8 +36,6 @@ class Dropout(Node):
          neural networks from overfitting." The journal of machine learning
          research 15.1 (2014): 1929-1958.
   """
-  data_pass_type = PASS_ONLY_ONE
-
   def __init__(self, prob, seed=None, **kwargs):
     super(Dropout, self).__init__(**kwargs)
     self.prob = prob

brainpy/nn/nodes/ANN/rnn_cells.py

@@ -11,6 +11,7 @@
                                 init_param,
                                 Initializer)
 from brainpy.nn.base import RecurrentNode
+from brainpy.nn.datatypes import MultipleData
 from brainpy.tools.checking import (check_integer,
                                     check_initializer,
                                     check_shape_consistency)
@@ -55,6 +56,7 @@ class VanillaRNN(RecurrentNode):
     Whether set the node is trainable.
 
   """
+  data_pass = MultipleData('sequence')
 
   def __init__(
       self,
@@ -169,6 +171,7 @@ class GRU(RecurrentNode):
          evaluation of gated recurrent neural networks on sequence modeling.
          arXiv preprint arXiv:1412.3555.
   """
+  data_pass = MultipleData('sequence')
 
   def __init__(
       self,
@@ -302,6 +305,7 @@ class LSTM(RecurrentNode):
          exploration of recurrent network architectures." In International conference
          on machine learning, pp. 2342-2350. PMLR, 2015.
   """
+  data_pass = MultipleData('sequence')
 
   def __init__(
       self,
@@ -391,16 +395,16 @@ def c(self, value):
 
 
 class ConvNDLSTM(RecurrentNode):
-  pass
+  data_pass = MultipleData('sequence')
 
 
 class Conv1DLSTM(ConvNDLSTM):
-  pass
+  data_pass = MultipleData('sequence')
 
 
 class Conv2DLSTM(ConvNDLSTM):
-  pass
+  data_pass = MultipleData('sequence')
 
 
 class Conv3DLSTM(ConvNDLSTM):
-  pass
+  data_pass = MultipleData('sequence')

brainpy/nn/nodes/RC/linear_readout.py

@@ -5,6 +5,7 @@
 import brainpy.math as bm
 from brainpy.errors import MathError
 from brainpy.initialize import Initializer
+from brainpy.nn.datatypes import MultipleData
 from brainpy.nn.nodes.base.dense import Dense
 from brainpy.tools.checking import check_shape_consistency
 
@@ -27,6 +28,7 @@ class LinearReadout(Dense):
   trainable: bool
     Default is true.
   """
+  data_pass = MultipleData('sequence')
 
   def __init__(self, num_unit: int, **kwargs):
     super(LinearReadout, self).__init__(num_unit=num_unit, **kwargs)

brainpy/nn/nodes/RC/nvar.py

@@ -7,6 +7,7 @@
 
 import brainpy.math as bm
 from brainpy.nn.base import RecurrentNode
+from brainpy.nn.datatypes import MultipleData
 from brainpy.tools.checking import (check_shape_consistency,
                                     check_integer,
                                     check_sequence)
@@ -61,6 +62,7 @@ class NVAR(RecurrentNode):
          https://doi.org/10.1038/s41467-021-25801-2
 
   """
+  data_pass = MultipleData('sequence')
 
   def __init__(
       self,

brainpy/nn/nodes/RC/reservoir.py

@@ -5,6 +5,7 @@
 import brainpy.math as bm
 from brainpy.initialize import Normal, ZeroInit, Initializer, init_param
 from brainpy.nn.base import RecurrentNode
+from brainpy.nn.datatypes import MultipleData
 from brainpy.tools.checking import (check_shape_consistency,
                                     check_float,
                                     check_initializer,
@@ -90,6 +91,7 @@ class Reservoir(RecurrentNode):
   .. [1] Lukoševičius, Mantas. "A practical guide to applying echo state networks."
          Neural networks: Tricks of the trade. Springer, Berlin, Heidelberg, 2012. 659-686.
   """
+  data_pass = MultipleData('sequence')
 
   def __init__(
       self,

brainpy/nn/nodes/base/activation.py

@@ -4,7 +4,6 @@
 
 from brainpy.math import activations
 from brainpy.nn.base import Node
-from brainpy.nn.constants import PASS_ONLY_ONE
 
 __all__ = [
   'Activation'
@@ -22,8 +21,6 @@ class Activation(Node):
     The settings for the activation function.
   """
 
-  data_pass_type = PASS_ONLY_ONE
-
   def __init__(self,
                activation: str = 'relu',
                fun_setting: Optional[Dict[str, Any]] = None,

brainpy/nn/nodes/base/dense.py

@@ -9,6 +9,7 @@
 from brainpy.errors import MathError
 from brainpy.initialize import XavierNormal, ZeroInit, Initializer, init_param
 from brainpy.nn.base import Node
+from brainpy.nn.datatypes import MultipleData
 from brainpy.tools.checking import (check_shape_consistency,
                                     check_initializer)
 from brainpy.types import Tensor
@@ -40,6 +41,8 @@ class GeneralDense(Node):
     Enable training this node or not. (default True)
   """
 
+  data_pass = MultipleData('sequence')
+
   def __init__(
       self,
       num_unit: int,
@@ -123,6 +126,7 @@ class Dense(GeneralDense):
   trainable: bool
     Enable training this node or not. (default True)
   """
+  data_pass = MultipleData('sequence')
 
   def __init__(
       self,