Refine JIT wrappers for new JAX for comaptiblity with jax>=0.8.2

brainpy/__init__.py

@@ -17,15 +17,14 @@
 __version__ = "2.7.5"
 __version_info__ = tuple(map(int, __version__.split(".")))
 
-
 from brainpy import _errors as errors
-from brainpy import mixin
 # fundamental supporting modules
 from brainpy import check, tools
 #  Part: Math Foundation  #
 # ----------------------- #
 # math foundation
 from brainpy import math
+from brainpy import mixin
 #  Part: Toolbox  #
 # --------------- #
 # modules of toolbox
@@ -109,11 +108,11 @@
 # ---------------- #
 from brainpy.train.base import (DSTrainer as DSTrainer, )
 from brainpy.train.back_propagation import (BPTT as BPTT,
-                                                     BPFF as BPFF, )
+                                            BPFF as BPFF, )
 from brainpy.train.online import (OnlineTrainer as OnlineTrainer,
-                                           ForceTrainer as ForceTrainer, )
+                                  ForceTrainer as ForceTrainer, )
 from brainpy.train.offline import (OfflineTrainer as OfflineTrainer,
-                                            RidgeTrainer as RidgeTrainer, )
+                                   RidgeTrainer as RidgeTrainer, )
 
 #  Part: Analysis  #
 # ---------------- #
@@ -147,42 +146,7 @@
 
 optimizers = optim
 
-try:
-    import brainpy.state as state
-except:
-    pass
-
 
+# New package
+from brainpy import state
 
-if __name__ == '__main__':
-    connect
-    initialize,  # weight initialization
-    optim,  # gradient descent optimizers
-    losses,  # loss functions
-    measure,  # methods for data analysis
-    inputs,  # methods for generating input currents
-    encoding,  # encoding schema
-    checkpoints,  # checkpoints
-    check,  # error checking
-    mixin,  # mixin classes
-    algorithms,  # online or offline training algorithms
-    check, tools, errors, math
-    BrainPyObject,
-    integrators, ode, sde, fde
-    Integrator, JointEq, IntegratorRunner, odeint, sdeint, fdeint
-    DynamicalSystem, DynSysGroup, Sequential, Dynamic, Projection
-    receive_update_input, receive_update_output, not_receive_update_input, not_receive_update_output
-    VarDelay
-    dnn, layers, dyn
-    NeuGroup, NeuGroupNS
-    share
-    reset_level, reset_state, save_state, load_state, clear_input
-    DSRunner, LoopOverTime, running
-    DSTrainer, BPTT, BPFF, OnlineTrainer, ForceTrainer,
-    OfflineTrainer, RidgeTrainer
-    analysis
-    visualize
-    train
-    channels, neurons, synapses, rates, synouts, synplast
-    Base
-    ArrayCollector, Collector, errors

brainpy/analysis/__init__.py

@@ -34,4 +34,3 @@
 from .highdim.slow_points import *
 from .lowdim.lowdim_bifurcation import *
 from .lowdim.lowdim_phase_plane import *
-

brainpy/analysis/highdim/slow_points.py

@@ -25,8 +25,8 @@
 from jax.scipy.optimize import minimize
 
 import brainpy.math as bm
-from brainpy._errors import AnalyzerError, UnsupportedError
 from brainpy import optim, losses
+from brainpy._errors import AnalyzerError, UnsupportedError
 from brainpy.analysis import utils, base, constants
 from brainpy.context import share
 from brainpy.deprecations import _input_deprecate_msg

brainpy/channels.py

