helion.language.rand4x#

helion.language.rand4x(seed, offsets, device=None)[source]#

hl.rand4x returns four independent uniform float32 tensors in [0, 1) per offset from a single Philox round (~4× cheaper than four separate hl.rand() calls). Mirrors Triton’s tl.rand4x.

Parameters:

seed (int | Tensor) – A single element int64 tensor or int literal
offsets (Tensor) – Int64 offset tensor fed into the Philox RNG. The output tensors have shape equal to offsets.shape.
device (device | None) – Device must match the current compile environment device

Returns:

Four float32 tensors, each with shape offsets.shape and values in [0, 1).

Return type:

tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]

Examples

Three sibling dropout masks per element with a single Philox call:

@helion.kernel
def triple_dropout_kernel(x: torch.Tensor) -> torch.Tensor:
    out = torch.empty_like(x)
    (n,) = x.shape
    for tile in hl.tile(n):
        base = hl.tile_index(tile).to(torch.int64)
        r0, r1, r2, _ = hl.rand4x(seed=42, offsets=base)
        keep = (r0 > 0.1) & (r1 > 0.2) & (r2 > 0.3)
        out[tile] = x[tile] * keep.to(x.dtype)
    return out