Skip to content

Revert "Add hl.rand op with seed arg lowering to tl.rand (#652)" #698

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Sep 26, 2025
Merged

Revert "Add hl.rand op with seed arg lowering to tl.rand (#652)" #698

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 0 additions & 1 deletion helion/language/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,6 @@
from .matmul_ops import dot as dot
from .memory_ops import load as load
from .memory_ops import store as store
from .random_ops import rand as rand
from .reduce_ops import reduce as reduce
from .scan_ops import associative_scan as associative_scan
from .scan_ops import cumprod as cumprod
Expand Down
121 changes: 0 additions & 121 deletions helion/language/random_ops.py
View file
Open in desktop

This file was deleted.

91 changes: 0 additions & 91 deletions test/test_rng.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -348,97 +348,6 @@ def randn_kernel_3d(x: torch.Tensor) -> torch.Tensor:
f"Slice {b_idx} std {slice_std} is not well distributed",
)

def test_hl_rand_1d(self):
@helion.kernel
def rand_kernel_tiled_1d(x: torch.Tensor, seed: int) -> torch.Tensor:
output = torch.zeros_like(x)
(m,) = x.shape
for (tile_m,) in hl.tile([m]):
output[tile_m] = hl.rand([tile_m], seed=seed)
return output

x_small = torch.ones(128, device=DEVICE)
_, output = code_and_output(rand_kernel_tiled_1d, (x_small, 42))
_, output2 = code_and_output(rand_kernel_tiled_1d, (x_small, 1337))

self.assertFalse(
torch.allclose(output, output2),
"Different seeds should produce different outputs",
)

_, output3 = code_and_output(rand_kernel_tiled_1d, (x_small, 42))
self.assertTrue(
torch.allclose(output, output3),
"Same seed should produce identical outputs",
)

# Check that all values are in [0, 1) range
self.assertTrue(torch.all(output >= 0.0), "All values should be >= 0")
self.assertTrue(torch.all(output < 1.0), "All values should be < 1")

def test_hl_rand_2d(self):
@helion.kernel
def rand_kernel_tiled_2d(x: torch.Tensor, seed: int) -> torch.Tensor:
output = torch.zeros_like(x)
m, n = x.shape
for tile_m, tile_n in hl.tile([m, n]):
output[tile_m, tile_n] = hl.rand([tile_m, tile_n], seed=seed)
return output

x_small = torch.ones(128, 128, device=DEVICE)
_, output = code_and_output(rand_kernel_tiled_2d, (x_small, 42))
_, output2 = code_and_output(rand_kernel_tiled_2d, (x_small, 1337))

self.assertFalse(
torch.allclose(output, output2),
"Different seeds should produce different outputs",
)

_, output3 = code_and_output(rand_kernel_tiled_2d, (x_small, 42))
self.assertTrue(
torch.allclose(output, output3),
"Same seed should produce identical outputs",
)

self.assertTrue(torch.all(output >= 0.0), "All values should be >= 0")
self.assertTrue(torch.all(output < 1.0), "All values should be < 1")

def test_hl_rand_3d(self):
@helion.kernel
def rand_kernel_tiled_3d(x: torch.Tensor, seed: int) -> torch.Tensor:
output = torch.zeros_like(x)
b, m, n = x.shape
for tile_b, tile_m, tile_n in hl.tile([b, m, n]):
output[tile_b, tile_m, tile_n] = hl.rand(
[tile_b, tile_m, tile_n], seed=seed
)
return output

x_small = torch.ones(16, 32, 64, device=DEVICE)
_, output = code_and_output(rand_kernel_tiled_3d, (x_small, 42))
_, output2 = code_and_output(rand_kernel_tiled_3d, (x_small, 1337))

self.assertFalse(
torch.allclose(output, output2),
"Different seeds should produce different outputs",
)

_, output3 = code_and_output(rand_kernel_tiled_3d, (x_small, 42))
self.assertTrue(
torch.allclose(output, output3),
"Same seed should produce identical outputs",
)

self.assertTrue(torch.all(output >= 0.0), "All values should be >= 0")
self.assertTrue(torch.all(output < 1.0), "All values should be < 1")

# Check distribution properties
mean_val = output.mean().item()
self.assertTrue(
0.4 < mean_val < 0.6,
f"Mean {mean_val:.3f} should be around 0.5 for uniform distribution",
)


if __name__ == "__main__":
unittest.main()
Loading