returnn.frontend.loop

While loop.

Why did we choose this API? To allow for both eager-based and graph-based frameworks, and also to avoid any magic happening here. https://github.com/rwth-i6/returnn/issues/1282 https://github.com/rwth-i6/returnn/issues/1324

returnn.frontend.loop.while_loop(cond: Callable[[S], bool | Tensor], body: Callable[[S], S], initial: S) → S[source]

It executes:

while cond(loop_vars):
    loop_vars = body(loop_vars)

And then it returns the final loop vars.

If you want to iterate over some axis, maybe of an existing tensor, or if you want to accumulate frames in each iteration, see scan().

Parameters:

  • cond

  • body

  • initial – initial loop vars

Returns:

final loop vars

returnn.frontend.loop.scan(*, spatial_dim: Dim | None = None, cond_dims: Sequence[Dim] | None = None, cond_before_body: bool = True, initial: S | None = None, xs: X | None = None, ys: Y | None = None, cond: Callable[[X, S], Tensor] | None = None, body: Callable[[X, S], Tuple[Y, S]], max_seq_len: int | Tensor | None = None, return_tensor_arrays: bool = False) → Tuple[Y, S, Dim][source]

Extended variant of while_loop().

Supports iterating over a given axis (spatial_dim), supports iterating over some input tensors (xs: X) on the given axis, and supports returning some frame-accumulated output tensors (ys: Y).

https://github.com/rwth-i6/returnn/issues/1324

Parameters:

  • spatial_dim – if None or unknown, need to provide cond. must be given if xs is not None.

  • cond_dims – if spatial_dim is not given, this must be given to know what shape to expect from cond. This will also be the shape of the dyn_size_ext of the resulting spatial_dim.

  • cond_before_body – if True, will execute cond before body, otherwise after. If True, corresponds to while cond(...): body(...), otherwise while True: body(...); if not cond(...): break. Note that cond is executed in any case at the end of the loop but with cond_before_body=True this final value would be ignored. Be careful though that you do not have any side-effects in cond.

  • initial – state/carry

  • xs – input, will be unstacked over spatial_dim. can be None.

  • ys – output, as templates, per iteration (excluding spatial_dim)

  • cond – if spatial_dim is None/unknown, need to provide this. The shape will be the same as the dyn_size_ext of the resulting spatial_dim. Unlike while_loop cond, does not need to be scalar. E.g. some shape like [B] is normal. Once it returns False for all entries, the loop will stop. Once it returns False for some entry, further values in further iterations for this entry will be ignored. We do not expect any side-effects in cond.

  • body

  • max_seq_len – If given, it is checked in addition to cond, and when reached, it stops the loop.

  • return_tensor_arrays – if True, will return TensorArray instead of Tensor for ys. Internally, we work with TensorArray anyway, so this avoids the final stack(). In case of beam search, it might make more sense to perform some post-processing on the TensorArray per entry, like selecting the right beam entries.

Returns:

outputs ys, final state, and the new spatial_dim