brainstate 0.1.10__py2.py3-none-any.whl → 0.2.1__py2.py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- brainstate/__init__.py +169 -58
- brainstate/_compatible_import.py +340 -148
- brainstate/_compatible_import_test.py +681 -0
- brainstate/_deprecation.py +210 -0
- brainstate/_deprecation_test.py +2319 -0
- brainstate/{util/error.py → _error.py} +45 -55
- brainstate/_state.py +1652 -1605
- brainstate/_state_test.py +52 -52
- brainstate/_utils.py +47 -47
- brainstate/environ.py +1495 -563
- brainstate/environ_test.py +1223 -62
- brainstate/graph/__init__.py +22 -29
- brainstate/graph/_node.py +240 -0
- brainstate/graph/_node_test.py +589 -0
- brainstate/graph/{_graph_operation.py → _operation.py} +1624 -1738
- brainstate/graph/_operation_test.py +1147 -0
- brainstate/mixin.py +1433 -365
- brainstate/mixin_test.py +1017 -77
- brainstate/nn/__init__.py +137 -135
- brainstate/nn/_activations.py +1100 -808
- brainstate/nn/_activations_test.py +354 -331
- brainstate/nn/_collective_ops.py +633 -514
- brainstate/nn/_collective_ops_test.py +774 -43
- brainstate/nn/_common.py +226 -178
- brainstate/nn/_common_test.py +154 -0
- brainstate/nn/_conv.py +2010 -501
- brainstate/nn/_conv_test.py +849 -238
- brainstate/nn/_delay.py +575 -588
- brainstate/nn/_delay_test.py +243 -238
- brainstate/nn/_dropout.py +618 -426
- brainstate/nn/_dropout_test.py +477 -100
- brainstate/nn/_dynamics.py +1267 -1343
- brainstate/nn/_dynamics_test.py +67 -78
- brainstate/nn/_elementwise.py +1298 -1119
- brainstate/nn/_elementwise_test.py +830 -169
- brainstate/nn/_embedding.py +408 -58
- brainstate/nn/_embedding_test.py +156 -0
- brainstate/nn/{_fixedprob.py → _event_fixedprob.py} +233 -239
- brainstate/nn/{_fixedprob_test.py → _event_fixedprob_test.py} +115 -114
- brainstate/nn/{_linear_mv.py → _event_linear.py} +83 -83
- brainstate/nn/{_linear_mv_test.py → _event_linear_test.py} +121 -120
- brainstate/nn/_exp_euler.py +254 -92
- brainstate/nn/_exp_euler_test.py +377 -35
- brainstate/nn/_linear.py +744 -424
- brainstate/nn/_linear_test.py +475 -107
- brainstate/nn/_metrics.py +1070 -0
- brainstate/nn/_metrics_test.py +611 -0
- brainstate/nn/_module.py +384 -377
- brainstate/nn/_module_test.py +40 -40
- brainstate/nn/_normalizations.py +1334 -975
- brainstate/nn/_normalizations_test.py +699 -73
- brainstate/nn/_paddings.py +1020 -0
- brainstate/nn/_paddings_test.py +723 -0
- brainstate/nn/_poolings.py +2239 -1177
- brainstate/nn/_poolings_test.py +953 -217
- brainstate/nn/{_rate_rnns.py → _rnns.py} +946 -554
- brainstate/nn/_rnns_test.py +593 -0
- brainstate/nn/_utils.py +216 -89
- brainstate/nn/_utils_test.py +402 -0
- brainstate/{init/_random_inits.py → nn/init.py} +809 -553
- brainstate/{init/_random_inits_test.py → nn/init_test.py} +180 -149
- brainstate/random/__init__.py +270 -24
- brainstate/random/_rand_funs.py +3938 -3616
- brainstate/random/_rand_funs_test.py +640 -567
- brainstate/random/_rand_seed.py +675 -210
- brainstate/random/_rand_seed_test.py +48 -48
- brainstate/random/_rand_state.py +1617 -1409
- brainstate/random/_rand_state_test.py +551 -0
- brainstate/transform/__init__.py +59 -0
- brainstate/transform/_ad_checkpoint.py +176 -0
- brainstate/{compile → transform}/_ad_checkpoint_test.py +49 -49
- brainstate/{augment → transform}/_autograd.py +1025 -778
- brainstate/{augment → transform}/_autograd_test.py +1289 -1289
- brainstate/transform/_conditions.py +316 -0
- brainstate/{compile → transform}/_conditions_test.py +220 -220
- brainstate/{compile → transform}/_error_if.py +94 -92
- brainstate/{compile → transform}/_error_if_test.py +52 -52
- brainstate/transform/_eval_shape.py +145 -0
- brainstate/{augment → transform}/_eval_shape_test.py +38 -38
- brainstate/{compile → transform}/_jit.py +399 -346
- brainstate/{compile → transform}/_jit_test.py +143 -143
- brainstate/{compile → transform}/_loop_collect_return.py +675 -536
- brainstate/{compile → transform}/_loop_collect_return_test.py +58 -58
- brainstate/{compile → transform}/_loop_no_collection.py +283 -184
- brainstate/{compile → transform}/_loop_no_collection_test.py +50 -50
- brainstate/transform/_make_jaxpr.py +2016 -0
- brainstate/transform/_make_jaxpr_test.py +1510 -0
- brainstate/transform/_mapping.py +529 -0
- brainstate/transform/_mapping_test.py +194 -0
- brainstate/{compile → transform}/_progress_bar.py +255 -202
- brainstate/{augment → transform}/_random.py +171 -151
- brainstate/{compile → transform}/_unvmap.py +256 -159
- brainstate/transform/_util.py +286 -0
- brainstate/typing.py +837 -304
- brainstate/typing_test.py +780 -0
- brainstate/util/__init__.py +27 -50
- brainstate/util/_others.py +1025 -0
- brainstate/util/_others_test.py +962 -0
- brainstate/util/_pretty_pytree.py +1301 -0
- brainstate/util/_pretty_pytree_test.py +675 -0
- brainstate/util/{pretty_repr.py → _pretty_repr.py} +462 -328
- brainstate/util/_pretty_repr_test.py +696 -0
- brainstate/util/filter.py +945 -469
- brainstate/util/filter_test.py +912 -0
- brainstate/util/struct.py +910 -523
- brainstate/util/struct_test.py +602 -0
- {brainstate-0.1.10.dist-info → brainstate-0.2.1.dist-info}/METADATA +108 -91
- brainstate-0.2.1.dist-info/RECORD +111 -0
- {brainstate-0.1.10.dist-info → brainstate-0.2.1.dist-info}/licenses/LICENSE +202 -202
- brainstate/augment/__init__.py +0 -30
- brainstate/augment/_eval_shape.py +0 -99
- brainstate/augment/_mapping.py +0 -1060
- brainstate/augment/_mapping_test.py +0 -597
- brainstate/compile/__init__.py +0 -38
- brainstate/compile/_ad_checkpoint.py +0 -204
- brainstate/compile/_conditions.py +0 -256
- brainstate/compile/_make_jaxpr.py +0 -888
- brainstate/compile/_make_jaxpr_test.py +0 -156
- brainstate/compile/_util.py +0 -147
- brainstate/functional/__init__.py +0 -27
- brainstate/graph/_graph_node.py +0 -244
- brainstate/graph/_graph_node_test.py +0 -73
- brainstate/graph/_graph_operation_test.py +0 -563
- brainstate/init/__init__.py +0 -26
- brainstate/init/_base.py +0 -52
- brainstate/init/_generic.py +0 -244
- brainstate/init/_regular_inits.py +0 -105
- brainstate/init/_regular_inits_test.py +0 -50
- brainstate/nn/_inputs.py +0 -608
- brainstate/nn/_ltp.py +0 -28
- brainstate/nn/_neuron.py +0 -705
- brainstate/nn/_neuron_test.py +0 -161
- brainstate/nn/_others.py +0 -46
- brainstate/nn/_projection.py +0 -486
- brainstate/nn/_rate_rnns_test.py +0 -63
- brainstate/nn/_readout.py +0 -209
- brainstate/nn/_readout_test.py +0 -53
- brainstate/nn/_stp.py +0 -236
- brainstate/nn/_synapse.py +0 -505
- brainstate/nn/_synapse_test.py +0 -131
- brainstate/nn/_synaptic_projection.py +0 -423
- brainstate/nn/_synouts.py +0 -162
- brainstate/nn/_synouts_test.py +0 -57
- brainstate/nn/metrics.py +0 -388
- brainstate/optim/__init__.py +0 -38
- brainstate/optim/_base.py +0 -64
- brainstate/optim/_lr_scheduler.py +0 -448
- brainstate/optim/_lr_scheduler_test.py +0 -50
- brainstate/optim/_optax_optimizer.py +0 -152
- brainstate/optim/_optax_optimizer_test.py +0 -53
- brainstate/optim/_sgd_optimizer.py +0 -1104
- brainstate/random/_random_for_unit.py +0 -52
- brainstate/surrogate.py +0 -1957
- brainstate/transform.py +0 -23
- brainstate/util/caller.py +0 -98
- brainstate/util/others.py +0 -540
- brainstate/util/pretty_pytree.py +0 -945
- brainstate/util/pretty_pytree_test.py +0 -159
- brainstate/util/pretty_table.py +0 -2954
- brainstate/util/scaling.py +0 -258
- brainstate-0.1.10.dist-info/RECORD +0 -130
- {brainstate-0.1.10.dist-info → brainstate-0.2.1.dist-info}/WHEEL +0 -0
- {brainstate-0.1.10.dist-info → brainstate-0.2.1.dist-info}/top_level.txt +0 -0
brainstate/nn/__init__.py
CHANGED
@@ -1,135 +1,137 @@
|
|
1
|
-
# Copyright 2024
|
2
|
-
#
|
3
|
-
# Licensed under the Apache License, Version 2.0 (the "License");
|
4
|
-
# you may not use this file except in compliance with the License.
|
5
|
-
# You may obtain a copy of the License at
|
6
|
-
#
|
7
|
-
# http://www.apache.org/licenses/LICENSE-2.0
|
8
|
-
#
|
9
|
-
# Unless required by applicable law or agreed to in writing, software
|
10
|
-
# distributed under the License is distributed on an "AS IS" BASIS,
|
11
|
-
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
12
|
-
# See the License for the specific language governing permissions and
|
13
|
-
# limitations under the License.
|
14
|
-
# ==============================================================================
|
15
|
-
|
16
|
-
|
17
|
-
from . import
|
18
|
-
from .
|
19
|
-
from ._collective_ops import
|
20
|
-
from .
|
21
|
-
from ._common import
|
22
|
-
from .
|
23
|
-
from ._conv import
|
24
|
-
from .
|
25
|
-
from ._delay import
|
26
|
-
from .
|
27
|
-
from ._dropout import
|
28
|
-
from .
|
29
|
-
from ._dynamics import
|
30
|
-
from .
|
31
|
-
from ._elementwise import
|
32
|
-
from .
|
33
|
-
from ._embedding import
|
34
|
-
from .
|
35
|
-
from .
|
36
|
-
from .
|
37
|
-
from .
|
38
|
-
from .
|
39
|
-
from .
|
40
|
-
from .
|
41
|
-
from ._linear import
|
42
|
-
from .
|
43
|
-
from .
|
44
|
-
from .
|
45
|
-
from .
|
46
|
-
from ._module import
|
47
|
-
from .
|
48
|
-
from .
|
49
|
-
from .
|
50
|
-
from .
|
51
|
-
from .
|
52
|
-
from .
|
53
|
-
from .
|
54
|
-
from .
|
55
|
-
from .
|
56
|
-
from .
|
57
|
-
|
58
|
-
|
59
|
-
|
60
|
-
|
61
|
-
|
62
|
-
|
63
|
-
|
64
|
-
|
65
|
-
|
66
|
-
|
67
|
-
|
68
|
-
|
69
|
-
|
70
|
-
|
71
|
-
|
72
|
-
|
73
|
-
|
74
|
-
|
75
|
-
|
76
|
-
|
77
|
-
|
78
|
-
|
79
|
-
|
80
|
-
|
81
|
-
|
82
|
-
|
83
|
-
|
84
|
-
|
85
|
-
|
86
|
-
|
87
|
-
|
88
|
-
|
89
|
-
|
90
|
-
|
91
|
-
|
92
|
-
|
93
|
-
|
94
|
-
|
95
|
-
|
96
|
-
|
97
|
-
|
98
|
-
|
99
|
-
|
100
|
-
|
101
|
-
|
102
|
-
|
103
|
-
|
104
|
-
|
105
|
-
|
106
|
-
|
107
|
-
|
108
|
-
|
109
|
-
|
110
|
-
|
111
|
-
|
112
|
-
|
113
|
-
|
114
|
-
|
115
|
-
|
116
|
-
|
117
|
-
|
118
|
-
|
119
|
-
|
120
|
-
|
121
|
-
|
122
|
-
|
123
|
-
|
124
|
-
|
125
|
-
|
126
|
-
|
127
|
-
|
128
|
-
|
129
|
-
|
130
|
-
|
131
|
-
|
132
|
-
|
133
|
-
|
134
|
-
|
135
|
-
|
1
|
+
# Copyright 2024 BrainX Ecosystem Limited. All Rights Reserved.
|
2
|
+
#
|
3
|
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
4
|
+
# you may not use this file except in compliance with the License.
|
5
|
+
# You may obtain a copy of the License at
|
6
|
+
#
|
7
|
+
# http://www.apache.org/licenses/LICENSE-2.0
|
8
|
+
#
|
9
|
+
# Unless required by applicable law or agreed to in writing, software
|
10
|
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
11
|
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
12
|
+
# See the License for the specific language governing permissions and
|
13
|
+
# limitations under the License.
|
14
|
+
# ==============================================================================
|
15
|
+
|
16
|
+
from . import init
|
17
|
+
from ._activations import *
|
18
|
+
from ._activations import __all__ as activation_all
|
19
|
+
from ._collective_ops import *
|
20
|
+
from ._collective_ops import __all__ as collective_ops_all
|
21
|
+
from ._common import *
|
22
|
+
from ._common import __all__ as common_all
|
23
|
+
from ._conv import *
|
24
|
+
from ._conv import __all__ as conv_all
|
25
|
+
from ._delay import *
|
26
|
+
from ._delay import __all__ as state_delay_all
|
27
|
+
from ._dropout import *
|
28
|
+
from ._dropout import __all__ as dropout_all
|
29
|
+
from ._dynamics import *
|
30
|
+
from ._dynamics import __all__ as dyn_all
|
31
|
+
from ._elementwise import *
|
32
|
+
from ._elementwise import __all__ as elementwise_all
|
33
|
+
from ._embedding import *
|
34
|
+
from ._embedding import __all__ as embed_all
|
35
|
+
from ._event_fixedprob import *
|
36
|
+
from ._event_fixedprob import __all__ as fixedprob_all
|
37
|
+
from ._event_linear import *
|
38
|
+
from ._event_linear import __all__ as linear_mv_all
|
39
|
+
from ._exp_euler import *
|
40
|
+
from ._exp_euler import __all__ as exp_euler_all
|
41
|
+
from ._linear import *
|
42
|
+
from ._linear import __all__ as linear_all
|
43
|
+
from ._metrics import *
|
44
|
+
from ._metrics import __all__ as metrics_all
|
45
|
+
from ._module import *
|
46
|
+
from ._module import __all__ as module_all
|
47
|
+
from ._normalizations import *
|
48
|
+
from ._normalizations import __all__ as normalizations_all
|
49
|
+
from ._paddings import *
|
50
|
+
from ._paddings import __all__ as paddings_all
|
51
|
+
from ._poolings import *
|
52
|
+
from ._poolings import __all__ as poolings_all
|
53
|
+
from ._rnns import *
|
54
|
+
from ._rnns import __all__ as rate_rnns
|
55
|
+
from ._utils import *
|
56
|
+
from ._utils import __all__ as utils_all
|
57
|
+
|
58
|
+
__all__ = ['init'] + activation_all + metrics_all
|
59
|
+
__all__ = __all__ + collective_ops_all + common_all + elementwise_all + module_all + exp_euler_all
|
60
|
+
__all__ = __all__ + utils_all + dyn_all + state_delay_all + conv_all
|
61
|
+
__all__ = __all__ + linear_all + normalizations_all + paddings_all + poolings_all + fixedprob_all + linear_mv_all
|
62
|
+
__all__ = __all__ + embed_all + dropout_all + elementwise_all
|
63
|
+
__all__ = __all__ + rate_rnns
|
64
|
+
|
65
|
+
del (
|
66
|
+
metrics_all,
|
67
|
+
activation_all,
|
68
|
+
collective_ops_all,
|
69
|
+
common_all,
|
70
|
+
module_all,
|
71
|
+
exp_euler_all,
|
72
|
+
utils_all,
|
73
|
+
dyn_all,
|
74
|
+
state_delay_all,
|
75
|
+
conv_all,
|
76
|
+
linear_all,
|
77
|
+
normalizations_all,
|
78
|
+
paddings_all,
|
79
|
+
poolings_all,
|
80
|
+
embed_all,
|
81
|
+
fixedprob_all,
|
82
|
+
linear_mv_all,
|
83
|
+
dropout_all,
|
84
|
+
elementwise_all,
|
85
|
+
rate_rnns,
|
86
|
+
)
|
87
|
+
|
88
|
+
# Deprecated names that redirect to brainpy
|
89
|
+
_DEPRECATED_NAMES = {
|
90
|
+
'SpikeTime': 'brainpy.SpikeTime',
|
91
|
+
'PoissonSpike': 'brainpy.PoissonSpike',
|
92
|
+
'PoissonEncoder': 'brainpy.PoissonEncoder',
|
93
|
+
'PoissonInput': 'brainpy.PoissonInput',
|
94
|
+
'poisson_input': 'brainpy.poisson_input',
|
95
|
+
'Neuron': 'brainpy.Neuron',
|
96
|
+
'IF': 'brainpy.IF',
|
97
|
+
'LIF': 'brainpy.LIF',
|
98
|
+
'LIFRef': 'brainpy.LIFRef',
|
99
|
+
'ALIF': 'brainpy.ALIF',
|
100
|
+
'LeakyRateReadout': 'brainpy.LeakyRateReadout',
|
101
|
+
'LeakySpikeReadout': 'brainpy.LeakySpikeReadout',
|
102
|
+
'STP': 'brainpy.STP',
|
103
|
+
'STD': 'brainpy.STD',
|
104
|
+
'Synapse': 'brainpy.Synapse',
|
105
|
+
'Expon': 'brainpy.Expon',
|
106
|
+
'DualExpon': 'brainpy.DualExpon',
|
107
|
+
'Alpha': 'brainpy.Alpha',
|
108
|
+
'AMPA': 'brainpy.AMPA',
|
109
|
+
'GABAa': 'brainpy.GABAa',
|
110
|
+
'COBA': 'brainpy.COBA',
|
111
|
+
'CUBA': 'brainpy.CUBA',
|
112
|
+
'MgBlock': 'brainpy.MgBlock',
|
113
|
+
'SynOut': 'brainpy.SynOut',
|
114
|
+
'AlignPostProj': 'brainpy.AlignPostProj',
|
115
|
+
'DeltaProj': 'brainpy.DeltaProj',
|
116
|
+
'CurrentProj': 'brainpy.CurrentProj',
|
117
|
+
'align_pre_projection': 'brainpy.align_pre_projection',
|
118
|
+
'Projection': 'brainpy.Projection',
|
119
|
+
'SymmetryGapJunction': 'brainpy.SymmetryGapJunction',
|
120
|
+
'AsymmetryGapJunction': 'brainpy.AsymmetryGapJunction',
|
121
|
+
}
|
122
|
+
|
123
|
+
|
124
|
+
def __getattr__(name: str):
|
125
|
+
if name in _DEPRECATED_NAMES:
|
126
|
+
import warnings
|
127
|
+
new_name = _DEPRECATED_NAMES[name]
|
128
|
+
warnings.warn(
|
129
|
+
f"'brainstate.nn.{name}' is deprecated and will be removed in a future version. "
|
130
|
+
f"Please use '{new_name}' instead.",
|
131
|
+
DeprecationWarning,
|
132
|
+
stacklevel=2
|
133
|
+
)
|
134
|
+
# Import and return the actual brainpy object
|
135
|
+
import brainpy
|
136
|
+
return getattr(brainpy, name)
|
137
|
+
raise AttributeError(f"module 'brainstate.nn' has no attribute '{name}'")
|