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.
Files changed (163) hide show
  1. brainstate/__init__.py +169 -58
  2. brainstate/_compatible_import.py +340 -148
  3. brainstate/_compatible_import_test.py +681 -0
  4. brainstate/_deprecation.py +210 -0
  5. brainstate/_deprecation_test.py +2319 -0
  6. brainstate/{util/error.py → _error.py} +45 -55
  7. brainstate/_state.py +1652 -1605
  8. brainstate/_state_test.py +52 -52
  9. brainstate/_utils.py +47 -47
  10. brainstate/environ.py +1495 -563
  11. brainstate/environ_test.py +1223 -62
  12. brainstate/graph/__init__.py +22 -29
  13. brainstate/graph/_node.py +240 -0
  14. brainstate/graph/_node_test.py +589 -0
  15. brainstate/graph/{_graph_operation.py → _operation.py} +1624 -1738
  16. brainstate/graph/_operation_test.py +1147 -0
  17. brainstate/mixin.py +1433 -365
  18. brainstate/mixin_test.py +1017 -77
  19. brainstate/nn/__init__.py +137 -135
  20. brainstate/nn/_activations.py +1100 -808
  21. brainstate/nn/_activations_test.py +354 -331
  22. brainstate/nn/_collective_ops.py +633 -514
  23. brainstate/nn/_collective_ops_test.py +774 -43
  24. brainstate/nn/_common.py +226 -178
  25. brainstate/nn/_common_test.py +154 -0
  26. brainstate/nn/_conv.py +2010 -501
  27. brainstate/nn/_conv_test.py +849 -238
  28. brainstate/nn/_delay.py +575 -588
  29. brainstate/nn/_delay_test.py +243 -238
  30. brainstate/nn/_dropout.py +618 -426
  31. brainstate/nn/_dropout_test.py +477 -100
  32. brainstate/nn/_dynamics.py +1267 -1343
  33. brainstate/nn/_dynamics_test.py +67 -78
  34. brainstate/nn/_elementwise.py +1298 -1119
  35. brainstate/nn/_elementwise_test.py +830 -169
  36. brainstate/nn/_embedding.py +408 -58
  37. brainstate/nn/_embedding_test.py +156 -0
  38. brainstate/nn/{_fixedprob.py → _event_fixedprob.py} +233 -239
  39. brainstate/nn/{_fixedprob_test.py → _event_fixedprob_test.py} +115 -114
  40. brainstate/nn/{_linear_mv.py → _event_linear.py} +83 -83
  41. brainstate/nn/{_linear_mv_test.py → _event_linear_test.py} +121 -120
  42. brainstate/nn/_exp_euler.py +254 -92
  43. brainstate/nn/_exp_euler_test.py +377 -35
  44. brainstate/nn/_linear.py +744 -424
  45. brainstate/nn/_linear_test.py +475 -107
  46. brainstate/nn/_metrics.py +1070 -0
  47. brainstate/nn/_metrics_test.py +611 -0
  48. brainstate/nn/_module.py +384 -377
  49. brainstate/nn/_module_test.py +40 -40
  50. brainstate/nn/_normalizations.py +1334 -975
  51. brainstate/nn/_normalizations_test.py +699 -73
  52. brainstate/nn/_paddings.py +1020 -0
  53. brainstate/nn/_paddings_test.py +723 -0
  54. brainstate/nn/_poolings.py +2239 -1177
  55. brainstate/nn/_poolings_test.py +953 -217
  56. brainstate/nn/{_rate_rnns.py → _rnns.py} +946 -554
  57. brainstate/nn/_rnns_test.py +593 -0
  58. brainstate/nn/_utils.py +216 -89
  59. brainstate/nn/_utils_test.py +402 -0
  60. brainstate/{init/_random_inits.py → nn/init.py} +809 -553
  61. brainstate/{init/_random_inits_test.py → nn/init_test.py} +180 -149
  62. brainstate/random/__init__.py +270 -24
  63. brainstate/random/_rand_funs.py +3938 -3616
  64. brainstate/random/_rand_funs_test.py +640 -567
  65. brainstate/random/_rand_seed.py +675 -210
  66. brainstate/random/_rand_seed_test.py +48 -48
  67. brainstate/random/_rand_state.py +1617 -1409
  68. brainstate/random/_rand_state_test.py +551 -0
  69. brainstate/transform/__init__.py +59 -0
  70. brainstate/transform/_ad_checkpoint.py +176 -0
  71. brainstate/{compile → transform}/_ad_checkpoint_test.py +49 -49
  72. brainstate/{augment → transform}/_autograd.py +1025 -778
  73. brainstate/{augment → transform}/_autograd_test.py +1289 -1289
  74. brainstate/transform/_conditions.py +316 -0
  75. brainstate/{compile → transform}/_conditions_test.py +220 -220
  76. brainstate/{compile → transform}/_error_if.py +94 -92
  77. brainstate/{compile → transform}/_error_if_test.py +52 -52
  78. brainstate/transform/_eval_shape.py +145 -0
  79. brainstate/{augment → transform}/_eval_shape_test.py +38 -38
  80. brainstate/{compile → transform}/_jit.py +399 -346
  81. brainstate/{compile → transform}/_jit_test.py +143 -143
  82. brainstate/{compile → transform}/_loop_collect_return.py +675 -536
  83. brainstate/{compile → transform}/_loop_collect_return_test.py +58 -58
  84. brainstate/{compile → transform}/_loop_no_collection.py +283 -184
  85. brainstate/{compile → transform}/_loop_no_collection_test.py +50 -50
  86. brainstate/transform/_make_jaxpr.py +2016 -0
  87. brainstate/transform/_make_jaxpr_test.py +1510 -0
  88. brainstate/transform/_mapping.py +529 -0
  89. brainstate/transform/_mapping_test.py +194 -0
  90. brainstate/{compile → transform}/_progress_bar.py +255 -202
  91. brainstate/{augment → transform}/_random.py +171 -151
  92. brainstate/{compile → transform}/_unvmap.py +256 -159
  93. brainstate/transform/_util.py +286 -0
  94. brainstate/typing.py +837 -304
  95. brainstate/typing_test.py +780 -0
  96. brainstate/util/__init__.py +27 -50
  97. brainstate/util/_others.py +1025 -0
  98. brainstate/util/_others_test.py +962 -0
  99. brainstate/util/_pretty_pytree.py +1301 -0
  100. brainstate/util/_pretty_pytree_test.py +675 -0
  101. brainstate/util/{pretty_repr.py → _pretty_repr.py} +462 -328
  102. brainstate/util/_pretty_repr_test.py +696 -0
  103. brainstate/util/filter.py +945 -469
  104. brainstate/util/filter_test.py +912 -0
  105. brainstate/util/struct.py +910 -523
  106. brainstate/util/struct_test.py +602 -0
  107. {brainstate-0.1.10.dist-info → brainstate-0.2.1.dist-info}/METADATA +108 -91
  108. brainstate-0.2.1.dist-info/RECORD +111 -0
  109. {brainstate-0.1.10.dist-info → brainstate-0.2.1.dist-info}/licenses/LICENSE +202 -202
  110. brainstate/augment/__init__.py +0 -30
  111. brainstate/augment/_eval_shape.py +0 -99
  112. brainstate/augment/_mapping.py +0 -1060
  113. brainstate/augment/_mapping_test.py +0 -597
  114. brainstate/compile/__init__.py +0 -38
  115. brainstate/compile/_ad_checkpoint.py +0 -204
  116. brainstate/compile/_conditions.py +0 -256
  117. brainstate/compile/_make_jaxpr.py +0 -888
  118. brainstate/compile/_make_jaxpr_test.py +0 -156
  119. brainstate/compile/_util.py +0 -147
  120. brainstate/functional/__init__.py +0 -27
  121. brainstate/graph/_graph_node.py +0 -244
  122. brainstate/graph/_graph_node_test.py +0 -73
  123. brainstate/graph/_graph_operation_test.py +0 -563
  124. brainstate/init/__init__.py +0 -26
  125. brainstate/init/_base.py +0 -52
  126. brainstate/init/_generic.py +0 -244
  127. brainstate/init/_regular_inits.py +0 -105
  128. brainstate/init/_regular_inits_test.py +0 -50
  129. brainstate/nn/_inputs.py +0 -608
  130. brainstate/nn/_ltp.py +0 -28
  131. brainstate/nn/_neuron.py +0 -705
  132. brainstate/nn/_neuron_test.py +0 -161
  133. brainstate/nn/_others.py +0 -46
  134. brainstate/nn/_projection.py +0 -486
  135. brainstate/nn/_rate_rnns_test.py +0 -63
  136. brainstate/nn/_readout.py +0 -209
  137. brainstate/nn/_readout_test.py +0 -53
  138. brainstate/nn/_stp.py +0 -236
  139. brainstate/nn/_synapse.py +0 -505
  140. brainstate/nn/_synapse_test.py +0 -131
  141. brainstate/nn/_synaptic_projection.py +0 -423
  142. brainstate/nn/_synouts.py +0 -162
  143. brainstate/nn/_synouts_test.py +0 -57
  144. brainstate/nn/metrics.py +0 -388
  145. brainstate/optim/__init__.py +0 -38
  146. brainstate/optim/_base.py +0 -64
  147. brainstate/optim/_lr_scheduler.py +0 -448
  148. brainstate/optim/_lr_scheduler_test.py +0 -50
  149. brainstate/optim/_optax_optimizer.py +0 -152
  150. brainstate/optim/_optax_optimizer_test.py +0 -53
  151. brainstate/optim/_sgd_optimizer.py +0 -1104
  152. brainstate/random/_random_for_unit.py +0 -52
  153. brainstate/surrogate.py +0 -1957
  154. brainstate/transform.py +0 -23
  155. brainstate/util/caller.py +0 -98
  156. brainstate/util/others.py +0 -540
  157. brainstate/util/pretty_pytree.py +0 -945
  158. brainstate/util/pretty_pytree_test.py +0 -159
  159. brainstate/util/pretty_table.py +0 -2954
  160. brainstate/util/scaling.py +0 -258
  161. brainstate-0.1.10.dist-info/RECORD +0 -130
  162. {brainstate-0.1.10.dist-info → brainstate-0.2.1.dist-info}/WHEEL +0 -0
  163. {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 BDP 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
-
17
- from . import metrics
18
- from ._collective_ops import *
19
- from ._collective_ops import __all__ as collective_ops_all
20
- from ._common import *
21
- from ._common import __all__ as common_all
22
- from ._conv import *
23
- from ._conv import __all__ as conv_all
24
- from ._delay import *
25
- from ._delay import __all__ as state_delay_all
26
- from ._dropout import *
27
- from ._dropout import __all__ as dropout_all
28
- from ._dynamics import *
29
- from ._dynamics import __all__ as dyn_all
30
- from ._elementwise import *
31
- from ._elementwise import __all__ as elementwise_all
32
- from ._embedding import *
33
- from ._embedding import __all__ as embed_all
34
- from ._exp_euler import *
35
- from ._exp_euler import __all__ as exp_euler_all
36
- from ._fixedprob import *
37
- from ._fixedprob import __all__ as fixedprob_all
38
- from ._inputs import *
39
- from ._inputs import __all__ as inputs_all
40
- from ._linear import *
41
- from ._linear import __all__ as linear_all
42
- from ._linear_mv import *
43
- from ._linear_mv import __all__ as linear_mv_all
44
- from ._ltp import *
45
- from ._ltp import __all__ as ltp_all
46
- from ._module import *
47
- from ._module import __all__ as module_all
48
- from ._neuron import *
49
- from ._neuron import __all__ as dyn_neuron_all
50
- from ._normalizations import *
51
- from ._normalizations import __all__ as normalizations_all
52
- from ._others import *
53
- from ._others import __all__ as _others_all
54
- from ._poolings import *
55
- from ._poolings import __all__ as poolings_all
56
- from ._projection import *
57
- from ._projection import __all__ as projection_all
58
- from ._rate_rnns import *
59
- from ._rate_rnns import __all__ as rate_rnns
60
- from ._readout import *
61
- from ._readout import __all__ as readout_all
62
- from ._stp import *
63
- from ._stp import __all__ as stp_all
64
- from ._synapse import *
65
- from ._synapse import __all__ as dyn_synapse_all
66
- from ._synaptic_projection import *
67
- from ._synaptic_projection import __all__ as _syn_proj_all
68
- from ._synouts import *
69
- from ._synouts import __all__ as synouts_all
70
- from ._utils import *
71
- from ._utils import __all__ as utils_all
72
-
73
- __all__ = (
74
- [
75
- 'metrics',
76
- ]
77
- + collective_ops_all
78
- + common_all
79
- + elementwise_all
80
- + module_all
81
- + exp_euler_all
82
- + utils_all
83
- + dyn_all
84
- + projection_all
85
- + state_delay_all
86
- + synouts_all
87
- + conv_all
88
- + linear_all
89
- + normalizations_all
90
- + poolings_all
91
- + fixedprob_all
92
- + linear_mv_all
93
- + embed_all
94
- + dropout_all
95
- + elementwise_all
96
- + dyn_neuron_all
97
- + dyn_synapse_all
98
- + inputs_all
99
- + rate_rnns
100
- + readout_all
101
- + stp_all
102
- + ltp_all
103
- + _syn_proj_all
104
- + _others_all
105
- )
106
-
107
- del (
108
- collective_ops_all,
109
- common_all,
110
- module_all,
111
- exp_euler_all,
112
- utils_all,
113
- dyn_all,
114
- projection_all,
115
- state_delay_all,
116
- synouts_all,
117
- conv_all,
118
- linear_all,
119
- normalizations_all,
120
- poolings_all,
121
- embed_all,
122
- fixedprob_all,
123
- linear_mv_all,
124
- dropout_all,
125
- elementwise_all,
126
- dyn_neuron_all,
127
- dyn_synapse_all,
128
- inputs_all,
129
- readout_all,
130
- rate_rnns,
131
- stp_all,
132
- ltp_all,
133
- _syn_proj_all,
134
- _others_all,
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}'")