issue/66: modified random sample test function

test/infiniop/causal_softmax.py

@@ -21,7 +21,6 @@ from libinfiniop import (
 #  Configuration (Internal Use Only)
 # ==============================================================================
 # These are not meant to be imported from other modules
-
 _TEST_CASES = [
     # x_shape, x_stride
     ((32, 512), None),

test/infiniop/random_sample.py

@@ -22,7 +22,6 @@ from libinfiniop import (
 #  Configuration (Internal Use Only)
 # ==============================================================================
 # These are not meant to be imported from other modules
-
 _TEST_CASES = [
     # voc, random_val, topp, topk, temperature
     (512, 0.8, 0.8, 3, 0.5),
@@ -59,6 +58,7 @@ infiniopRandomSampleDescriptor_t = POINTER(RandomSampleDescriptor)
 
 
 def random_sample(data, random_val, topp, topk, voc, temperature, torch_device):
+    if topp > 0 and topk > 1:
         indices = torch.zeros([topk], dtype=torch.int64)
         dataNp = data.clone().detach()
         sorted_indices = torch.arange(voc)
@@ -102,9 +102,7 @@ def random_sample(data, random_val, topp, topk, voc, temperature, torch_device):
             sum_s += dataNp[i]
             if random_val < sum_s:
                 return indices[i]
-
-
-def random_sample_0(data):
+    else:
         return torch.argmax(data)
 
 
@@ -124,12 +122,10 @@ def test(
     data = torch.arange(voc).float() * 0.0001
     _perm = torch.randperm(voc)
     data = data[_perm].to(x_dtype).to(torch_device)
-    if topp > 0 and topk > 1:
+
     ans = random_sample(
-            data.to("cpu"), random_val, topp, topk, voc, temperature, "cpu"
-        )
-    else:
-        ans = random_sample_0(data)
+        data, random_val, topp, topk, voc, temperature, torch_device
+    )  # 这个函数在device速度可能会很慢，可以通过data.to("cpu")方式加快计算过程
 
     indices = torch.zeros([1], dtype=torch.int64).to(torch_device)