tf.keras.ops.scan

Scan a function over leading array axes while carrying along state.

tf.keras.ops.scan(
    f, init, xs=None, length=None, reverse=False, unroll=1
)

When the type of xs is an array type or None, and the type of ys is an array type, the semantics of scan() are given roughly by this Python implementation:

def scan(f, init, xs, length=None):
    if xs is None:
        xs = [None] * length
    carry = init
    ys = []
    for x in xs:
        carry, y = f(carry, x)
        ys.append(y)
    return carry, np.stack(ys)

The loop-carried value carry (init) must hold a fixed shape and dtype across all iterations.

In TensorFlow, y must match carry in shape and dtype. This is not required in other backends.

f Callable defines the logic for each loop iteration. This 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. This callable returns a pair where the first represents a new value for the loop carry and the second represents a slice of the output.
init The initial loop carry value. This can be a scalar, tensor, or any nested structure. It must match the structure of the first element returned by f.
xs Optional value to scan along its leading axis. This can be a tensor or any nested structure. If xs is not provided, you must specify length to define the number of loop iterations. Defaults to None.
length Optional integer specifying the number of loop iterations. If length is not provided, it defaults to the sizes of leading axis of the arrays in xs. Defaults to None.
reverse Optional boolean specifying whether to run the scan iteration forward or in reverse, equivalent to reversing the leading axes of the arrays in both xs and in ys.
unroll Optional positive integer or boolean specifying 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 (unroll=True) or left completely unrolled (unroll=False). Note that unrolling is only supported by JAX and TensorFlow backends.

A pair where the first element represents the final loop carry value and the second element represents the stacked outputs of f when scanned over the leading axis of the inputs.

Examples:

sum_fn = lambda c, x: (c + x, c + x)
init = keras.ops.array(0)
xs = keras.ops.array([1, 2, 3, 4, 5])
carry, result = keras.ops.scan(sum_fn, init, xs)
carry
15
result
[1, 3, 6, 10, 15]