Refine JIT wrappers for new JAX for comaptiblity with jax>=0.8.2 (#809)

* refactor: clean up imports and remove unnecessary newlines

* refactor: remove unused imports and clean up main execution block

* feat: add brainpy_state module and update dependencies

Author: chaoming0625

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 *
-