This document is relevant for: Trn2, Trn3

nki.language.sequential_range#

nki.language.sequential_range(start, stop=None, step=1)[source]#

Create a sequence for fully unrolled loop iteration.

Create a sequence of numbers for use as loop iterators in NKI, resulting in a fully unrolled loop. Prefer static_range instead.

Warning

This API is deprecated and will be removed in future releases.

Parameters:

  • start – start value (or stop if stop is None).

  • stop – stop value (exclusive).

  • step – step size.

Returns:

an iterator yielding integer values from start to stop.

Examples:

import nki.language as nl

# nki.language.sequential_range
for i in nl.sequential_range(input_tensor.shape[1] // 512):
    offset = i * 512
    tile = nl.load(input_tensor[0:128, offset:offset+512])
    result = nl.multiply(tile, tile)
    nl.store(out_tensor[0:128, offset:offset+512], result)

This document is relevant for: Trn2, Trn3