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