Merge pull request #5 from InfiniTensor/use-custom-gauss

Use custom `gauss`

Author: voltjia

tests/test_add.py

@@ -1,11 +1,9 @@
-import random
-
 import pytest
 import torch
 
 import ntops
 from tests.skippers import skip_if_cuda_not_available
-from tests.utils import generate_arguments
+from tests.utils import gauss, generate_arguments
 
 
 @skip_if_cuda_not_available
@@ -15,7 +13,7 @@ def test_cuda(shape, dtype, atol, rtol):
 
     input = torch.randn(shape, dtype=dtype, device=device)
     other = torch.randn(shape, dtype=dtype, device=device)
-    alpha = random.gauss()
+    alpha = gauss()
 
     ninetoothed_output = ntops.add(input, other, alpha)
     reference_output = input + alpha * other

tests/test_addmm.py

@@ -1,11 +1,10 @@
-import random
-
 import pytest
 import torch
 
 import ntops
 from tests.skippers import skip_if_cuda_not_available
 from tests.test_mm import generate_arguments
+from tests.utils import gauss
 
 
 @skip_if_cuda_not_available
@@ -16,8 +15,8 @@ def test_cuda(m, n, k, dtype, atol, rtol):
     input = torch.randn((m, n), dtype=dtype, device=device)
     x = torch.randn((m, k), dtype=dtype, device=device)
     y = torch.randn((k, n), dtype=dtype, device=device)
-    beta = random.gauss()
-    alpha = random.gauss()
+    beta = gauss()
+    alpha = gauss()
 
     ninetoothed_output = ntops.addmm(input, x, y, beta=beta, alpha=alpha)
     reference_output = torch.addmm(input, x, y, beta=beta, alpha=alpha)

tests/utils.py

@@ -21,6 +21,10 @@ def generate_arguments():
     return "shape, dtype, atol, rtol", arguments
 
 
+def gauss(mu=0.0, sigma=1.0):
+    return random.gauss(mu, sigma)
+
+
 def _random_shape(ndim, min_num_elements=2**8, max_num_elements=2**10):
     num_elements = random.randint(min_num_elements, max_num_elements)