brainstate 0.2.0__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 (112) hide show
  1. brainstate/__init__.py +169 -169
  2. brainstate/_compatible_import.py +340 -340
  3. brainstate/_compatible_import_test.py +681 -681
  4. brainstate/_deprecation.py +210 -210
  5. brainstate/_deprecation_test.py +2319 -2319
  6. brainstate/_error.py +45 -45
  7. brainstate/_state.py +1652 -1652
  8. brainstate/_state_test.py +52 -52
  9. brainstate/_utils.py +47 -47
  10. brainstate/environ.py +1495 -1495
  11. brainstate/environ_test.py +1223 -1223
  12. brainstate/graph/__init__.py +22 -22
  13. brainstate/graph/_node.py +240 -240
  14. brainstate/graph/_node_test.py +589 -589
  15. brainstate/graph/_operation.py +1624 -1624
  16. brainstate/graph/_operation_test.py +1147 -1147
  17. brainstate/mixin.py +1433 -1433
  18. brainstate/mixin_test.py +1017 -1017
  19. brainstate/nn/__init__.py +137 -137
  20. brainstate/nn/_activations.py +1100 -1100
  21. brainstate/nn/_activations_test.py +354 -354
  22. brainstate/nn/_collective_ops.py +633 -633
  23. brainstate/nn/_collective_ops_test.py +774 -774
  24. brainstate/nn/_common.py +226 -226
  25. brainstate/nn/_common_test.py +154 -154
  26. brainstate/nn/_conv.py +2010 -2010
  27. brainstate/nn/_conv_test.py +849 -849
  28. brainstate/nn/_delay.py +575 -575
  29. brainstate/nn/_delay_test.py +243 -243
  30. brainstate/nn/_dropout.py +618 -618
  31. brainstate/nn/_dropout_test.py +477 -477
  32. brainstate/nn/_dynamics.py +1267 -1267
  33. brainstate/nn/_dynamics_test.py +67 -67
  34. brainstate/nn/_elementwise.py +1298 -1298
  35. brainstate/nn/_elementwise_test.py +829 -829
  36. brainstate/nn/_embedding.py +408 -408
  37. brainstate/nn/_embedding_test.py +156 -156
  38. brainstate/nn/_event_fixedprob.py +233 -233
  39. brainstate/nn/_event_fixedprob_test.py +115 -115
  40. brainstate/nn/_event_linear.py +83 -83
  41. brainstate/nn/_event_linear_test.py +121 -121
  42. brainstate/nn/_exp_euler.py +254 -254
  43. brainstate/nn/_exp_euler_test.py +377 -377
  44. brainstate/nn/_linear.py +744 -744
  45. brainstate/nn/_linear_test.py +475 -475
  46. brainstate/nn/_metrics.py +1070 -1070
  47. brainstate/nn/_metrics_test.py +611 -611
  48. brainstate/nn/_module.py +384 -384
  49. brainstate/nn/_module_test.py +40 -40
  50. brainstate/nn/_normalizations.py +1334 -1334
  51. brainstate/nn/_normalizations_test.py +699 -699
  52. brainstate/nn/_paddings.py +1020 -1020
  53. brainstate/nn/_paddings_test.py +722 -722
  54. brainstate/nn/_poolings.py +2239 -2239
  55. brainstate/nn/_poolings_test.py +952 -952
  56. brainstate/nn/_rnns.py +946 -946
  57. brainstate/nn/_rnns_test.py +592 -592
  58. brainstate/nn/_utils.py +216 -216
  59. brainstate/nn/_utils_test.py +401 -401
  60. brainstate/nn/init.py +809 -809
  61. brainstate/nn/init_test.py +180 -180
  62. brainstate/random/__init__.py +270 -270
  63. brainstate/random/_rand_funs.py +3938 -3938
  64. brainstate/random/_rand_funs_test.py +640 -640
  65. brainstate/random/_rand_seed.py +675 -675
  66. brainstate/random/_rand_seed_test.py +48 -48
  67. brainstate/random/_rand_state.py +1617 -1617
  68. brainstate/random/_rand_state_test.py +551 -551
  69. brainstate/transform/__init__.py +59 -59
  70. brainstate/transform/_ad_checkpoint.py +176 -176
  71. brainstate/transform/_ad_checkpoint_test.py +49 -49
  72. brainstate/transform/_autograd.py +1025 -1025
  73. brainstate/transform/_autograd_test.py +1289 -1289
  74. brainstate/transform/_conditions.py +316 -316
  75. brainstate/transform/_conditions_test.py +220 -220
  76. brainstate/transform/_error_if.py +94 -94
  77. brainstate/transform/_error_if_test.py +52 -52
  78. brainstate/transform/_eval_shape.py +145 -145
  79. brainstate/transform/_eval_shape_test.py +38 -38
  80. brainstate/transform/_jit.py +399 -399
  81. brainstate/transform/_jit_test.py +143 -143
  82. brainstate/transform/_loop_collect_return.py +675 -675
  83. brainstate/transform/_loop_collect_return_test.py +58 -58
  84. brainstate/transform/_loop_no_collection.py +283 -283
  85. brainstate/transform/_loop_no_collection_test.py +50 -50
  86. brainstate/transform/_make_jaxpr.py +2016 -2016
  87. brainstate/transform/_make_jaxpr_test.py +1510 -1510
  88. brainstate/transform/_mapping.py +529 -529
  89. brainstate/transform/_mapping_test.py +194 -194
  90. brainstate/transform/_progress_bar.py +255 -255
  91. brainstate/transform/_random.py +171 -171
  92. brainstate/transform/_unvmap.py +256 -256
  93. brainstate/transform/_util.py +286 -286
  94. brainstate/typing.py +837 -837
  95. brainstate/typing_test.py +780 -780
  96. brainstate/util/__init__.py +27 -27
  97. brainstate/util/_others.py +1024 -1024
  98. brainstate/util/_others_test.py +962 -962
  99. brainstate/util/_pretty_pytree.py +1301 -1301
  100. brainstate/util/_pretty_pytree_test.py +675 -675
  101. brainstate/util/_pretty_repr.py +462 -462
  102. brainstate/util/_pretty_repr_test.py +696 -696
  103. brainstate/util/filter.py +945 -945
  104. brainstate/util/filter_test.py +911 -911
  105. brainstate/util/struct.py +910 -910
  106. brainstate/util/struct_test.py +602 -602
  107. {brainstate-0.2.0.dist-info → brainstate-0.2.1.dist-info}/METADATA +108 -108
  108. brainstate-0.2.1.dist-info/RECORD +111 -0
  109. {brainstate-0.2.0.dist-info → brainstate-0.2.1.dist-info}/licenses/LICENSE +202 -202
  110. brainstate-0.2.0.dist-info/RECORD +0 -111
  111. {brainstate-0.2.0.dist-info → brainstate-0.2.1.dist-info}/WHEEL +0 -0
  112. {brainstate-0.2.0.dist-info → brainstate-0.2.1.dist-info}/top_level.txt +0 -0
