braincore.transform.for_loop

braincore.transform.for_loop(f, *xs, length=None, reverse=False, unroll=1, pbar=None)

for-loop control flow with State.

Parameters:

• f – a Python function to be scanned of type c -> a -> (c, b), meaning that f accepts two arguments where the first is a value of the loop carry and the second is a slice of xs along its leading axis, and that f returns a pair where the first element represents a new value for the loop carry and the second represents a slice of the output.

• xs – the value of type [a] over which to scan along the leading axis, where [a] can be an array or any pytree (nested Python tuple/list/dict) thereof with consistent leading axis sizes.

• length (Optional[int]) – optional integer specifying the number of loop iterations, which must agree with the sizes of leading axes of the arrays in xs (but can be used to perform scans where no input xs are needed).

• 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 | bool) – optional positive int or bool specifying, in the underlying operation of the scan primitive, how many scan iterations to unroll within a single iteration of a loop. If an integer is provided, it determines how many unrolled loop iterations to run within a single rolled iteration of the loop. If a boolean is provided, it will determine if the loop is completely unrolled (i.e. unroll=True) or left completely unrolled (i.e. unroll=False).

• pbar (Optional[ProgressBar]) – optional ProgressBar instance to display the progress of the scan operation.

Returns:

The return represents the stacked outputs of the second output of f when scanned over the leading axis of the inputs.