add Activation and Flatten class (#291)

add Activation and Flatten class

Author: chaoming0625

brainpy/dyn/layers/__init__.py

@@ -7,4 +7,5 @@
 from .rnncells import *
 from .conv import *
 from .normalization import *
-from .pooling import *
+from .pooling import *
+from .activate import *

brainpy/dyn/layers/activate.py

@@ -0,0 +1,34 @@
+from brainpy.dyn.base import DynamicalSystem
+from typing import Optional
+from brainpy.modes import Mode
+from typing import Callable
+
+
+class Activation(DynamicalSystem):
+  r"""Applies a activation to the inputs
+
+  Parameters:
+  ----------
+  activate_fun: Callable
+    The function of Activation
+  name: str, Optional
+    The name of the object
+  mode: Mode
+    Enable training this node or not. (default True).
+  """
+
+  def __init__(self,
+               activate_fun: Callable,
+               name: Optional[str] = None,
+               mode: Optional[Mode] = None,
+               **kwargs,
+      ):
+    super().__init__(name, mode)
+    self.activate_fun = activate_fun
+    self.kwargs = kwargs
+
+  def update(self, sha, x):
+    return self.activate_fun(x, **self.kwargs)
+
+  def reset_state(self, batch_size=None):
+    pass

brainpy/dyn/layers/linear.py

@@ -9,12 +9,13 @@
 from brainpy.dyn.base import DynamicalSystem
 from brainpy.errors import MathError
 from brainpy.initialize import XavierNormal, ZeroInit, Initializer, parameter
-from brainpy.modes import Mode, TrainingMode, training
+from brainpy.modes import  Mode, TrainingMode, BatchingMode, training, batching
 from brainpy.tools.checking import check_initializer
 from brainpy.types import Array
 
 __all__ = [
   'Dense',
+  'Flatten'
 ]
 
 
@@ -188,3 +189,29 @@ def offline_fit(self,
       bias, Wff = bm.split(weights, [1])
       self.W.value = Wff
       self.b.value = bias[0]
+
+
+class Flatten(DynamicalSystem):
+  r"""Flattens a contiguous range of dims into 2D or 1D.
+
+  Parameters:
+  ----------
+  name: str, Optional
+    The name of the object
+  mode: Mode
+    Enable training this node or not. (default True)
+  """
+  def __init__(self,
+               name: Optional[str] = None,
+               mode: Optional[Mode] = batching,
+      ):
+    super().__init__(name, mode)
+  
+  def update(self, shr, x):
+    if isinstance(self.mode, BatchingMode):
+      return x.reshape((x.shape[0], -1))
+    else:
+      return x.flatten()
+  
+  def reset_state(self, batch_size=None):
+    pass