brainpy.math.controls.for_loop#

brainpy.math.controls.for_loop(body_fun, dyn_vars, operands, reverse=False, unroll=1)[source]#

for-loop control flow with Variable.

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.

  • dyn_vars (Variable, sequence of Variable, dict) – The instances of Variable.

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

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

Returns

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

Return type

Any