Skip to content

Commit 1847bd3

Browse files
chaoming0625chaoming0625
authored
chore: JAX>=0.10.0 compatibility (#824)
* chore: simplify imports by using brainstate._compatible_import for JAX compatibility * chore: update import statement for jaxpr_to_python_code to use brainstate.transform * chore: bump version to 2.7.8 * chore: update changelog for version 2.7.8 with JAX compatibility improvements and packaging updates * test: skip test1 if brainstate version is less than 0.3.1 * test: skip test1 if brainstate version is less than 0.3.1 due to method availability
1 parent f861c4a commit 1847bd3

File tree

9 files changed

+4470
-5574
lines changed

9 files changed

+4470
-5574
lines changed

‎brainpy/__init__.py‎

Lines changed: 146 additions & 146 deletions
Original file line numberDiff line numberDiff line change
@@ -1,146 +1,146 @@
1-
# -*- coding: utf-8 -*-
2-
# Copyright 2025 BrainX Ecosystem Limited. All Rights Reserved.
3-
#
4-
# Licensed under the Apache License, Version 2.0 (the "License");
5-
# you may not use this file except in compliance with the License.
6-
# You may obtain a copy of the License at
7-
#
8-
# http://www.apache.org/licenses/LICENSE-2.0
9-
#
10-
# Unless required by applicable law or agreed to in writing, software
11-
# distributed under the License is distributed on an "AS IS" BASIS,
12-
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13-
# See the License for the specific language governing permissions and
14-
# limitations under the License.
15-
# ==============================================================================
16-
17-
__version__ = "2.7.7"
18-
__version_info__ = tuple(map(int, __version__.split(".")))
19-
20-
from brainpy import _errors as errors
21-
# fundamental supporting modules
22-
from brainpy import check, tools
23-
# Part: Math Foundation #
24-
# ----------------------- #
25-
# math foundation
26-
from brainpy import math
27-
from brainpy import mixin
28-
# Part: Toolbox #
29-
# --------------- #
30-
# modules of toolbox
31-
from . import (
32-
connect, # synaptic connection
33-
initialize, # weight initialization
34-
optim, # gradient descent optimizers
35-
losses, # loss functions
36-
measure, # methods for data analysis
37-
inputs, # methods for generating input currents
38-
encoding, # encoding schema
39-
checkpoints, # checkpoints
40-
check, # error checking
41-
algorithms, # online or offline training algorithms
42-
)
43-
from .math import BrainPyObject
44-
45-
# convenient alias
46-
conn = connect
47-
init = initialize
48-
49-
# numerical integrators
50-
from brainpy import integrators
51-
from brainpy.integrators import ode, sde, fde
52-
from brainpy.integrators.base import (Integrator as Integrator)
53-
from brainpy.integrators.joint_eq import (JointEq as JointEq)
54-
from brainpy.integrators.runner import (IntegratorRunner as IntegratorRunner)
55-
from brainpy.integrators.ode.generic import (odeint as odeint)
56-
from brainpy.integrators.sde.generic import (sdeint as sdeint)
57-
from brainpy.integrators.fde.generic import (fdeint as fdeint)
58-
59-
# Part: Models #
60-
# -------------- #
61-
62-
# base classes
63-
from brainpy.dynsys import (
64-
DynamicalSystem as DynamicalSystem,
65-
DynSysGroup as DynSysGroup, # collectors
66-
Sequential as Sequential,
67-
Dynamic as Dynamic, # category
68-
Projection as Projection,
69-
receive_update_input, # decorators
70-
receive_update_output,
71-
not_receive_update_input,
72-
not_receive_update_output,
73-
)
74-
75-
DynamicalSystemNS = DynamicalSystem
76-
Network = DynSysGroup
77-
# delays
78-
from brainpy.delay import (
79-
VarDelay as VarDelay,
80-
)
81-
82-
# building blocks
83-
from brainpy import (
84-
dnn, layers, # module for dnn layers
85-
dyn, # module for modeling dynamics
86-
)
87-
88-
NeuGroup = NeuGroupNS = dyn.NeuDyn
89-
dyn.DynamicalSystem = DynamicalSystem
90-
91-
# common tools
92-
from brainpy.context import (share as share)
93-
from brainpy.helpers import (
94-
reset_level as reset_level,
95-
reset_state as reset_state,
96-
save_state as save_state,
97-
load_state as load_state,
98-
clear_input as clear_input
99-
)
100-
101-
# Part: Running #
102-
# --------------- #
103-
from brainpy.runners import (DSRunner as DSRunner)
104-
from brainpy.transform import (LoopOverTime as LoopOverTime, )
105-
from brainpy import (running as running)
106-
107-
# Part: Training #
108-
# ---------------- #
109-
from brainpy.train.base import (DSTrainer as DSTrainer, )
110-
from brainpy.train.back_propagation import (BPTT as BPTT,
111-
BPFF as BPFF, )
112-
from brainpy.train.online import (OnlineTrainer as OnlineTrainer,
113-
ForceTrainer as ForceTrainer, )
114-
from brainpy.train.offline import (OfflineTrainer as OfflineTrainer,
115-
RidgeTrainer as RidgeTrainer, )
116-
117-
# Part: Analysis #
118-
# ---------------- #
119-
from brainpy import (analysis as analysis)
120-
121-
# Part: Others #
122-
# ---------------- #
123-
import brainpy.visualization as visualize
124-
125-
# Part: Deprecations #
126-
# -------------------- #
127-
from brainpy import train
128-
from brainpy import (
129-
channels, # channel models
130-
neurons, # neuron groups
131-
synapses, # synapses
132-
rates, # rate models
133-
synouts, # synaptic output
134-
synplast, # synaptic plasticity
135-
)
136-
from brainpy.math.object_transform.base import Base as Base
137-
138-
from brainpy.math.object_transform.collectors import (
139-
ArrayCollector as ArrayCollector,
140-
Collector as Collector,
141-
)
142-
143-
optimizers = optim
144-
145-
# New package
146-
from brainpy import state
1+
# -*- coding: utf-8 -*-
2+
# Copyright 2025 BrainX Ecosystem Limited. All Rights Reserved.
3+
#
4+
# Licensed under the Apache License, Version 2.0 (the "License");
5+
# you may not use this file except in compliance with the License.
6+
# You may obtain a copy of the License at
7+
#
8+
# http://www.apache.org/licenses/LICENSE-2.0
9+
#
10+
# Unless required by applicable law or agreed to in writing, software
11+
# distributed under the License is distributed on an "AS IS" BASIS,
12+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
# See the License for the specific language governing permissions and
14+
# limitations under the License.
15+
# ==============================================================================
16+
17+
__version__ = "2.7.8"
18+
__version_info__ = tuple(map(int, __version__.split(".")))
19+
20+
from brainpy import _errors as errors
21+
# fundamental supporting modules
22+
from brainpy import check, tools
23+
# Part: Math Foundation #
24+
# ----------------------- #
25+
# math foundation
26+
from brainpy import math
27+
from brainpy import mixin
28+
# Part: Toolbox #
29+
# --------------- #
30+
# modules of toolbox
31+
from . import (
32+
connect, # synaptic connection
33+
initialize, # weight initialization
34+
optim, # gradient descent optimizers
35+
losses, # loss functions
36+
measure, # methods for data analysis
37+
inputs, # methods for generating input currents
38+
encoding, # encoding schema
39+
checkpoints, # checkpoints
40+
check, # error checking
41+
algorithms, # online or offline training algorithms
42+
)
43+
from .math import BrainPyObject
44+
45+
# convenient alias
46+
conn = connect
47+
init = initialize
48+
49+
# numerical integrators
50+
from brainpy import integrators
51+
from brainpy.integrators import ode, sde, fde
52+
from brainpy.integrators.base import (Integrator as Integrator)
53+
from brainpy.integrators.joint_eq import (JointEq as JointEq)
54+
from brainpy.integrators.runner import (IntegratorRunner as IntegratorRunner)
55+
from brainpy.integrators.ode.generic import (odeint as odeint)
56+
from brainpy.integrators.sde.generic import (sdeint as sdeint)
57+
from brainpy.integrators.fde.generic import (fdeint as fdeint)
58+
59+
# Part: Models #
60+
# -------------- #
61+
62+
# base classes
63+
from brainpy.dynsys import (
64+
DynamicalSystem as DynamicalSystem,
65+
DynSysGroup as DynSysGroup, # collectors
66+
Sequential as Sequential,
67+
Dynamic as Dynamic, # category
68+
Projection as Projection,
69+
receive_update_input, # decorators
70+
receive_update_output,
71+
not_receive_update_input,
72+
not_receive_update_output,
73+
)
74+
75+
DynamicalSystemNS = DynamicalSystem
76+
Network = DynSysGroup
77+
# delays
78+
from brainpy.delay import (
79+
VarDelay as VarDelay,
80+
)
81+
82+
# building blocks
83+
from brainpy import (
84+
dnn, layers, # module for dnn layers
85+
dyn, # module for modeling dynamics
86+
)
87+
88+
NeuGroup = NeuGroupNS = dyn.NeuDyn
89+
dyn.DynamicalSystem = DynamicalSystem
90+
91+
# common tools
92+
from brainpy.context import (share as share)
93+
from brainpy.helpers import (
94+
reset_level as reset_level,
95+
reset_state as reset_state,
96+
save_state as save_state,
97+
load_state as load_state,
98+
clear_input as clear_input
99+
)
100+
101+
# Part: Running #
102+
# --------------- #
103+
from brainpy.runners import (DSRunner as DSRunner)
104+
from brainpy.transform import (LoopOverTime as LoopOverTime, )
105+
from brainpy import (running as running)
106+
107+
# Part: Training #
108+
# ---------------- #
109+
from brainpy.train.base import (DSTrainer as DSTrainer, )
110+
from brainpy.train.back_propagation import (BPTT as BPTT,
111+
BPFF as BPFF, )
112+
from brainpy.train.online import (OnlineTrainer as OnlineTrainer,
113+
ForceTrainer as ForceTrainer, )
114+
from brainpy.train.offline import (OfflineTrainer as OfflineTrainer,
115+
RidgeTrainer as RidgeTrainer, )
116+
117+
# Part: Analysis #
118+
# ---------------- #
119+
from brainpy import (analysis as analysis)
120+
121+
# Part: Others #
122+
# ---------------- #
123+
import brainpy.visualization as visualize
124+
125+
# Part: Deprecations #
126+
# -------------------- #
127+
from brainpy import train
128+
from brainpy import (
129+
channels, # channel models
130+
neurons, # neuron groups
131+
synapses, # synapses
132+
rates, # rate models
133+
synouts, # synaptic output
134+
synplast, # synaptic plasticity
135+
)
136+
from brainpy.math.object_transform.base import Base as Base
137+
138+
from brainpy.math.object_transform.collectors import (
139+
ArrayCollector as ArrayCollector,
140+
Collector as Collector,
141+
)
142+
143+
optimizers = optim
144+
145+
# New package
146+
from brainpy import state

0 commit comments

Comments
 (0)