update installation setup (#269)

update installation setup

Author: chaoming0625

brainpy/__init__.py

@@ -1,17 +1,38 @@
 # -*- coding: utf-8 -*-
 
-__version__ = "2.2.3"
+__version__ = "2.2.3.1"
 
 
 try:
   import jaxlib
   del jaxlib
 except ModuleNotFoundError:
   raise ModuleNotFoundError(
-    'Please install jaxlib. See '
-    'https://brainpy.readthedocs.io/en/latest/quickstart/installation.html#dependency-2-jax '
-    'for installation instructions.'
-  ) from None
+    '''
+
+BrainPy needs jaxlib, please install jaxlib. 
+
+1. If you are using Windows system, install jaxlib through
+
+   >>> pip install jaxlib -f https://whls.blob.core.windows.net/unstable/index.html
+
+2. If you are using macOS platform, install jaxlib through
+
+   >>> pip install jaxlib -f https://storage.googleapis.com/jax-releases/jax_releases.html
+
+3. If you are using Linux platform, install jaxlib through
+
+   >>> pip install jaxlib -f https://storage.googleapis.com/jax-releases/jax_releases.html
+
+4. If you are using Linux + CUDA platform, install jaxlib through
+
+   >>> pip install jaxlib -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+
+Note that the versions of "jax" and "jaxlib" should be consistent, like "jax=0.3.14", "jaxlib=0.3.14".  
+
+More detail installation instruction, please see https://brainpy.readthedocs.io/en/latest/quickstart/installation.html#dependency-2-jax 
+    
+    ''') from None
 
 
 # fundamental modules

brainpy/dyn/layers/conv.py

@@ -6,7 +6,7 @@
 import brainpy.math as bm
 from brainpy.dyn.base import DynamicalSystem
 from brainpy.initialize import XavierNormal, ZeroInit, parameter
-from brainpy.modes import Mode, TrainingMode, NormalMode, training, check
+from brainpy.modes import Mode, TrainingMode, training
 
 __all__ = [
   'GeneralConv',

brainpy/dyn/layers/normalization.py

@@ -9,7 +9,7 @@
 import brainpy.math as bm
 from brainpy.initialize import ZeroInit, OneInit, Initializer, parameter
 from brainpy.dyn.base import DynamicalSystem
-from brainpy.modes import Mode, TrainingMode, NormalMode, training, check
+from brainpy.modes import Mode, TrainingMode, NormalMode, training, check_mode
 
 __all__ = [
   'BatchNorm',

brainpy/dyn/layers/nvar.py

@@ -8,7 +8,7 @@
 
 import brainpy.math as bm
 from brainpy.dyn.base import DynamicalSystem
-from brainpy.modes import Mode, NormalMode, BatchingMode, batching, check
+from brainpy.modes import Mode, NormalMode, BatchingMode, batching, check_mode
 from brainpy.tools.checking import (check_integer, check_sequence)
 
 __all__ = [
@@ -73,7 +73,7 @@ def __init__(
       name: str = None,
   ):
     super(NVAR, self).__init__(mode=mode, name=name)
-    check(self.mode, (BatchingMode, NormalMode), self.__class__.__name__)
+    check_mode(self.mode, (BatchingMode, NormalMode), self.__class__.__name__)
 
     # parameters
     order = tuple() if order is None else order

brainpy/dyn/layers/pooling.py

@@ -4,7 +4,7 @@
 import jax.lax
 import brainpy.math as bm
 from brainpy.dyn.base import DynamicalSystem
-from brainpy.modes import Mode, TrainingMode, NormalMode, training, check
+from brainpy.modes import Mode, training
 
 __all__ = [
   'Pool',

brainpy/dyn/neurons/biological_models.py

@@ -8,7 +8,7 @@
 from brainpy.integrators.joint_eq import JointEq
 from brainpy.integrators.ode import odeint
 from brainpy.integrators.sde import sdeint
-from brainpy.modes import Mode, BatchingMode, TrainingMode, NormalMode, normal, check
+from brainpy.modes import Mode, BatchingMode, TrainingMode, NormalMode, normal, check_mode
 from brainpy.tools.checking import check_initializer
 from brainpy.types import Shape, Array
 
@@ -219,7 +219,7 @@ def __init__(
                              keep_size=keep_size,
                              name=name,
                              mode=mode)
-    check(self.mode, (BatchingMode, NormalMode), self.__class__.__name__)
+    check_mode(self.mode, (BatchingMode, NormalMode), self.__class__.__name__)
 
     # parameters
     self.ENa = parameter(ENa, self.varshape, allow_none=False)
@@ -427,7 +427,7 @@ def __init__(
                                       keep_size=keep_size,
                                       name=name,
                                       mode=mode)
-    check(self.mode, (BatchingMode, NormalMode), self.__class__)
+    check_mode(self.mode, (BatchingMode, NormalMode), self.__class__)
 
     # params
     self.V_Ca = parameter(V_Ca, self.varshape, allow_none=False)
@@ -685,7 +685,7 @@ def __init__(
                                             keep_size=keep_size,
                                             name=name,
                                             mode=mode)
-    check(self.mode, (NormalMode, BatchingMode), self.__class__)
+    check_mode(self.mode, (NormalMode, BatchingMode), self.__class__)
 
     # conductance parameters
     self.gAHP = parameter(gAHP, self.varshape, allow_none=False)
@@ -994,7 +994,7 @@ def __init__(
   ):
     # initialization
     super(WangBuzsakiModel, self).__init__(size=size, keep_size=keep_size, name=name, mode=mode)
-    check(self.mode, (BatchingMode, NormalMode), self.__class__)
+    check_mode(self.mode, (BatchingMode, NormalMode), self.__class__)
 
     # parameters
     self.ENa = parameter(ENa, self.varshape, allow_none=False)

brainpy/dyn/neurons/input_groups.py

@@ -11,6 +11,7 @@
 from brainpy.modes import Mode, BatchingMode, normal
 from brainpy.types import Shape, Array
 
+
 __all__ = [
   'InputGroup',
   'OutputGroup',
@@ -205,3 +206,4 @@ def reset(self, batch_size=None):
 
   def reset_state(self, batch_size=None):
     self.spike.value = variable(lambda s: bm.zeros(s, dtype=bool), batch_size, self.varshape)
+