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.
- brainstate/__init__.py +169 -169
- brainstate/_compatible_import.py +340 -340
- brainstate/_compatible_import_test.py +681 -681
- brainstate/_deprecation.py +210 -210
- brainstate/_deprecation_test.py +2319 -2319
- brainstate/_error.py +45 -45
- brainstate/_state.py +1652 -1652
- brainstate/_state_test.py +52 -52
- brainstate/_utils.py +47 -47
- brainstate/environ.py +1495 -1495
- brainstate/environ_test.py +1223 -1223
- brainstate/graph/__init__.py +22 -22
- brainstate/graph/_node.py +240 -240
- brainstate/graph/_node_test.py +589 -589
- brainstate/graph/_operation.py +1624 -1624
- brainstate/graph/_operation_test.py +1147 -1147
- brainstate/mixin.py +1433 -1433
- brainstate/mixin_test.py +1017 -1017
- brainstate/nn/__init__.py +137 -137
- brainstate/nn/_activations.py +1100 -1100
- brainstate/nn/_activations_test.py +354 -354
- brainstate/nn/_collective_ops.py +633 -633
- brainstate/nn/_collective_ops_test.py +774 -774
- brainstate/nn/_common.py +226 -226
- brainstate/nn/_common_test.py +154 -154
- brainstate/nn/_conv.py +2010 -2010
- brainstate/nn/_conv_test.py +849 -849
- brainstate/nn/_delay.py +575 -575
- brainstate/nn/_delay_test.py +243 -243
- brainstate/nn/_dropout.py +618 -618
- brainstate/nn/_dropout_test.py +477 -477
- brainstate/nn/_dynamics.py +1267 -1267
- brainstate/nn/_dynamics_test.py +67 -67
- brainstate/nn/_elementwise.py +1298 -1298
- brainstate/nn/_elementwise_test.py +829 -829
- brainstate/nn/_embedding.py +408 -408
- brainstate/nn/_embedding_test.py +156 -156
- brainstate/nn/_event_fixedprob.py +233 -233
- brainstate/nn/_event_fixedprob_test.py +115 -115
- brainstate/nn/_event_linear.py +83 -83
- brainstate/nn/_event_linear_test.py +121 -121
- brainstate/nn/_exp_euler.py +254 -254
- brainstate/nn/_exp_euler_test.py +377 -377
- brainstate/nn/_linear.py +744 -744
- brainstate/nn/_linear_test.py +475 -475
- brainstate/nn/_metrics.py +1070 -1070
- brainstate/nn/_metrics_test.py +611 -611
- brainstate/nn/_module.py +384 -384
- brainstate/nn/_module_test.py +40 -40
- brainstate/nn/_normalizations.py +1334 -1334
- brainstate/nn/_normalizations_test.py +699 -699
- brainstate/nn/_paddings.py +1020 -1020
- brainstate/nn/_paddings_test.py +722 -722
- brainstate/nn/_poolings.py +2239 -2239
- brainstate/nn/_poolings_test.py +952 -952
- brainstate/nn/_rnns.py +946 -946
- brainstate/nn/_rnns_test.py +592 -592
- brainstate/nn/_utils.py +216 -216
- brainstate/nn/_utils_test.py +401 -401
- brainstate/nn/init.py +809 -809
- brainstate/nn/init_test.py +180 -180
- brainstate/random/__init__.py +270 -270
- brainstate/random/_rand_funs.py +3938 -3938
- brainstate/random/_rand_funs_test.py +640 -640
- brainstate/random/_rand_seed.py +675 -675
- brainstate/random/_rand_seed_test.py +48 -48
- brainstate/random/_rand_state.py +1617 -1617
- brainstate/random/_rand_state_test.py +551 -551
- brainstate/transform/__init__.py +59 -59
- brainstate/transform/_ad_checkpoint.py +176 -176
- brainstate/transform/_ad_checkpoint_test.py +49 -49
- brainstate/transform/_autograd.py +1025 -1025
- brainstate/transform/_autograd_test.py +1289 -1289
- brainstate/transform/_conditions.py +316 -316
- brainstate/transform/_conditions_test.py +220 -220
- brainstate/transform/_error_if.py +94 -94
- brainstate/transform/_error_if_test.py +52 -52
- brainstate/transform/_eval_shape.py +145 -145
- brainstate/transform/_eval_shape_test.py +38 -38
- brainstate/transform/_jit.py +399 -399
- brainstate/transform/_jit_test.py +143 -143
- brainstate/transform/_loop_collect_return.py +675 -675
- brainstate/transform/_loop_collect_return_test.py +58 -58
- brainstate/transform/_loop_no_collection.py +283 -283
- brainstate/transform/_loop_no_collection_test.py +50 -50
- brainstate/transform/_make_jaxpr.py +2016 -2016
- brainstate/transform/_make_jaxpr_test.py +1510 -1510
- brainstate/transform/_mapping.py +529 -529
- brainstate/transform/_mapping_test.py +194 -194
- brainstate/transform/_progress_bar.py +255 -255
- brainstate/transform/_random.py +171 -171
- brainstate/transform/_unvmap.py +256 -256
- brainstate/transform/_util.py +286 -286
- brainstate/typing.py +837 -837
- brainstate/typing_test.py +780 -780
- brainstate/util/__init__.py +27 -27
- brainstate/util/_others.py +1024 -1024
- brainstate/util/_others_test.py +962 -962
- brainstate/util/_pretty_pytree.py +1301 -1301
- brainstate/util/_pretty_pytree_test.py +675 -675
- brainstate/util/_pretty_repr.py +462 -462
- brainstate/util/_pretty_repr_test.py +696 -696
- brainstate/util/filter.py +945 -945
- brainstate/util/filter_test.py +911 -911
- brainstate/util/struct.py +910 -910
- brainstate/util/struct_test.py +602 -602
- {brainstate-0.2.0.dist-info → brainstate-0.2.1.dist-info}/METADATA +108 -108
- brainstate-0.2.1.dist-info/RECORD +111 -0
- {brainstate-0.2.0.dist-info → brainstate-0.2.1.dist-info}/licenses/LICENSE +202 -202
- brainstate-0.2.0.dist-info/RECORD +0 -111
- {brainstate-0.2.0.dist-info → brainstate-0.2.1.dist-info}/WHEEL +0 -0
- {brainstate-0.2.0.dist-info → brainstate-0.2.1.dist-info}/top_level.txt +0 -0
@@ -1,316 +1,316 @@
|
|
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 collections.abc import Callable, Sequence
|
17
|
-
|
18
|
-
import jax
|
19
|
-
import jax.numpy as jnp
|
20
|
-
import numpy as np
|
21
|
-
|
22
|
-
from brainstate._compatible_import import to_concrete_aval, Tracer
|
23
|
-
from brainstate._utils import set_module_as
|
24
|
-
from ._error_if import jit_error_if
|
25
|
-
from ._make_jaxpr import StatefulFunction
|
26
|
-
from ._util import wrap_single_fun_in_multi_branches
|
27
|
-
|
28
|
-
__all__ = [
|
29
|
-
'cond', 'switch', 'ifelse',
|
30
|
-
]
|
31
|
-
|
32
|
-
|
33
|
-
@set_module_as('brainstate.transform')
|
34
|
-
def cond(pred, true_fun: Callable, false_fun: Callable, *operands):
|
35
|
-
"""
|
36
|
-
Conditionally apply ``true_fun`` or ``false_fun``.
|
37
|
-
|
38
|
-
Parameters
|
39
|
-
----------
|
40
|
-
pred : bool or array-like
|
41
|
-
Boolean scalar selecting which branch to execute. Numeric inputs are
|
42
|
-
treated as ``True`` when non-zero.
|
43
|
-
true_fun : Callable
|
44
|
-
Function that receives ``*operands`` when ``pred`` is ``True``.
|
45
|
-
false_fun : Callable
|
46
|
-
Function that receives ``*operands`` when ``pred`` is ``False``.
|
47
|
-
*operands : Any
|
48
|
-
Operands forwarded to either branch. May be any pytree of arrays,
|
49
|
-
scalars, or nested containers thereof.
|
50
|
-
|
51
|
-
Returns
|
52
|
-
-------
|
53
|
-
Any
|
54
|
-
Value returned by the selected branch with the same pytree structure
|
55
|
-
as produced by ``true_fun`` or ``false_fun``.
|
56
|
-
|
57
|
-
Notes
|
58
|
-
-----
|
59
|
-
Provided the arguments are correctly typed, :func:`cond` has semantics
|
60
|
-
that match the following Python implementation, where ``pred`` must be a
|
61
|
-
scalar:
|
62
|
-
|
63
|
-
.. code-block:: python
|
64
|
-
|
65
|
-
>>> def cond(pred, true_fun, false_fun, *operands):
|
66
|
-
... if pred:
|
67
|
-
... return true_fun(*operands)
|
68
|
-
... return false_fun(*operands)
|
69
|
-
|
70
|
-
In contrast with :func:`jax.lax.select`, using :func:`cond` indicates that only
|
71
|
-
one branch runs (subject to compiler rewrites and optimizations). When
|
72
|
-
transformed with :func:`~jax.vmap` over a batch of predicates, :func:`cond` is
|
73
|
-
converted to :func:`~jax.lax.select`.
|
74
|
-
|
75
|
-
Examples
|
76
|
-
--------
|
77
|
-
.. code-block:: python
|
78
|
-
|
79
|
-
>>> import brainstate
|
80
|
-
>>>
|
81
|
-
>>> def branch_true(x):
|
82
|
-
... return x + 1
|
83
|
-
>>>
|
84
|
-
>>> def branch_false(x):
|
85
|
-
... return x - 1
|
86
|
-
>>>
|
87
|
-
>>> brainstate.transform.cond(True, branch_true, branch_false, 3)
|
88
|
-
"""
|
89
|
-
if not (callable(true_fun) and callable(false_fun)):
|
90
|
-
raise TypeError("true_fun and false_fun arguments should be callable.")
|
91
|
-
|
92
|
-
if pred is None:
|
93
|
-
raise TypeError("cond predicate is None")
|
94
|
-
if isinstance(pred, Sequence) or np.ndim(pred) != 0:
|
95
|
-
raise TypeError(f"Pred must be a scalar, got {pred} of " +
|
96
|
-
(f"type {type(pred)}" if isinstance(pred, Sequence) else f"shape {np.shape(pred)}."))
|
97
|
-
|
98
|
-
# check pred
|
99
|
-
try:
|
100
|
-
pred_dtype = jax.dtypes.result_type(pred)
|
101
|
-
except TypeError as err:
|
102
|
-
raise TypeError("Pred type must be either boolean or number, got {}.".format(pred)) from err
|
103
|
-
if pred_dtype.kind != 'b':
|
104
|
-
if pred_dtype.kind in 'iuf':
|
105
|
-
pred = pred != 0
|
106
|
-
else:
|
107
|
-
raise TypeError("Pred type must be either boolean or number, got {}.".format(pred_dtype))
|
108
|
-
|
109
|
-
# not jit
|
110
|
-
if jax.config.jax_disable_jit and not isinstance(to_concrete_aval(pred), Tracer):
|
111
|
-
if pred:
|
112
|
-
return true_fun(*operands)
|
113
|
-
else:
|
114
|
-
return false_fun(*operands)
|
115
|
-
|
116
|
-
# evaluate jaxpr
|
117
|
-
stateful_true = StatefulFunction(true_fun, name='cond:true').make_jaxpr(*operands)
|
118
|
-
stateful_false = StatefulFunction(false_fun, name='conda:false').make_jaxpr(*operands)
|
119
|
-
|
120
|
-
# state trace and state values
|
121
|
-
state_trace = (stateful_true.get_state_trace(*operands) +
|
122
|
-
stateful_false.get_state_trace(*operands))
|
123
|
-
read_state_vals = state_trace.get_read_state_values(True)
|
124
|
-
write_state_vals = state_trace.get_write_state_values(True)
|
125
|
-
|
126
|
-
# wrap the functions
|
127
|
-
true_fun = wrap_single_fun_in_multi_branches(
|
128
|
-
stateful_true, state_trace, read_state_vals, True, stateful_true.get_arg_cache_key(*operands)
|
129
|
-
)
|
130
|
-
false_fun = wrap_single_fun_in_multi_branches(
|
131
|
-
stateful_false, state_trace, read_state_vals, True, stateful_false.get_arg_cache_key(*operands)
|
132
|
-
)
|
133
|
-
|
134
|
-
# cond
|
135
|
-
write_state_vals, out = jax.lax.cond(pred, true_fun, false_fun, write_state_vals, *operands)
|
136
|
-
|
137
|
-
# assign the written state values and restore the read state values
|
138
|
-
state_trace.assign_state_vals_v2(read_state_vals, write_state_vals)
|
139
|
-
return out
|
140
|
-
|
141
|
-
|
142
|
-
@set_module_as('brainstate.transform')
|
143
|
-
def switch(index, branches: Sequence[Callable], *operands):
|
144
|
-
"""
|
145
|
-
Apply exactly one branch from ``branches`` based on ``index``.
|
146
|
-
|
147
|
-
Parameters
|
148
|
-
----------
|
149
|
-
index : int or array-like
|
150
|
-
Scalar integer specifying which branch to execute.
|
151
|
-
branches : Sequence[Callable]
|
152
|
-
Sequence of callables; each receives ``*operands``.
|
153
|
-
*operands : Any
|
154
|
-
Operands forwarded to the selected branch. May be any pytree of arrays,
|
155
|
-
scalars, or nested containers thereof.
|
156
|
-
|
157
|
-
Returns
|
158
|
-
-------
|
159
|
-
Any
|
160
|
-
Value returned by the selected branch with the same pytree structure
|
161
|
-
as the selected callable.
|
162
|
-
|
163
|
-
Notes
|
164
|
-
-----
|
165
|
-
If ``index`` is out of bounds, it is clamped to ``[0, len(branches) - 1]``.
|
166
|
-
Conceptually, :func:`switch` behaves like:
|
167
|
-
|
168
|
-
.. code-block:: python
|
169
|
-
|
170
|
-
>>> def switch(index, branches, *operands):
|
171
|
-
... safe_index = clamp(0, index, len(branches) - 1)
|
172
|
-
... return branches[safe_index](*operands)
|
173
|
-
|
174
|
-
Internally this wraps XLA's `Conditional <https://www.tensorflow.org/xla/operation_semantics#conditional>`_
|
175
|
-
operator. When transformed with :func:`~jax.vmap` over a batch of predicates,
|
176
|
-
:func:`switch` is converted to :func:`~jax.lax.select`.
|
177
|
-
|
178
|
-
Examples
|
179
|
-
--------
|
180
|
-
.. code-block:: python
|
181
|
-
|
182
|
-
>>> import brainstate
|
183
|
-
>>>
|
184
|
-
>>> branches = (
|
185
|
-
... lambda x: x - 1,
|
186
|
-
... lambda x: x,
|
187
|
-
... lambda x: x + 1,
|
188
|
-
... )
|
189
|
-
>>>
|
190
|
-
>>> brainstate.transform.switch(2, branches, 3)
|
191
|
-
"""
|
192
|
-
# check branches
|
193
|
-
if not all(callable(branch) for branch in branches):
|
194
|
-
raise TypeError("branches argument should be a sequence of callables.")
|
195
|
-
|
196
|
-
# check index
|
197
|
-
if len(np.shape(index)) != 0:
|
198
|
-
raise TypeError(f"Branch index must be scalar, got {index} of shape {np.shape(index)}.")
|
199
|
-
try:
|
200
|
-
index_dtype = jax.dtypes.result_type(index)
|
201
|
-
except TypeError as err:
|
202
|
-
msg = f"Index type must be an integer, got {index}."
|
203
|
-
raise TypeError(msg) from err
|
204
|
-
if index_dtype.kind not in 'iu':
|
205
|
-
raise TypeError(f"Index type must be an integer, got {index} as {index_dtype}")
|
206
|
-
|
207
|
-
# format branches
|
208
|
-
branches = tuple(branches)
|
209
|
-
if len(branches) == 0:
|
210
|
-
raise ValueError("Empty branch sequence")
|
211
|
-
elif len(branches) == 1:
|
212
|
-
return branches[0](*operands)
|
213
|
-
|
214
|
-
# format index
|
215
|
-
index = jax.lax.convert_element_type(index, np.int32)
|
216
|
-
lo = np.array(0, np.int32)
|
217
|
-
hi = np.array(len(branches) - 1, np.int32)
|
218
|
-
index = jax.lax.clamp(lo, index, hi)
|
219
|
-
|
220
|
-
# not jit
|
221
|
-
if jax.config.jax_disable_jit and isinstance(jax.core.core.get_aval(index), jax.core.ConcreteArray):
|
222
|
-
return branches[int(index)](*operands)
|
223
|
-
|
224
|
-
# evaluate jaxpr
|
225
|
-
wrapped_branches = [StatefulFunction(branch, name='switch').make_jaxpr(*operands) for branch in branches]
|
226
|
-
|
227
|
-
# wrap the functions
|
228
|
-
state_trace = (wrapped_branches[0].get_state_trace(*operands) +
|
229
|
-
wrapped_branches[1].get_state_trace(*operands))
|
230
|
-
state_trace.merge(*[wrapped_branch.get_state_trace(*operands)
|
231
|
-
for wrapped_branch in wrapped_branches[2:]])
|
232
|
-
read_state_vals = state_trace.get_read_state_values(True)
|
233
|
-
write_state_vals = state_trace.get_write_state_values(True)
|
234
|
-
branches = [
|
235
|
-
wrap_single_fun_in_multi_branches(
|
236
|
-
wrapped_branch, state_trace, read_state_vals, True, wrapped_branch.get_arg_cache_key(*operands)
|
237
|
-
)
|
238
|
-
for wrapped_branch in wrapped_branches
|
239
|
-
]
|
240
|
-
|
241
|
-
# switch
|
242
|
-
write_state_vals, out = jax.lax.switch(index, branches, write_state_vals, *operands)
|
243
|
-
|
244
|
-
# write back state values or restore them
|
245
|
-
state_trace.assign_state_vals_v2(read_state_vals, write_state_vals)
|
246
|
-
return out
|
247
|
-
|
248
|
-
|
249
|
-
@set_module_as('brainstate.transform')
|
250
|
-
def ifelse(conditions, branches, *operands, check_cond: bool = True):
|
251
|
-
"""
|
252
|
-
Represent multi-way ``if``/``elif``/``else`` control flow.
|
253
|
-
|
254
|
-
Parameters
|
255
|
-
----------
|
256
|
-
conditions : Sequence[bool] or Array
|
257
|
-
Sequence of mutually exclusive boolean predicates. When ``check_cond`` is
|
258
|
-
``True``, exactly one entry must evaluate to ``True``.
|
259
|
-
branches : Sequence[Callable]
|
260
|
-
Sequence of branch callables evaluated lazily. Must have the same length as
|
261
|
-
``conditions``, contain at least two callables, and each branch receives
|
262
|
-
``*operands`` when selected.
|
263
|
-
*operands : Any
|
264
|
-
Operands forwarded to the selected branch as positional arguments.
|
265
|
-
check_cond : bool, default=True
|
266
|
-
Whether to verify that exactly one condition evaluates to ``True``.
|
267
|
-
|
268
|
-
Returns
|
269
|
-
-------
|
270
|
-
Any
|
271
|
-
Value produced by the branch corresponding to the active condition.
|
272
|
-
|
273
|
-
Notes
|
274
|
-
-----
|
275
|
-
When ``check_cond`` is ``True``, exactly one condition must evaluate to ``True``.
|
276
|
-
A common pattern is to make the final condition ``True`` to encode a default
|
277
|
-
branch.
|
278
|
-
|
279
|
-
Examples
|
280
|
-
--------
|
281
|
-
.. code-block:: python
|
282
|
-
|
283
|
-
>>> import brainstate
|
284
|
-
>>>
|
285
|
-
>>> def describe(a):
|
286
|
-
... return brainstate.transform.ifelse(
|
287
|
-
... conditions=[a > 5, a > 0, True],
|
288
|
-
... branches=[
|
289
|
-
... lambda: "greater than five",
|
290
|
-
... lambda: "positive",
|
291
|
-
... lambda: "non-positive",
|
292
|
-
... ],
|
293
|
-
... )
|
294
|
-
>>>
|
295
|
-
>>> describe(7)
|
296
|
-
>>> describe(-1)
|
297
|
-
"""
|
298
|
-
# check branches
|
299
|
-
if not all(callable(branch) for branch in branches):
|
300
|
-
raise TypeError("branches argument should be a sequence of callables.")
|
301
|
-
|
302
|
-
# format branches
|
303
|
-
branches = tuple(branches)
|
304
|
-
if len(branches) == 0:
|
305
|
-
raise ValueError("Empty branch sequence")
|
306
|
-
elif len(branches) == 1:
|
307
|
-
return branches[0](*operands)
|
308
|
-
if len(conditions) != len(branches):
|
309
|
-
raise ValueError("The number of conditions should be equal to the number of branches.")
|
310
|
-
|
311
|
-
# format index
|
312
|
-
conditions = jnp.asarray(conditions, np.int32)
|
313
|
-
if check_cond:
|
314
|
-
jit_error_if(jnp.sum(conditions) != 1, "Only one condition can be True. But got {}.", err_arg=conditions)
|
315
|
-
index = jnp.where(conditions, size=1, fill_value=len(conditions) - 1)[0][0]
|
316
|
-
return switch(index, branches, *operands)
|
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 collections.abc import Callable, Sequence
|
17
|
+
|
18
|
+
import jax
|
19
|
+
import jax.numpy as jnp
|
20
|
+
import numpy as np
|
21
|
+
|
22
|
+
from brainstate._compatible_import import to_concrete_aval, Tracer
|
23
|
+
from brainstate._utils import set_module_as
|
24
|
+
from ._error_if import jit_error_if
|
25
|
+
from ._make_jaxpr import StatefulFunction
|
26
|
+
from ._util import wrap_single_fun_in_multi_branches
|
27
|
+
|
28
|
+
__all__ = [
|
29
|
+
'cond', 'switch', 'ifelse',
|
30
|
+
]
|
31
|
+
|
32
|
+
|
33
|
+
@set_module_as('brainstate.transform')
|
34
|
+
def cond(pred, true_fun: Callable, false_fun: Callable, *operands):
|
35
|
+
"""
|
36
|
+
Conditionally apply ``true_fun`` or ``false_fun``.
|
37
|
+
|
38
|
+
Parameters
|
39
|
+
----------
|
40
|
+
pred : bool or array-like
|
41
|
+
Boolean scalar selecting which branch to execute. Numeric inputs are
|
42
|
+
treated as ``True`` when non-zero.
|
43
|
+
true_fun : Callable
|
44
|
+
Function that receives ``*operands`` when ``pred`` is ``True``.
|
45
|
+
false_fun : Callable
|
46
|
+
Function that receives ``*operands`` when ``pred`` is ``False``.
|
47
|
+
*operands : Any
|
48
|
+
Operands forwarded to either branch. May be any pytree of arrays,
|
49
|
+
scalars, or nested containers thereof.
|
50
|
+
|
51
|
+
Returns
|
52
|
+
-------
|
53
|
+
Any
|
54
|
+
Value returned by the selected branch with the same pytree structure
|
55
|
+
as produced by ``true_fun`` or ``false_fun``.
|
56
|
+
|
57
|
+
Notes
|
58
|
+
-----
|
59
|
+
Provided the arguments are correctly typed, :func:`cond` has semantics
|
60
|
+
that match the following Python implementation, where ``pred`` must be a
|
61
|
+
scalar:
|
62
|
+
|
63
|
+
.. code-block:: python
|
64
|
+
|
65
|
+
>>> def cond(pred, true_fun, false_fun, *operands):
|
66
|
+
... if pred:
|
67
|
+
... return true_fun(*operands)
|
68
|
+
... return false_fun(*operands)
|
69
|
+
|
70
|
+
In contrast with :func:`jax.lax.select`, using :func:`cond` indicates that only
|
71
|
+
one branch runs (subject to compiler rewrites and optimizations). When
|
72
|
+
transformed with :func:`~jax.vmap` over a batch of predicates, :func:`cond` is
|
73
|
+
converted to :func:`~jax.lax.select`.
|
74
|
+
|
75
|
+
Examples
|
76
|
+
--------
|
77
|
+
.. code-block:: python
|
78
|
+
|
79
|
+
>>> import brainstate
|
80
|
+
>>>
|
81
|
+
>>> def branch_true(x):
|
82
|
+
... return x + 1
|
83
|
+
>>>
|
84
|
+
>>> def branch_false(x):
|
85
|
+
... return x - 1
|
86
|
+
>>>
|
87
|
+
>>> brainstate.transform.cond(True, branch_true, branch_false, 3)
|
88
|
+
"""
|
89
|
+
if not (callable(true_fun) and callable(false_fun)):
|
90
|
+
raise TypeError("true_fun and false_fun arguments should be callable.")
|
91
|
+
|
92
|
+
if pred is None:
|
93
|
+
raise TypeError("cond predicate is None")
|
94
|
+
if isinstance(pred, Sequence) or np.ndim(pred) != 0:
|
95
|
+
raise TypeError(f"Pred must be a scalar, got {pred} of " +
|
96
|
+
(f"type {type(pred)}" if isinstance(pred, Sequence) else f"shape {np.shape(pred)}."))
|
97
|
+
|
98
|
+
# check pred
|
99
|
+
try:
|
100
|
+
pred_dtype = jax.dtypes.result_type(pred)
|
101
|
+
except TypeError as err:
|
102
|
+
raise TypeError("Pred type must be either boolean or number, got {}.".format(pred)) from err
|
103
|
+
if pred_dtype.kind != 'b':
|
104
|
+
if pred_dtype.kind in 'iuf':
|
105
|
+
pred = pred != 0
|
106
|
+
else:
|
107
|
+
raise TypeError("Pred type must be either boolean or number, got {}.".format(pred_dtype))
|
108
|
+
|
109
|
+
# not jit
|
110
|
+
if jax.config.jax_disable_jit and not isinstance(to_concrete_aval(pred), Tracer):
|
111
|
+
if pred:
|
112
|
+
return true_fun(*operands)
|
113
|
+
else:
|
114
|
+
return false_fun(*operands)
|
115
|
+
|
116
|
+
# evaluate jaxpr
|
117
|
+
stateful_true = StatefulFunction(true_fun, name='cond:true').make_jaxpr(*operands)
|
118
|
+
stateful_false = StatefulFunction(false_fun, name='conda:false').make_jaxpr(*operands)
|
119
|
+
|
120
|
+
# state trace and state values
|
121
|
+
state_trace = (stateful_true.get_state_trace(*operands) +
|
122
|
+
stateful_false.get_state_trace(*operands))
|
123
|
+
read_state_vals = state_trace.get_read_state_values(True)
|
124
|
+
write_state_vals = state_trace.get_write_state_values(True)
|
125
|
+
|
126
|
+
# wrap the functions
|
127
|
+
true_fun = wrap_single_fun_in_multi_branches(
|
128
|
+
stateful_true, state_trace, read_state_vals, True, stateful_true.get_arg_cache_key(*operands)
|
129
|
+
)
|
130
|
+
false_fun = wrap_single_fun_in_multi_branches(
|
131
|
+
stateful_false, state_trace, read_state_vals, True, stateful_false.get_arg_cache_key(*operands)
|
132
|
+
)
|
133
|
+
|
134
|
+
# cond
|
135
|
+
write_state_vals, out = jax.lax.cond(pred, true_fun, false_fun, write_state_vals, *operands)
|
136
|
+
|
137
|
+
# assign the written state values and restore the read state values
|
138
|
+
state_trace.assign_state_vals_v2(read_state_vals, write_state_vals)
|
139
|
+
return out
|
140
|
+
|
141
|
+
|
142
|
+
@set_module_as('brainstate.transform')
|
143
|
+
def switch(index, branches: Sequence[Callable], *operands):
|
144
|
+
"""
|
145
|
+
Apply exactly one branch from ``branches`` based on ``index``.
|
146
|
+
|
147
|
+
Parameters
|
148
|
+
----------
|
149
|
+
index : int or array-like
|
150
|
+
Scalar integer specifying which branch to execute.
|
151
|
+
branches : Sequence[Callable]
|
152
|
+
Sequence of callables; each receives ``*operands``.
|
153
|
+
*operands : Any
|
154
|
+
Operands forwarded to the selected branch. May be any pytree of arrays,
|
155
|
+
scalars, or nested containers thereof.
|
156
|
+
|
157
|
+
Returns
|
158
|
+
-------
|
159
|
+
Any
|
160
|
+
Value returned by the selected branch with the same pytree structure
|
161
|
+
as the selected callable.
|
162
|
+
|
163
|
+
Notes
|
164
|
+
-----
|
165
|
+
If ``index`` is out of bounds, it is clamped to ``[0, len(branches) - 1]``.
|
166
|
+
Conceptually, :func:`switch` behaves like:
|
167
|
+
|
168
|
+
.. code-block:: python
|
169
|
+
|
170
|
+
>>> def switch(index, branches, *operands):
|
171
|
+
... safe_index = clamp(0, index, len(branches) - 1)
|
172
|
+
... return branches[safe_index](*operands)
|
173
|
+
|
174
|
+
Internally this wraps XLA's `Conditional <https://www.tensorflow.org/xla/operation_semantics#conditional>`_
|
175
|
+
operator. When transformed with :func:`~jax.vmap` over a batch of predicates,
|
176
|
+
:func:`switch` is converted to :func:`~jax.lax.select`.
|
177
|
+
|
178
|
+
Examples
|
179
|
+
--------
|
180
|
+
.. code-block:: python
|
181
|
+
|
182
|
+
>>> import brainstate
|
183
|
+
>>>
|
184
|
+
>>> branches = (
|
185
|
+
... lambda x: x - 1,
|
186
|
+
... lambda x: x,
|
187
|
+
... lambda x: x + 1,
|
188
|
+
... )
|
189
|
+
>>>
|
190
|
+
>>> brainstate.transform.switch(2, branches, 3)
|
191
|
+
"""
|
192
|
+
# check branches
|
193
|
+
if not all(callable(branch) for branch in branches):
|
194
|
+
raise TypeError("branches argument should be a sequence of callables.")
|
195
|
+
|
196
|
+
# check index
|
197
|
+
if len(np.shape(index)) != 0:
|
198
|
+
raise TypeError(f"Branch index must be scalar, got {index} of shape {np.shape(index)}.")
|
199
|
+
try:
|
200
|
+
index_dtype = jax.dtypes.result_type(index)
|
201
|
+
except TypeError as err:
|
202
|
+
msg = f"Index type must be an integer, got {index}."
|
203
|
+
raise TypeError(msg) from err
|
204
|
+
if index_dtype.kind not in 'iu':
|
205
|
+
raise TypeError(f"Index type must be an integer, got {index} as {index_dtype}")
|
206
|
+
|
207
|
+
# format branches
|
208
|
+
branches = tuple(branches)
|
209
|
+
if len(branches) == 0:
|
210
|
+
raise ValueError("Empty branch sequence")
|
211
|
+
elif len(branches) == 1:
|
212
|
+
return branches[0](*operands)
|
213
|
+
|
214
|
+
# format index
|
215
|
+
index = jax.lax.convert_element_type(index, np.int32)
|
216
|
+
lo = np.array(0, np.int32)
|
217
|
+
hi = np.array(len(branches) - 1, np.int32)
|
218
|
+
index = jax.lax.clamp(lo, index, hi)
|
219
|
+
|
220
|
+
# not jit
|
221
|
+
if jax.config.jax_disable_jit and isinstance(jax.core.core.get_aval(index), jax.core.ConcreteArray):
|
222
|
+
return branches[int(index)](*operands)
|
223
|
+
|
224
|
+
# evaluate jaxpr
|
225
|
+
wrapped_branches = [StatefulFunction(branch, name='switch').make_jaxpr(*operands) for branch in branches]
|
226
|
+
|
227
|
+
# wrap the functions
|
228
|
+
state_trace = (wrapped_branches[0].get_state_trace(*operands) +
|
229
|
+
wrapped_branches[1].get_state_trace(*operands))
|
230
|
+
state_trace.merge(*[wrapped_branch.get_state_trace(*operands)
|
231
|
+
for wrapped_branch in wrapped_branches[2:]])
|
232
|
+
read_state_vals = state_trace.get_read_state_values(True)
|
233
|
+
write_state_vals = state_trace.get_write_state_values(True)
|
234
|
+
branches = [
|
235
|
+
wrap_single_fun_in_multi_branches(
|
236
|
+
wrapped_branch, state_trace, read_state_vals, True, wrapped_branch.get_arg_cache_key(*operands)
|
237
|
+
)
|
238
|
+
for wrapped_branch in wrapped_branches
|
239
|
+
]
|
240
|
+
|
241
|
+
# switch
|
242
|
+
write_state_vals, out = jax.lax.switch(index, branches, write_state_vals, *operands)
|
243
|
+
|
244
|
+
# write back state values or restore them
|
245
|
+
state_trace.assign_state_vals_v2(read_state_vals, write_state_vals)
|
246
|
+
return out
|
247
|
+
|
248
|
+
|
249
|
+
@set_module_as('brainstate.transform')
|
250
|
+
def ifelse(conditions, branches, *operands, check_cond: bool = True):
|
251
|
+
"""
|
252
|
+
Represent multi-way ``if``/``elif``/``else`` control flow.
|
253
|
+
|
254
|
+
Parameters
|
255
|
+
----------
|
256
|
+
conditions : Sequence[bool] or Array
|
257
|
+
Sequence of mutually exclusive boolean predicates. When ``check_cond`` is
|
258
|
+
``True``, exactly one entry must evaluate to ``True``.
|
259
|
+
branches : Sequence[Callable]
|
260
|
+
Sequence of branch callables evaluated lazily. Must have the same length as
|
261
|
+
``conditions``, contain at least two callables, and each branch receives
|
262
|
+
``*operands`` when selected.
|
263
|
+
*operands : Any
|
264
|
+
Operands forwarded to the selected branch as positional arguments.
|
265
|
+
check_cond : bool, default=True
|
266
|
+
Whether to verify that exactly one condition evaluates to ``True``.
|
267
|
+
|
268
|
+
Returns
|
269
|
+
-------
|
270
|
+
Any
|
271
|
+
Value produced by the branch corresponding to the active condition.
|
272
|
+
|
273
|
+
Notes
|
274
|
+
-----
|
275
|
+
When ``check_cond`` is ``True``, exactly one condition must evaluate to ``True``.
|
276
|
+
A common pattern is to make the final condition ``True`` to encode a default
|
277
|
+
branch.
|
278
|
+
|
279
|
+
Examples
|
280
|
+
--------
|
281
|
+
.. code-block:: python
|
282
|
+
|
283
|
+
>>> import brainstate
|
284
|
+
>>>
|
285
|
+
>>> def describe(a):
|
286
|
+
... return brainstate.transform.ifelse(
|
287
|
+
... conditions=[a > 5, a > 0, True],
|
288
|
+
... branches=[
|
289
|
+
... lambda: "greater than five",
|
290
|
+
... lambda: "positive",
|
291
|
+
... lambda: "non-positive",
|
292
|
+
... ],
|
293
|
+
... )
|
294
|
+
>>>
|
295
|
+
>>> describe(7)
|
296
|
+
>>> describe(-1)
|
297
|
+
"""
|
298
|
+
# check branches
|
299
|
+
if not all(callable(branch) for branch in branches):
|
300
|
+
raise TypeError("branches argument should be a sequence of callables.")
|
301
|
+
|
302
|
+
# format branches
|
303
|
+
branches = tuple(branches)
|
304
|
+
if len(branches) == 0:
|
305
|
+
raise ValueError("Empty branch sequence")
|
306
|
+
elif len(branches) == 1:
|
307
|
+
return branches[0](*operands)
|
308
|
+
if len(conditions) != len(branches):
|
309
|
+
raise ValueError("The number of conditions should be equal to the number of branches.")
|
310
|
+
|
311
|
+
# format index
|
312
|
+
conditions = jnp.asarray(conditions, np.int32)
|
313
|
+
if check_cond:
|
314
|
+
jit_error_if(jnp.sum(conditions) != 1, "Only one condition can be True. But got {}.", err_arg=conditions)
|
315
|
+
index = jnp.where(conditions, size=1, fill_value=len(conditions) - 1)[0][0]
|
316
|
+
return switch(index, branches, *operands)
|