scan

Contents

scan#

class brainpy.math.scan(body_fun, init, operands, reverse=False, unroll=1, remat=False, progress_bar=False)[source]#

scan control flow with Variable.

Similar to jax.lax.scan.

Added in version 2.4.7.

All returns in body function will be gathered as the return of the whole loop.

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.

  • init (Any) – An initial loop carry value of type c, which can be a scalar, array, or any pytree (nested Python tuple/list/dict) thereof, representing the initial loop carry value. This value must have the same structure as the first element of the pair returned by body_fun.

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

  • 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 (Union[bool, ProgressBar, int]) –

    Whether and how to display a progress bar during execution:

    • False (default): No progress bar

    • True: Display progress bar with default settings

    • ProgressBar instance: Display progress bar with custom settings

    • int: Display progress bar updating every N iterations (treated as freq parameter)

    See for_loop() for detailed examples of ProgressBar usage.

    Added in version 2.4.2.

    Changed in version 2.7.3: Now accepts ProgressBar instances and integers for advanced customization.

Returns:

outs – A two-element tuple (final_carry, stacked_ys):

  • final_carry: the loop carry value returned by the last iteration of body_fun (same structure as init).

  • stacked_ys: the per-iteration outputs of body_fun stacked along a new leading axis.

Return type:

tuple

Contents