@@ -17,11 +17,9 @@
 This module has been deprecated since brainpy>=2.4.0. Use ``brainpy.dyn`` module instead.
 """
 
-
 from .dyn.channels import *
 from .dyn.ions import *
 
-
 if __name__ == '__main__':
     IL
     Potassium

brainpy/check.py

@@ -647,4 +647,3 @@ def true_err_fun(arg, transforms):
     cond(remove_vmap(as_jax(pred)),
          lambda: jax.pure_callback(true_err_fun, None),
          lambda: None)
-

brainpy/connect/base.py

@@ -20,8 +20,8 @@
 import jax.numpy as jnp
 import numpy as onp
 
-from brainpy._errors import ConnectorError
 from brainpy import tools, math as bm
+from brainpy._errors import ConnectorError
 
 __all__ = [
     # the connection types

brainpy/connect/custom_conn.py

@@ -17,8 +17,8 @@
 import jax.numpy as jnp
 import numpy as np
 
-from brainpy._errors import ConnectorError
 from brainpy import math as bm, tools
+from brainpy._errors import ConnectorError
 from .base import *
 
 __all__ = [

brainpy/context.py

@@ -21,6 +21,7 @@
 from typing import Any, Union
 
 import brainstate
+
 from brainpy.math.defaults import env
 from brainpy.tools.dicts import DotDict
 

brainpy/delay.py

@@ -24,13 +24,13 @@
 import jax.numpy as jnp
 import numpy as np
 
-from brainpy.mixin import ParamDesc, ReturnInfo, JointType, SupportAutoDelay
 from brainpy import check, math as bm
 from brainpy.check import jit_error
 from brainpy.context import share
 from brainpy.dynsys import DynamicalSystem
 from brainpy.initialize import variable_
 from brainpy.math.delayvars import ROTATE_UPDATE, CONCAT_UPDATE
+from brainpy.mixin import ParamDesc, ReturnInfo, JointType, SupportAutoDelay
 
 __all__ = [
     'Delay',

brainpy/dnn/__init__.py

@@ -22,4 +22,3 @@
 from .linear import *
 from .normalization import *
 from .pooling import *
-

brainpy/dnn/conv.py

@@ -15,9 +15,9 @@
 # ==============================================================================
 from typing import Union, Tuple, Optional, Sequence, Callable
 
+import brainstate
 from jax import lax
 
-import brainstate
 from brainpy import math as bm, tools
 from brainpy.dnn.base import Layer
 from brainpy.initialize import Initializer, XavierNormal, ZeroInit, parameter

brainpy/dnn/linear.py

@@ -19,18 +19,18 @@
 import jax
 import jax.numpy as jnp
 import numpy as np
-from brainevent._csr_impl_plasticity import csr_on_pre, csr2csc_on_post
-from brainevent._dense_impl_plasticity import dense_on_pre, dense_on_post
+from brainevent import csr_on_pre, csr2csc_on_post
+from brainevent import dense_on_pre, dense_on_post
 
-from brainpy._errors import MathError
-from brainpy.mixin import SupportOnline, SupportOffline, SupportSTDP
 from brainpy import connect, initialize as init
 from brainpy import math as bm
+from brainpy._errors import MathError
 from brainpy.check import is_initializer
 from brainpy.connect import csr2csc
 from brainpy.context import share
 from brainpy.dnn.base import Layer
 from brainpy.initialize import XavierNormal, ZeroInit, Initializer, parameter
+from brainpy.mixin import SupportOnline, SupportOffline, SupportSTDP
 from brainpy.types import ArrayType, Sharding
 
 __all__ = [

brainpy/dyn/base.py

@@ -13,8 +13,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-from brainpy.mixin import SupportAutoDelay, ParamDesc
 from brainpy.dynsys import Dynamic
+from brainpy.mixin import SupportAutoDelay, ParamDesc
 
 __all__ = [
     'NeuDyn', 'SynDyn', 'IonChaDyn',

brainpy/dyn/channels/__init__.py

@@ -23,4 +23,3 @@
 from .potassium_compatible import *
 from .sodium import *
 from .sodium_compatible import *
-

brainpy/dyn/channels/base.py

@@ -13,9 +13,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-from brainpy.mixin import TreeNode
 from brainpy.dyn.base import IonChaDyn
 from brainpy.dyn.neurons.hh import HHTypedNeuron
+from brainpy.mixin import TreeNode
 
 __all__ = [
     'IonChannel',

brainpy/dyn/channels/potassium_calcium.py

@@ -20,12 +20,12 @@
 from typing import Union, Callable, Optional
 
 import brainpy.math as bm
-from brainpy.mixin import JointType
 from brainpy.context import share
 from brainpy.dyn.ions.calcium import Calcium
 from brainpy.dyn.ions.potassium import Potassium
 from brainpy.initialize import Initializer, parameter, variable
 from brainpy.integrators.ode.generic import odeint
+from brainpy.mixin import JointType
 from brainpy.types import Shape, ArrayType
 from .calcium import CalciumChannel
 from .potassium import PotassiumChannel

brainpy/dyn/ions/base.py

@@ -15,12 +15,13 @@
 # ==============================================================================
 from typing import Union, Optional, Dict, Sequence, Callable
 
+from brainstate.mixin import _JointGenericAlias
+
 import brainpy.math as bm
-from brainpy.mixin import Container, TreeNode
 from brainpy.dyn.base import IonChaDyn
 from brainpy.dyn.neurons.hh import HHTypedNeuron
+from brainpy.mixin import Container, TreeNode
 from brainpy.types import Shape
-from brainstate.mixin import _JointGenericAlias
 
 __all__ = [
     'MixIons',

brainpy/dyn/neurons/hh.py

@@ -17,14 +17,14 @@
 from typing import Union, Callable, Optional
 
 import brainpy.math as bm
-from brainpy.mixin import Container, TreeNode
 from brainpy.check import is_initializer
 from brainpy.context import share
 from brainpy.dyn.base import NeuDyn, IonChaDyn
 from brainpy.initialize import OneInit
 from brainpy.initialize import Uniform, variable_, noise as init_noise
 from brainpy.integrators import JointEq
 from brainpy.integrators import odeint, sdeint
+from brainpy.mixin import Container, TreeNode
 from brainpy.types import ArrayType
 from brainpy.types import Shape
 

brainpy/dyn/others/input.py

@@ -19,12 +19,12 @@
 import jax
 import jax.numpy as jnp
 
-from brainpy.mixin import ReturnInfo
 from brainpy import math as bm
 from brainpy.context import share
 from brainpy.dyn.base import NeuDyn
 from brainpy.dyn.utils import get_spk_type
 from brainpy.initialize import parameter, variable_
+from brainpy.mixin import ReturnInfo
 from brainpy.types import Shape, ArrayType
 
 __all__ = [

brainpy/dyn/outs/base.py

@@ -15,8 +15,8 @@
 from typing import Optional
 
 import brainpy.math as bm
-from brainpy.mixin import ParamDesc, BindCondData
 from brainpy.dynsys import DynamicalSystem
+from brainpy.mixin import ParamDesc, BindCondData
 
 __all__ = [
     'SynOut'

brainpy/dyn/projections/align_post.py

@@ -14,11 +14,11 @@
 # ==============================================================================
 from typing import Optional, Callable, Union
 
-from brainpy.mixin import (JointType, ParamDescriber, SupportAutoDelay, BindCondData, AlignPost)
 from brainpy import math as bm, check
 from brainpy.delay import (delay_identifier,
-                                    register_delay_by_return)
+                           register_delay_by_return)
 from brainpy.dynsys import DynamicalSystem, Projection
+from brainpy.mixin import (JointType, ParamDescriber, SupportAutoDelay, BindCondData, AlignPost)
 
 __all__ = [
     'HalfProjAlignPostMg', 'FullProjAlignPostMg',

brainpy/dyn/projections/align_pre.py

@@ -14,10 +14,10 @@
 # ==============================================================================
 from typing import Optional, Union
 
-from brainpy.mixin import (JointType, ParamDescriber, SupportAutoDelay, BindCondData)
 from brainpy import math as bm, check
 from brainpy.delay import (Delay, DelayAccess, init_delay_by_return, register_delay_by_return)
 from brainpy.dynsys import DynamicalSystem, Projection
+from brainpy.mixin import (JointType, ParamDescriber, SupportAutoDelay, BindCondData)
 from .utils import _get_return
 
 __all__ = [

brainpy/dyn/projections/base.py

@@ -12,8 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-from brainpy.mixin import ReturnInfo
 from brainpy import math as bm
+from brainpy.mixin import ReturnInfo
 
 
 def _get_return(return_info):

brainpy/dyn/projections/delta.py

@@ -14,10 +14,10 @@
 # ==============================================================================
 from typing import Optional, Union
 
-from brainpy.mixin import (JointType, SupportAutoDelay)
 from brainpy import math as bm, check
 from brainpy.delay import (delay_identifier, register_delay_by_return)
 from brainpy.dynsys import DynamicalSystem, Projection
+from brainpy.mixin import (JointType, SupportAutoDelay)
 
 __all__ = [
     'HalfProjDelta', 'FullProjDelta',

brainpy/dyn/projections/inputs.py

@@ -16,11 +16,11 @@
 from typing import Any
 from typing import Union, Optional
 
-from brainpy.mixin import SupportAutoDelay
 from brainpy import check, math as bm
 from brainpy.context import share
 from brainpy.dynsys import Dynamic
 from brainpy.dynsys import Projection
+from brainpy.mixin import SupportAutoDelay
 from brainpy.types import Shape
 
 __all__ = [

brainpy/dyn/projections/plasticity.py

@@ -14,12 +14,12 @@
 # ==============================================================================
 from typing import Optional, Callable, Union
 
-from brainpy.mixin import (JointType, ParamDescriber, SupportAutoDelay,
-                           BindCondData, AlignPost, SupportSTDP)
 from brainpy import math as bm, check
 from brainpy.delay import register_delay_by_return
 from brainpy.dyn.synapses.abstract_models import Expon
 from brainpy.dynsys import DynamicalSystem, Projection
+from brainpy.mixin import (JointType, ParamDescriber, SupportAutoDelay,
+                           BindCondData, AlignPost, SupportSTDP)
 from brainpy.types import ArrayType
 from .align_post import (align_post_add_bef_update, )
 from .align_pre import (align_pre2_add_bef_update, )

brainpy/dyn/projections/utils.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 from brainpy.mixin import ReturnInfo
-from brainpy import math as bm
+import brainpy.math as bm
 
 
 def _get_return(return_info):

brainpy/dyn/projections/vanilla.py

@@ -14,9 +14,9 @@
 # ==============================================================================
 from typing import Optional
 
-from brainpy.mixin import (JointType, BindCondData)
 from brainpy import math as bm, check
 from brainpy.dynsys import DynamicalSystem, Projection
+from brainpy.mixin import (JointType, BindCondData)
 
 __all__ = [
     'VanillaProj',

brainpy/dyn/rates/populations.py

@@ -23,11 +23,11 @@
 from brainpy.dyn.base import NeuDyn
 from brainpy.dyn.others.noise import OUProcess
 from brainpy.initialize import (Initializer,
-                                         Uniform,
-                                         parameter,
-                                         variable,
-                                         variable_,
-                                         ZeroInit)
+                                Uniform,
+                                parameter,
+                                variable,
+                                variable_,
+                                ZeroInit)
 from brainpy.integrators.joint_eq import JointEq
 from brainpy.integrators.ode.generic import odeint
 from brainpy.types import Shape, ArrayType