brainpy.math.for_loop#
- brainpy.math.for_loop(body_fun, operands, dyn_vars=None, out_vars=None, child_objs=None, reverse=False, unroll=1, remat=False, jit=True)[source]#
for-loop
control flow withVariable
.Changed in version 2.3.0:
dyn_vars
has been changed into a default argument. Please change your call fromfor_loop(fun, dyn_vars, operands)
tofor_loop(fun, operands, dyn_vars)
.Simply speaking, all dynamically changed variables used in the body function should be labeld in
dyn_vars
argument. All returns in body function will be gathered as the return of the whole loop.>>> import brainpy.math as bm >>> a = bm.Variable(bm.zeros(1)) >>> b = bm.Variable(bm.ones(1)) >>> # first example >>> def body(x): >>> a.value += x >>> b.value *= x >>> return a.value >>> a_hist = bm.for_loop(body, dyn_vars=[a, b], operands=bm.arange(1, 5)) >>> a_hist DeviceArray([[ 1.], [ 3.], [ 6.], [10.]], dtype=float32) >>> a Variable([10.], dtype=float32) >>> b Variable([24.], dtype=float32) >>> >>> # another example >>> def body(x, y): >>> a.value += x >>> b.value *= y >>> return a.value >>> a_hist = bm.for_loop(body, >>> dyn_vars=[a, b], >>> operands=(bm.arange(1, 5), bm.arange(2, 6))) >>> a_hist [[11.] [13.] [16.] [20.]]
New in version 2.1.11.
- Parameters:
body_fun (callable) – A Python function to be scanned. This function accepts one argument and returns one output. The argument denotes a slice of
operands
along its leading axis, and that output represents a slice of the return value.operands (Any) – The value over which to scan along the leading axis, where
operands
can be an array or any pytree (nested Python tuple/list/dict) thereof with consistent leading axis sizes. If body function body_func receives multiple arguments, operands should be a tuple/list whose length is equal to the number of arguments.dyn_vars (Variable, sequence of Variable, dict) – The instances of
Variable
.out_vars (PyTree, Optional) –
The variables to output in each step.
New in version 2.3.1: Support to return outputs with the specification of out_vars, which means no longer need to accumulate values through the function returns.
reverse (bool) – Optional boolean specifying whether to run the scan iteration forward (the default) or in reverse, equivalent to reversing the leading axes of the arrays in both
xs
and inys
.unroll (int) – Optional positive int specifying, in the underlying operation of the scan primitive, how many scan iterations to unroll within a single iteration of a loop.
child_objs (optional, dict, sequence of BrainPyObject, BrainPyObject) –
The children objects used in the target function.
New in version 2.3.1.
- Returns:
outs – The stacked outputs of
body_fun
when scanned over the leading axis of the inputs.- Return type:
Any