File: C:/Users/fred/anaconda3/Lib/site-packages/dask/array/tests/test_cupy_random.py
from __future__ import annotations
import pytest
pytestmark = pytest.mark.gpu
import numpy as np
import dask.array as da
from dask import config
from dask.array.utils import assert_eq
cupy = pytest.importorskip("cupy")
@pytest.mark.parametrize("backend", ["cupy", "numpy"])
@pytest.mark.parametrize("rs", [None, cupy.random.RandomState, np.random.RandomState])
def test_random_all_RandomState(backend, rs):
# RandomState argument takes priority over backend
if rs == cupy.random.RandomState:
expect = cupy.ndarray
elif rs == np.random.RandomState:
expect = np.ndarray
elif backend == "cupy":
expect = cupy.ndarray
else:
expect = np.ndarray
def rnd_test(func, *args, **kwargs):
a = func(*args, **kwargs)
assert type(a._meta) == expect
assert_eq(a, a) # Check that _meta and computed arrays match types
with config.set({"array.backend": backend}):
rs = da.random.RandomState(RandomState=rs)
rnd_test(rs.beta, 1, 2, size=5, chunks=3)
rnd_test(rs.binomial, 10, 0.5, size=5, chunks=3)
rnd_test(rs.chisquare, 1, size=5, chunks=3)
rnd_test(rs.exponential, 1, size=5, chunks=3)
rnd_test(rs.f, 1, 2, size=5, chunks=3)
rnd_test(rs.gamma, 5, 1, size=5, chunks=3)
rnd_test(rs.geometric, 1, size=5, chunks=3)
rnd_test(rs.gumbel, 1, size=5, chunks=3)
rnd_test(rs.hypergeometric, 1, 2, 3, size=5, chunks=3)
rnd_test(rs.laplace, size=5, chunks=3)
rnd_test(rs.logistic, size=5, chunks=3)
rnd_test(rs.lognormal, size=5, chunks=3)
rnd_test(rs.logseries, 0.5, size=5, chunks=3)
# No RandomState for multinomial in CuPy
# rnd_test(rs.multinomial, 20, [1 / 6.] * 6, size=5, chunks=3)
rnd_test(rs.negative_binomial, 5, 0.5, size=5, chunks=3)
rnd_test(rs.noncentral_chisquare, 2, 2, size=5, chunks=3)
rnd_test(rs.noncentral_f, 2, 2, 3, size=5, chunks=3)
rnd_test(rs.normal, 2, 2, size=5, chunks=3)
rnd_test(rs.pareto, 1, size=5, chunks=3)
rnd_test(rs.poisson, size=5, chunks=3)
rnd_test(rs.power, 1, size=5, chunks=3)
rnd_test(rs.rayleigh, size=5, chunks=3)
rnd_test(rs.random_sample, size=5, chunks=3)
rnd_test(rs.triangular, 1, 2, 3, size=5, chunks=3)
rnd_test(rs.uniform, size=5, chunks=3)
rnd_test(rs.vonmises, 2, 3, size=5, chunks=3)
rnd_test(rs.wald, 1, 2, size=5, chunks=3)
rnd_test(rs.weibull, 2, size=5, chunks=3)
rnd_test(rs.zipf, 2, size=5, chunks=3)
rnd_test(rs.standard_cauchy, size=5, chunks=3)
rnd_test(rs.standard_exponential, size=5, chunks=3)
rnd_test(rs.standard_gamma, 2, size=5, chunks=3)
rnd_test(rs.standard_normal, size=5, chunks=3)
rnd_test(rs.standard_t, 2, size=5, chunks=3)
@pytest.mark.parametrize("backend", ["cupy", "numpy"])
def test_random_all_direct_calls(backend):
if backend == "cupy":
expect = cupy.ndarray
else:
expect = np.ndarray
def rnd_test(func, *args, **kwargs):
a = func(*args, **kwargs)
assert type(a._meta) == expect
assert_eq(a, a) # Check that _meta and computed arrays match types
with config.set({"array.backend": backend}):
# check calling random functions directly
rnd_test(da.random.beta, 1, 2, size=5, chunks=3)
rnd_test(da.random.binomial, 10, 0.5, size=5, chunks=3)
rnd_test(da.random.chisquare, 1, size=5, chunks=3)
rnd_test(da.random.exponential, 1, size=5, chunks=3)
rnd_test(da.random.f, 1, 2, size=5, chunks=3)
rnd_test(da.random.gamma, 5, 1, size=5, chunks=3)
rnd_test(da.random.geometric, 1, size=5, chunks=3)
rnd_test(da.random.gumbel, 1, size=5, chunks=3)
rnd_test(da.random.hypergeometric, 1, 2, 3, size=5, chunks=3)
rnd_test(da.random.laplace, size=5, chunks=3)
rnd_test(da.random.logistic, size=5, chunks=3)
rnd_test(da.random.lognormal, size=5, chunks=3)
rnd_test(da.random.logseries, 0.5, size=5, chunks=3)
# No RandomState for multinomial in CuPy
if backend != "cupy":
rnd_test(da.random.multinomial, 20, [1 / 6.0] * 6, size=5, chunks=3)
rnd_test(da.random.negative_binomial, 5, 0.5, size=5, chunks=3)
rnd_test(da.random.noncentral_chisquare, 2, 2, size=5, chunks=3)
rnd_test(da.random.noncentral_f, 2, 2, 3, size=5, chunks=3)
rnd_test(da.random.normal, 2, 2, size=5, chunks=3)
rnd_test(da.random.pareto, 1, size=5, chunks=3)
rnd_test(da.random.poisson, size=5, chunks=3)
rnd_test(da.random.power, 1, size=5, chunks=3)
rnd_test(da.random.rayleigh, size=5, chunks=3)
rnd_test(da.random.randint, low=10, size=5, chunks=3)
rnd_test(da.random.random, size=5, chunks=3)
rnd_test(da.random.random_sample, size=5, chunks=3)
rnd_test(da.random.triangular, 1, 2, 3, size=5, chunks=3)
rnd_test(da.random.uniform, size=5, chunks=3)
rnd_test(da.random.vonmises, 2, 3, size=5, chunks=3)
rnd_test(da.random.wald, 1, 2, size=5, chunks=3)
rnd_test(da.random.weibull, 2, size=5, chunks=3)
rnd_test(da.random.zipf, 2, size=5, chunks=3)
rnd_test(da.random.standard_cauchy, size=5, chunks=3)
rnd_test(da.random.standard_exponential, size=5, chunks=3)
rnd_test(da.random.standard_gamma, 2, size=5, chunks=3)
rnd_test(da.random.standard_normal, size=5, chunks=3)
rnd_test(da.random.standard_t, 2, size=5, chunks=3)
@pytest.mark.parametrize("backend", ["cupy", "numpy"])
@pytest.mark.parametrize("gen", [None, cupy.random.default_rng, np.random.default_rng])
@pytest.mark.parametrize("shape", [(2), (2, 3), (2, 3, 4), (2, 3, 4, 2)], ids=type)
def test_random_all_Generator(backend, gen, shape):
# Generator argument takes priority over backend
if gen == cupy.random.default_rng:
expect = cupy.ndarray
elif gen == np.random.default_rng:
expect = np.ndarray
elif backend == "cupy":
expect = cupy.ndarray
else:
expect = np.ndarray
def rnd_test(func, *args, **kwargs):
a = func(*args, **kwargs)
assert type(a._meta) == expect
assert_eq(a, a) # Check that _meta and computed arrays match types
with config.set({"array.backend": backend}):
generator = gen(5) if gen else None
rng = da.random.default_rng(generator)
rnd_test(rng.beta, 1, 2, size=shape, chunks=3)
rnd_test(rng.binomial, 10, 0.5, size=shape, chunks=3)
rnd_test(rng.chisquare, 1, size=shape, chunks=3)
rnd_test(rng.exponential, 1, size=shape, chunks=3)
rnd_test(rng.f, 1, 2, size=shape, chunks=3)
rnd_test(rng.gamma, 5, 1, size=shape, chunks=3)
rnd_test(rng.geometric, 1, size=shape, chunks=3)
rnd_test(rng.hypergeometric, 1, 2, 3, size=shape, chunks=3)
rnd_test(rng.integers, 1, high=10, size=shape, chunks=3)
rnd_test(rng.logseries, 0.5, size=shape, chunks=3)
rnd_test(rng.poisson, 1, size=shape, chunks=3)
rnd_test(rng.power, 1, size=shape, chunks=3)
rnd_test(rng.random, size=shape, chunks=3)
rnd_test(rng.standard_exponential, size=shape, chunks=3)
rnd_test(rng.standard_gamma, 2, size=shape, chunks=3)
rnd_test(rng.standard_normal, size=shape, chunks=3)
@pytest.mark.parametrize("backend", ["cupy", "numpy"])
def test_random_Generator_processes(backend):
with config.set({"array.backend": backend}):
# Check that matching seeds produce consistent results
# with scheduler="processes"
state = da.random.default_rng(5)
x = state.standard_normal(size=(2, 3), chunks=3)
state = da.random.default_rng(5)
y = state.standard_normal(size=(2, 3), chunks=3)
assert_eq(x, y, scheduler="processes")
def test_cupy_unsupported():
with config.set({"array.backend": "cupy"}):
# permutation supported for np-backed BitGenerator
x = da.arange(12, chunks=3)
da.random.default_rng(np.random.PCG64()).permutation(x).compute()
# permutation not supported for default cupy BitGenerator
with pytest.raises(NotImplementedError):
da.random.default_rng().permutation(x).compute()
# choice not supported for cupy-backed Generator
with pytest.raises(NotImplementedError):
da.random.default_rng().choice(10).compute()
@pytest.mark.parametrize("shape", [(2, 3), (2, 3, 4), (2, 3, 4, 2)])
def test_random_shapes(shape):
rs = da.random.RandomState(RandomState=cupy.random.RandomState)
x = rs.poisson(size=shape, chunks=3)
assert type(x._meta) == cupy.ndarray
assert_eq(x, x) # Check that _meta and computed arrays match types
assert x._meta.shape == (0,) * len(shape)
assert x.shape == shape
rng = da.random.default_rng(cupy.random.default_rng())
x = rng.poisson(1.0, size=shape, chunks=3)
assert type(x._meta) == cupy.ndarray
assert_eq(x, x) # Check that _meta and computed arrays match types
assert x._meta.shape == (0,) * len(shape)
assert x.shape == shape