helion.language.randint#

helion.language.randint(shape, low, high, seed, device=None)[source]#

hl.randint provides a Philox-based pseudorandom integer generator (PRNG) that operates independently of PyTorch’s global random seed. Instead, it requires an explicit seed argument. Offsets are derived from the full logical sizes of the tiles specified in the shape argument.

Parameters:

shape (list[object]) – A list of sizes for the output tensor
low (int) – Lowest integer to be drawn from the distribution (inclusive)
high (int) – One above the highest integer to be drawn from the distribution (exclusive)
seed (int | Tensor) – A single element int64 tensor or int literal
device (device | None) – Device must match the current compile environment device

Returns:

A device tensor of int32 dtype filled with random integers in [low, high)

Return type:

torch.Tensor

Examples

@helion.kernel
def process_kernel(x: torch.Tensor) -> torch.Tensor:
    output = torch.zeros(x.shape, dtype=torch.int32, device=x.device)
    (m,) = x.shape
    for tile_m in hl.tile(m):
        output[tile_m] = hl.randint([tile_m], low=0, high=10, seed=42)
    return output