def fn():
rstate = random_initialize(12345678)
r0, v0 = random_uint32(rstate, (2, 2))
r1, v1 = random_uint32(r0, (2, 2))
r2, v2 = random_uint32(r1, (1,))
r3, v3 = random_uint32(r2, ())
return r3, v0, v1, v2, v3
def test_with_random(a, b):
rstate = random_initialize(123)
rstate, v0 = random_uint32(rstate, (2, 3))
_, v1 = random_uint32(rstate, (2, 3))
v0 = array_cast(v0, xtype.f64)
v1 = array_cast(v1, xtype.f64)
return a * np.sum(v0).item() + b * np.sum(v1).item()
def test_random():
rstate = random_initialize(1234)
r0, v0 = random_uint32(rstate, (2, 3))
r1, v1 = random_uint32(rstate, (2, 3))
assert isinstance(rstate, np.random.RandomState)
assert isinstance(r0, np.random.RandomState)
assert isinstance(r1, np.random.RandomState)
assert isinstance(v0, np.ndarray)
assert isinstance(v1, np.ndarray)
assert v0.dtype == "uint32"
assert v1.dtype == "uint32"
assert v0.shape == (2, 3)
assert v1.shape == (2, 3)
assert not np.allclose(v0, v1)
def only_compute(rstate, x):
"""Use rstate but return only computed value."""
_, _val = random_uint32(rstate, (2, 3))
_val = P.array_cast(_val, xtype.i64)
return x * np.sum(_val).item()
def rstate_and_compute(rstate, x):
"""Uses and returns a rstate with a computed value."""
_rs, _val = random_uint32(rstate, (2, 3))
_val = P.array_cast(_val, xtype.i64)
return _rs, x * np.sum(_val).item()
def gen_scalar(rng):
return random_uint32(rng, ())
def gen_1(rng):
return random_uint32(rng, (1,))
def gen_2_2(rng):
return random_uint32(rng, (2, 2))
def f() -> RandomStateWrapper:
rstate: RandomStateWrapper = random_initialize(10)
r0, _ = random_uint32(rstate, ())
return r0
