@triton.jit
def rand_init(out_ptr, n, seed, BLOCK_SIZE: tl.constexpr):
    pid = tl.program_id(axis=0)
    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n
    random = tl.rand(seed, offsets)
    tl.store(out_ptr + offsets, random, mask=mask)