@@ -1,256 +1,256 @@
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
- import jax
17
- import jax.core
18
- import jax.interpreters.batching as batching
19
- import jax.interpreters.mlir as mlir
20
- import jax.numpy as jnp
21
-
22
- from brainstate._compatible_import import Primitive
23
- from brainstate._utils import set_module_as
24
-
25
- __all__ = [
26
- "unvmap",
27
- ]
28
-
29
-
30
- @set_module_as('brainstate.transform')
31
- def unvmap(x, op: str = 'any'):
32
- """
33
- Remove a leading vmap dimension by aggregating batched values.
34
-
35
- Parameters
36
- ----------
37
- x : Any
38
- Value produced inside a :func:`jax.vmap`-transformed function.
39
- op : {'all', 'any', 'none', 'max'}, default='any'
40
- Reduction to apply across the vmapped axis. ``'none'`` returns ``x`` without
41
- reduction, while ``'max'`` computes the maximum element.
42
-
43
- Returns
44
- -------
45
- Any
46
- Result of applying the requested reduction with vmap metadata removed.
47
-
48
- Raises
49
- ------
50
- ValueError
51
- If ``op`` is not one of ``'all'``, ``'any'``, ``'none'``, or ``'max'``.
52
-
53
- Examples
54
- --------
55
- .. code-block:: python
56
-
57
- >>> import jax.numpy as jnp
58
- >>> import brainstate
59
- >>>
60
- >>> xs = jnp.array([[True, False], [True, True]])
61
- >>> brainstate.transform.unvmap(xs, op='all')
62
- """
63
- if op == 'all':
64
- return unvmap_all(x)
65
- elif op == 'any':
66
- return unvmap_any(x)
67
- elif op == 'none':
68
- return _without_vmap(x)
69
- elif op == 'max':
70
- return unvmap_max(x)
71
- else:
72
- raise ValueError(f'Do not support type: {op}')
73
-
74
-
75
- # unvmap_all
76
-
77
- unvmap_all_p = Primitive("unvmap_all")
78
-
79
-
80
- def unvmap_all(x):
81
- """
82
- Evaluate :func:`jax.numpy.all` while ignoring vmapped batch dimensions.
83
-
84
- Parameters
85
- ----------
86
- x : Any
87
- Input array or pytree produced under :func:`jax.vmap`.
88
-
89
- Returns
90
- -------
91
- jax.Array
92
- Scalar boolean result of ``jnp.all(x)``.
93
-
94
- Examples
95
- --------
96
- .. code-block:: python
97
-
98
- >>> import jax.numpy as jnp
99
- >>> import brainstate
100
- >>>
101
- >>> values = jnp.array([[True, False], [True, True]])
102
- >>> brainstate.transform.unvmap(values, op='all')
103
- """
104
- return unvmap_all_p.bind(x)
105
-
106
-
107
- def _unvmap_all_impl(x):
108
- return jnp.all(x)
109
-
110
-
111
- def _unvmap_all_abstract_eval(x):
112
- return jax.core.ShapedArray(shape=(), dtype=jax.numpy.bool_.dtype) # pyright: ignore
113
-
114
-
115
- def _unvmap_all_batch(x, batch_axes):
116
- (x,) = x
117
- return unvmap_all(x), batching.not_mapped
118
-
119
-
120
- unvmap_all_p.def_impl(_unvmap_all_impl)
121
- unvmap_all_p.def_abstract_eval(_unvmap_all_abstract_eval)
122
- batching.primitive_batchers[unvmap_all_p] = _unvmap_all_batch # pyright: ignore
123
- mlir.register_lowering(
124
- unvmap_all_p,
125
- mlir.lower_fun(_unvmap_all_impl, multiple_results=False),
126
- )
127
-
128
- # unvmap_any
129
-
130
- unvmap_any_p = Primitive("unvmap_any")
131
-
132
-
133
- def unvmap_any(x):
134
- """
135
- Evaluate :func:`jax.numpy.any` while ignoring vmapped batch dimensions.
136
-
137
- Parameters
138
- ----------
139
- x : Any
140
- Input array or pytree produced under :func:`jax.vmap`.
141
-
142
- Returns
143
- -------
144
- jax.Array
145
- Scalar boolean result of ``jnp.any(x)``.
146
-
147
- Examples
148
- --------
149
- .. code-block:: python
150
-
151
- >>> import jax.numpy as jnp
152
- >>> import brainstate
153
- >>>
154
- >>> values = jnp.array([[False, False], [False, True]])
155
- >>> brainstate.transform.unvmap(values, op='any')
156
- """
157
- return unvmap_any_p.bind(x)
158
-
159
-
160
- def _unvmap_any_impl(x):
161
- return jnp.any(x)
162
-
163
-
164
- def _unvmap_any_abstract_eval(x):
165
- return jax.core.ShapedArray(shape=(), dtype=jax.numpy.bool_.dtype) # pyright: ignore
166
-
167
-
168
- def _unvmap_any_batch(x, batch_axes):
169
- (x,) = x
170
- return unvmap_any(x), batching.not_mapped
171
-
172
-
173
- unvmap_any_p.def_impl(_unvmap_any_impl)
174
- unvmap_any_p.def_abstract_eval(_unvmap_any_abstract_eval)
175
- batching.primitive_batchers[unvmap_any_p] = _unvmap_any_batch # pyright: ignore
176
- mlir.register_lowering(
177
- unvmap_any_p,
178
- mlir.lower_fun(_unvmap_any_impl, multiple_results=False),
179
- )
180
-
181
- # unvmap_max
182
-
183
- unvmap_max_p = Primitive("unvmap_max")
184
-
185
-
186
- def unvmap_max(x):
187
- """
188
- Evaluate :func:`jax.numpy.max` while ignoring vmapped batch dimensions.
189
-
190
- Parameters
191
- ----------
192
- x : Any
193
- Input array or pytree produced under :func:`jax.vmap`.
194
-
195
- Returns
196
- -------
197
- jax.Array
198
- Scalar containing the maximum value of ``x`` with the same dtype.
199
-
200
- Examples
201
- --------
202
- .. code-block:: python
203
-
204
- >>> import jax.numpy as jnp
205
- >>> import brainstate
206
- >>>
207
- >>> values = jnp.array([[1.0, 2.0], [0.5, 3.5]])
208
- >>> brainstate.transform.unvmap(values, op='max')
209
- """
210
- return unvmap_max_p.bind(x)
211
-
212
-
213
- def _unvmap_max_impl(x):
214
- return jnp.max(x)
215
-
216
-
217
- def _unvmap_max_abstract_eval(x):
218
- return jax.core.ShapedArray(shape=(), dtype=x.dtype)
219
-
220
-
221
- def _unvmap_max_batch(x, batch_axes):
222
- (x,) = x
223
- return unvmap_max(x), batching.not_mapped
224
-
225
-
226
- unvmap_max_p.def_impl(_unvmap_max_impl)
227
- unvmap_max_p.def_abstract_eval(_unvmap_max_abstract_eval)
228
- batching.primitive_batchers[unvmap_max_p] = _unvmap_max_batch # pyright: ignore
229
- mlir.register_lowering(
230
- unvmap_max_p,
231
- mlir.lower_fun(_unvmap_max_impl, multiple_results=False),
232
- )
233
-
234
-
235
- def _without_vmap(x):
236
- return _no_vmap_prim.bind(x)
237
-
238
-
239
- def _without_vmap_imp(x):
240
- return x
241
-
242
-
243
- def _without_vmap_abs(x):
244
- return x
245
-
246
-
247
- def _without_vmap_batch(x, batch_axes):
248
- (x,) = x
249
- return _without_vmap(x), batching.not_mapped
250
-
251
-
252
- _no_vmap_prim = Primitive('no_vmap')
253
- _no_vmap_prim.def_impl(_without_vmap_imp)
254
- _no_vmap_prim.def_abstract_eval(_without_vmap_abs)
255
- batching.primitive_batchers[_no_vmap_prim] = _without_vmap_batch
256
- mlir.register_lowering(_no_vmap_prim, mlir.lower_fun(_without_vmap_imp, multiple_results=False))
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
+ import jax
17
+ import jax.core
18
+ import jax.interpreters.batching as batching
19
+ import jax.interpreters.mlir as mlir
20
+ import jax.numpy as jnp
21
+
22
+ from brainstate._compatible_import import Primitive
23
+ from brainstate._utils import set_module_as
24
+
25
+ __all__ = [
26
+ "unvmap",
27
+ ]
28
+
29
+
30
+ @set_module_as('brainstate.transform')
31
+ def unvmap(x, op: str = 'any'):
32
+ """
33
+ Remove a leading vmap dimension by aggregating batched values.
34
+
35
+ Parameters
36
+ ----------
37
+ x : Any
38
+ Value produced inside a :func:`jax.vmap`-transformed function.
39
+ op : {'all', 'any', 'none', 'max'}, default='any'
40
+ Reduction to apply across the vmapped axis. ``'none'`` returns ``x`` without
41
+ reduction, while ``'max'`` computes the maximum element.
42
+
43
+ Returns
44
+ -------
45
+ Any
46
+ Result of applying the requested reduction with vmap metadata removed.
47
+
48
+ Raises
49
+ ------
50
+ ValueError
51
+ If ``op`` is not one of ``'all'``, ``'any'``, ``'none'``, or ``'max'``.
52
+
53
+ Examples
54
+ --------
55
+ .. code-block:: python
56
+
57
+ >>> import jax.numpy as jnp
58
+ >>> import brainstate
59
+ >>>
60
+ >>> xs = jnp.array([[True, False], [True, True]])
61
+ >>> brainstate.transform.unvmap(xs, op='all')
62
+ """
63
+ if op == 'all':
64
+ return unvmap_all(x)
65
+ elif op == 'any':
66
+ return unvmap_any(x)
67
+ elif op == 'none':
68
+ return _without_vmap(x)
69
+ elif op == 'max':
70
+ return unvmap_max(x)
71
+ else:
72
+ raise ValueError(f'Do not support type: {op}')
73
+
74
+
75
+ # unvmap_all
76
+
77
+ unvmap_all_p = Primitive("unvmap_all")
78
+
79
+
80
+ def unvmap_all(x):
81
+ """
82
+ Evaluate :func:`jax.numpy.all` while ignoring vmapped batch dimensions.
83
+
84
+ Parameters
85
+ ----------
86
+ x : Any
87
+ Input array or pytree produced under :func:`jax.vmap`.
88
+
89
+ Returns
90
+ -------
91
+ jax.Array
92
+ Scalar boolean result of ``jnp.all(x)``.
93
+
94
+ Examples
95
+ --------
96
+ .. code-block:: python
97
+
98
+ >>> import jax.numpy as jnp
99
+ >>> import brainstate
100
+ >>>
101
+ >>> values = jnp.array([[True, False], [True, True]])
102
+ >>> brainstate.transform.unvmap(values, op='all')
103
+ """
104
+ return unvmap_all_p.bind(x)
105
+
106
+
107
+ def _unvmap_all_impl(x):
108
+ return jnp.all(x)
109
+
110
+
111
+ def _unvmap_all_abstract_eval(x):
112
+ return jax.core.ShapedArray(shape=(), dtype=jax.numpy.bool_.dtype) # pyright: ignore
113
+
114
+
115
+ def _unvmap_all_batch(x, batch_axes):
116
+ (x,) = x
117
+ return unvmap_all(x), batching.not_mapped
118
+
119
+
120
+ unvmap_all_p.def_impl(_unvmap_all_impl)
121
+ unvmap_all_p.def_abstract_eval(_unvmap_all_abstract_eval)
122
+ batching.primitive_batchers[unvmap_all_p] = _unvmap_all_batch # pyright: ignore
123
+ mlir.register_lowering(
124
+ unvmap_all_p,
125
+ mlir.lower_fun(_unvmap_all_impl, multiple_results=False),
126
+ )
127
+
128
+ # unvmap_any
129
+
130
+ unvmap_any_p = Primitive("unvmap_any")
131
+
132
+
133
+ def unvmap_any(x):
134
+ """
135
+ Evaluate :func:`jax.numpy.any` while ignoring vmapped batch dimensions.
136
+
137
+ Parameters
138
+ ----------
139
+ x : Any
140
+ Input array or pytree produced under :func:`jax.vmap`.
141
+
142
+ Returns
143
+ -------
144
+ jax.Array
145
+ Scalar boolean result of ``jnp.any(x)``.
146
+
147
+ Examples
148
+ --------
149
+ .. code-block:: python
150
+
151
+ >>> import jax.numpy as jnp
152
+ >>> import brainstate
153
+ >>>
154
+ >>> values = jnp.array([[False, False], [False, True]])
155
+ >>> brainstate.transform.unvmap(values, op='any')
156
+ """
157
+ return unvmap_any_p.bind(x)
158
+
159
+
160
+ def _unvmap_any_impl(x):
161
+ return jnp.any(x)
162
+
163
+
164
+ def _unvmap_any_abstract_eval(x):
165
+ return jax.core.ShapedArray(shape=(), dtype=jax.numpy.bool_.dtype) # pyright: ignore
166
+
167
+
168
+ def _unvmap_any_batch(x, batch_axes):
169
+ (x,) = x
170
+ return unvmap_any(x), batching.not_mapped
171
+
172
+
173
+ unvmap_any_p.def_impl(_unvmap_any_impl)
174
+ unvmap_any_p.def_abstract_eval(_unvmap_any_abstract_eval)
175
+ batching.primitive_batchers[unvmap_any_p] = _unvmap_any_batch # pyright: ignore
176
+ mlir.register_lowering(
177
+ unvmap_any_p,
178
+ mlir.lower_fun(_unvmap_any_impl, multiple_results=False),
179
+ )
180
+
181
+ # unvmap_max
182
+
183
+ unvmap_max_p = Primitive("unvmap_max")
184
+
185
+
186
+ def unvmap_max(x):
187
+ """
188
+ Evaluate :func:`jax.numpy.max` while ignoring vmapped batch dimensions.
189
+
190
+ Parameters
191
+ ----------
192
+ x : Any
193
+ Input array or pytree produced under :func:`jax.vmap`.
194
+
195
+ Returns
196
+ -------
197
+ jax.Array
198
+ Scalar containing the maximum value of ``x`` with the same dtype.
199
+
200
+ Examples
201
+ --------
202
+ .. code-block:: python
203
+
204
+ >>> import jax.numpy as jnp
205
+ >>> import brainstate
206
+ >>>
207
+ >>> values = jnp.array([[1.0, 2.0], [0.5, 3.5]])
208
+ >>> brainstate.transform.unvmap(values, op='max')
209
+ """
210
+ return unvmap_max_p.bind(x)
211
+
212
+
213
+ def _unvmap_max_impl(x):
214
+ return jnp.max(x)
215
+
216
+
217
+ def _unvmap_max_abstract_eval(x):
218
+ return jax.core.ShapedArray(shape=(), dtype=x.dtype)
219
+
220
+
221
+ def _unvmap_max_batch(x, batch_axes):
222
+ (x,) = x
223
+ return unvmap_max(x), batching.not_mapped
224
+
225
+
226
+ unvmap_max_p.def_impl(_unvmap_max_impl)
227
+ unvmap_max_p.def_abstract_eval(_unvmap_max_abstract_eval)
228
+ batching.primitive_batchers[unvmap_max_p] = _unvmap_max_batch # pyright: ignore
229
+ mlir.register_lowering(
230
+ unvmap_max_p,
231
+ mlir.lower_fun(_unvmap_max_impl, multiple_results=False),
232
+ )
233
+
234
+
235
+ def _without_vmap(x):
236
+ return _no_vmap_prim.bind(x)
237
+
238
+
239
+ def _without_vmap_imp(x):
240
+ return x
241
+
242
+
243
+ def _without_vmap_abs(x):
244
+ return x
245
+
246
+
247
+ def _without_vmap_batch(x, batch_axes):
248
+ (x,) = x
249
+ return _without_vmap(x), batching.not_mapped
250
+
251
+
252
+ _no_vmap_prim = Primitive('no_vmap')
253
+ _no_vmap_prim.def_impl(_without_vmap_imp)
254
+ _no_vmap_prim.def_abstract_eval(_without_vmap_abs)
255
+ batching.primitive_batchers[_no_vmap_prim] = _without_vmap_batch
256
+ mlir.register_lowering(_no_vmap_prim, mlir.lower_fun(_without_vmap_imp, multiple_results=False))