Replace torch.rand with torch.randn

Author: voltjia

add.py

@@ -11,8 +11,8 @@
     dtype = torch.float16
     device = "cuda"
 
-    input = torch.rand(size, dtype=dtype, device=device)
-    other = torch.rand(size, dtype=dtype, device=device)
+    input = torch.randn(size, dtype=dtype, device=device)
+    other = torch.randn(size, dtype=dtype, device=device)
 
     ninetoothed_output = ops.ninetoothed.torch.add(input, other)
     torch_output = input + other

swiglu.py

@@ -20,9 +20,9 @@ def torch_swiglu(
     dtype = torch.float16
     device = "cuda"
 
-    a = torch.rand(shape, dtype=dtype, device=device)
-    b = torch.rand(shape, dtype=dtype, device=device)
-    c = torch.rand(shape, dtype=dtype, device=device)
+    a = torch.randn(shape, dtype=dtype, device=device)
+    b = torch.randn(shape, dtype=dtype, device=device)
+    c = torch.randn(shape, dtype=dtype, device=device)
 
     ninetoothed_output = ops.ninetoothed.torch.swiglu(a, b)
     torch_output = torch_swiglu(a, b)
@@ -58,8 +58,8 @@ def torch_swiglu(
     def benchmark(m, n, provider):
         shape = (m, n)
 
-        a = torch.rand(shape, dtype=dtype, device=device)
-        b = torch.rand(shape, dtype=dtype, device=device)
+        a = torch.randn(shape, dtype=dtype, device=device)
+        b = torch.randn(shape, dtype=dtype, device=device)
 
         if provider == "ninetoothed":
             ms = triton.testing.do_bench(lambda: ops.ninetoothed.torch.swiglu(a, b))