update apis

Author: chaoming0625

brainpy/dyn/base.py

@@ -996,7 +996,7 @@ def _init_weights(
   ) -> Union[float, Array]:
     if comp_method not in ['sparse', 'dense']:
       raise ValueError(f'"comp_method" must be in "sparse" and "dense", but we got {comp_method}')
-    if sparse_data not in ['csr', 'ij']:
+    if sparse_data not in ['csr', 'ij', 'coo']:
       raise ValueError(f'"sparse_data" must be in "csr" and "ij", but we got {sparse_data}')
     if self.conn is None:
       raise ValueError(f'Must provide "conn" when initialize the model {self.name}')
@@ -1014,7 +1014,7 @@ def _init_weights(
       if comp_method == 'sparse':
         if sparse_data == 'csr':
           conn_mask = self.conn.require('pre2post')
-        elif sparse_data == 'ij':
+        elif sparse_data in ['ij', 'coo']:
           conn_mask = self.conn.require('post_ids', 'pre_ids')
         else:
           ValueError(f'Unknown sparse data type: {sparse_data}')

brainpy/dyn/layers/activate.py

@@ -1,28 +1,30 @@
-from brainpy.dyn.base import DynamicalSystem
-from typing import Optional
-from brainpy.modes import Mode
 from typing import Callable
+from typing import Optional
+
+from brainpy.dyn.base import DynamicalSystem
+from brainpy.modes import Mode, training
 
 
 class Activation(DynamicalSystem):
-  r"""Applies a activation to the inputs
+  r"""Applies an activation function to the inputs
 
   Parameters:
   ----------
-  activate_fun: Callable
+  activate_fun: Callable, function
     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,
-      ):
+  def __init__(
+      self,
+      activate_fun: Callable,
+      name: Optional[str] = None,
+      mode: Mode = training,
+      **kwargs,
+  ):
     super().__init__(name, mode)
     self.activate_fun = activate_fun
     self.kwargs = kwargs

brainpy/dyn/layers/linear.py

@@ -9,7 +9,7 @@
 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, BatchingMode, training, batching
+from brainpy.modes import Mode, TrainingMode, BatchingMode, training, batching
 from brainpy.tools.checking import check_initializer
 from brainpy.types import Array
 
@@ -201,17 +201,19 @@ class Flatten(DynamicalSystem):
   mode: Mode
     Enable training this node or not. (default True)
   """
-  def __init__(self,
-               name: Optional[str] = None,
-               mode: Optional[Mode] = batching,
-      ):
+
+  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
+    pass