fix: xumpy.random() honors dtype now

ndsl/xumpy/alloc.py

@@ -2,6 +2,7 @@
 
 import numpy as np
 import numpy.typing as npt
+from numpy._typing import _SupportsDType
 
 from ndsl.config import Backend
 from ndsl.dsl.typing import Float
@@ -13,6 +14,12 @@
 
 # Taking a page from cupy's playbook to have tuple & ndarray
 _ShapeLike = SupportsIndex | Sequence[SupportsIndex]
+_DTypeLikeFloat32 = (
+    np.dtype[np.float32] | _SupportsDType[np.dtype[np.float32]] | type[np.float32]
+)
+_DTypeLikeFloat64 = (
+    np.dtype[np.float64] | _SupportsDType[np.dtype[np.float64]] | type[np.float64]
+)
 
 
 def zeros(
@@ -59,8 +66,11 @@ def full(
 def random(
     shape: _ShapeLike,
     backend: Backend,
-    dtype: npt.DTypeLike = Float,
+    dtype: _DTypeLikeFloat32 | _DTypeLikeFloat64 = Float,  # type: ignore [valid-type]
 ) -> np.ndarray | cp.ndarray:
     if backend.is_gpu_backend():
-        cp.random.rand(*shape)
-    return np.random.rand(*shape)
+        gen = cp.random.default_rng()
+        return gen.random(shape, dtype, None)
+
+    gen = np.random.default_rng()
+    return gen.random(shape, dtype, None)

tests/xumpy/test_alloc.py

@@ -0,0 +1,31 @@
+import numpy as np
+import pytest
+
+from ndsl import xumpy as xp
+from ndsl.config import Backend
+
+
+@pytest.mark.parametrize("dtype", [None, np.float32, np.float64])
+def test_random(dtype) -> None:
+    shape = (2, 3, 5)
+    rand_array = xp.random(shape, Backend.python())
+    assert rand_array.shape == shape
+    assert (rand_array != xp.random(shape, Backend.python())).all()
+
+
+def test_ones() -> None:
+    shape = (2, 3, 5)
+    assert (np.ones(shape) == xp.ones(shape, Backend.python())).all()
+
+
+def test_zeros() -> None:
+    shape = (2, 3, 5)
+    assert (np.zeros(shape) == xp.zeros(shape, Backend.python())).all()
+
+
+def test_full() -> None:
+    shape = (2, 3, 5)
+    value = 42.42
+    assert (
+        np.full(shape, value) == xp.full(shape, value=value, backend=Backend.python())
+    ).all()

tests/xumpy/test_count_nonzeros.py

@@ -0,0 +1,12 @@
+import numpy as np
+
+from ndsl import xumpy as xp
+from ndsl.config import Backend
+
+
+def test_count_nonzero():
+    shape = (2, 3, 5)
+    rand_array = xp.random(shape, Backend.python())
+    rand_array[1, 1, :] = 0
+
+    assert np.count_nonzero(rand_array) == xp.count_nonzero(rand_array)

tests/xumpy/test_minmax.py

@@ -0,0 +1,17 @@
+import numpy as np
+
+from ndsl import xumpy as xp
+from ndsl.config import Backend
+
+
+def test_minmax():
+    shape = (2, 3, 5)
+    rand_array = xp.random(shape, Backend.python())
+
+    assert (np.max(rand_array, axis=1) == xp.max(rand_array, axis=1)).all()
+    assert (np.min(rand_array, axis=1) == xp.min(rand_array, axis=1)).all()
+
+    out_buffer = xp.empty(shape, Backend.python())
+    xp.max_on_horizontal_plane(rand_array, out_buffer)
+
+    assert (np.max(rand_array, axis=(0, 1)) == out_buffer).all()