test: generate fallback randint data on cpu

Author: zkjh

tests/utils.py

@@ -74,10 +74,22 @@ def rand_strided(shape, strides, *, dtype=None, device=None):
 
 def randint_strided(low, high, shape, strides, *, dtype=None, device=None):
     output = empty_strided(shape, strides, dtype=dtype, device=device)
-
-    output.as_strided(
+    output_flat = output.as_strided(
         (output.untyped_storage().size() // output.element_size(),), (1,)
-    ).random_(low, high)
+    )
+
+    try:
+        output_flat.random_(low, high)
+    except RuntimeError as exc:
+        if "random_" not in str(exc):
+            raise
+
+        cpu_output = empty_strided(shape, strides, dtype=dtype, device="cpu")
+        cpu_flat = cpu_output.as_strided(
+            (cpu_output.untyped_storage().size() // cpu_output.element_size(),), (1,)
+        )
+        cpu_flat.random_(low, high)
+        output_flat.copy_(cpu_flat.to(device=output.device))
 
     return output