Random number construction fails within helion loop

[tomasruizt]

Describe the bug
Calling torch.rand with helion tiles as sizes fails, e.g torch.rand((tile_m, tile_n)).
This behavior contrasts with other supported torch APIs like torch.ones((tile_m, tile_n))

To Reproduce
Below is the simple matrix-multiply example. I added 1 + noise to each tile before saving. The 1s work, but the noise construction fails.

import torch
import helion
import helion.language as hl
from torch import Tensor
from typing import Callable


@helion.kernel(autotune_effort="none")
def matmul(
    x: Tensor,
    y: Tensor,
    epilogue: Callable[[Tensor, tuple[Tensor, ...]], Tensor] = lambda acc, tile: acc,
) -> Tensor:
    """
    Performs matrix multiplication of x and y with an optional epilogue function.
    Args:
        x (Tensor): Left matrix of shape [m, k].
        y (Tensor): Right matrix of shape [k, n].
        epilogue (Callable, optional): Function applied to the accumulator and tile indices
            after the matmul. Defaults to identity (no change).
    Returns:
        Tensor: Resulting matrix of shape [m, n].
    """
    m, k = x.size()
    k2, n = y.size()
    assert k == k2, f"size mismatch {k} != {k2}"
    out = torch.empty(
        [m, n], dtype=torch.promote_types(x.dtype, y.dtype), device=x.device
    )
    for tile_m, tile_n in hl.tile([m, n]):
        acc = hl.zeros([tile_m, tile_n], dtype=torch.float32)
        for tile_k in hl.tile(k):
            acc = torch.addmm(acc, x[tile_m, tile_k], y[tile_k, tile_n])
            ones = torch.ones((tile_m, tile_n), dtype=torch.float32)  # <<< succeeds
            noise = torch.rand((tile_m, tile_n), dtype=torch.float32)  # <<< fails
            acc = acc + ones + noise
        out[tile_m, tile_n] = epilogue(acc, (tile_m, tile_n))
    return out


torch.set_default_device("cuda:0")
x = torch.randn(10, 10)
y = torch.randn(10, 10)
x = matmul(x, y)
print(x.shape)
print("finished")

Expected behavior
torch.rand() should work just like torch.ones()

Versions
helion=0.2.1
torch=2.9.0

Running on RTX 3090 GPU.

[yf225]

We currently support torch.rand_like and torch.randn_like via #530. We welcome contributions to support torch.rand.

[tomasruizt]

@yf225 Thank you for the quick response. I tried to use torch.rand_like, but I get the same error: AssertionError: Not enough symbol arguments for dimensions. You can reproduce this by changing the single noise line with.

noise = torch.rand_like(ones, dtype=torch.float32)

[yf225]

@yf225 Thank you for the quick response. I tried to use torch.rand_like, but I get the same error: AssertionError: Not enough symbol arguments for dimensions. You can reproduce this by changing the single noise line with.

noise = torch.rand_like(ones, dtype=torch.float32)

@tomasruizt Thanks for the report! #1057 should fix it.

[tomasruizt]

@yf225 The matmul example works now! 🥳 However, I find another error constructing noise in the matmul_layernorm example. This time the error is a bit different.

...
  File "/home/tomasruiz/code/helion/helion/_compiler/inductor_lowering.py", line 1581, in _codegen_rng_op
    assert block_id is not None
           ^^^^^^^^^^^^^^^^^^^^
AssertionError

I took the example from the docs (https://helionlang.com/examples/matmul_layernorm.html) and added the single line constructing noise:

import torch
import helion
import helion.language as hl


@helion.kernel(autotune_effort="none")
def matmul_layernorm(
    x: torch.Tensor, y: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor
) -> torch.Tensor:
    m, k = x.size()
    k2 = y.size(0)
    n = hl.specialize(y.size(1))
    assert k == k2, f"size mismatch {k} != {k2}"
    assert weight.size(0) == n, f"weight size mismatch {weight.size(0)} != {n}"
    assert bias.size(0) == n, f"bias size mismatch {bias.size(0)} != {n}"
    out = torch.empty([m, n], dtype=torch.promote_types(x.dtype, y.dtype), device=x.device)
    for tile_m in hl.tile(m):
        acc = hl.zeros([tile_m, n], dtype=torch.float32)
        for tile_k in hl.tile(k):
            mm = torch.matmul(x[tile_m, tile_k], y[tile_k, :])
            acc = acc + mm
        noise = torch.rand_like(acc, dtype=torch.float32)  # <<< fails
        eps = 1e-5
        sum_vals = acc.sum(dim=-1, keepdim=True)
        mean = sum_vals / n
        centered = acc - mean
        var = (centered * centered).sum(dim=-1, keepdim=True) / n
        normalized = centered * torch.rsqrt(var + eps)
        acc = normalized * (weight[:].to(torch.float32)) + (bias[:].to(torch.float32))
        out[tile_m, :] = acc
    return out


if __name__ == "__main__":
    torch.set_default_device("cuda:0")
    m, k, n = 32, 64, 200
    x = torch.randn([m, k], dtype=torch.float16)
    y = torch.randn([k, n], dtype=torch.float16)
    weight = torch.randn([n], dtype=torch.float16)
    bias = torch.randn([n], dtype=torch.float16)
    out = matmul_layernorm(x, y, weight, bias)
    print("finished")

[tomasruizt]

Hi @yf225! Perhaps this issue should be reopened to reflect the still remaining issue?