Ctrl+K

for_loop

Contents

for_loop#

class brainpy.math.for_loop(body_fun, operands, reverse=False, unroll=1, remat=False, jit=None, progress_bar=False, unroll_kwargs=None, dyn_vars=None, child_objs=None)[source]#

for-loop control flow with Variable.

New in version 2.1.11.

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).

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, 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, operands=(bm.arange(1, 5), bm.arange(2, 6)))
>>> a_hist
[[11.]
 [13.]
 [16.]
 [20.]]
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.

  • remat (bool) – Make fun recompute internal linearization points when differentiated.

  • jit (bool) – Whether to just-in-time compile the function.

  • 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.

  • progress_bar (bool) –

    Whether we use the progress bar to report the running progress.

    New in version 2.4.2.

  • dyn_vars (Variable, sequence of Variable, dict) –

    The instances of Variable.

    Deprecated since version 2.4.0: No longer need to provide dyn_vars. This function is capable of automatically collecting the dynamical variables used in the target func.

  • child_objs (optional, dict, sequence of BrainPyObject, BrainPyObject) –

    The children objects used in the target function.

    New in version 2.3.1.

    Deprecated since version 2.4.0: No longer need to provide child_objs. This function is capable of automatically collecting the children objects used in the target func.

  • unroll_kwargs (dict) – The keyword arguments without unrolling.

Returns:

outs – The stacked outputs of body_fun when scanned over the leading axis of the inputs.

Return type:

Any

Contents