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)[source]#

for-loop control flow with Variable.

Changed in version 2.3.0: dyn_vars has been changed into a default argument. Please change your call from for_loop(fun, dyn_vars, operands) to for_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 in ys.

  • 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