def test_PoissonRNG():
m1 = RNG(seed=111, device="xpu0")
m2 = RNG(seed=111, device="xpu1")
m3 = RNG(seed=222, device="xpu0")
lam = Tensor([[2, 3, 4], [9, 10, 11]], dtype=np.float32)
out1 = m1.poisson(lam.to("xpu0"), size=(100, ))
out2 = m2.poisson(lam.to("xpu1"), size=(100, ))
out3 = m3.poisson(lam.to("xpu0"), size=(100, ))
np.testing.assert_allclose(out1.numpy(), out2.numpy(), atol=1e-6)
assert out1.device == "xpu0" and out2.device == "xpu1"
assert not (out1.numpy() == out3.numpy()).all()
out = m1.poisson(lam.to("xpu0"), size=(20, 30))
out_shp = out.shape
expected_shape = (20, 30) + lam._tuple_shape
if isinstance(out_shp, tuple):
assert out_shp == expected_shape
else:
assert all(out.shape.numpy() == np.array(expected_shape))
lam = lam.numpy()
assert (np.abs(out.mean(axis=(0, 1)).numpy() - lam) /
np.sqrt(lam)).mean() < 0.1
assert np.abs(np.std(out.numpy(), axis=(0, 1)) - np.sqrt(lam)).mean() < 0.1
def test_ShuffleRNG():
g = []
def cb(grad):
g.append(grad)
n, m = 6, 3
arr = np.arange(n * m)
out0 = Tensor(arr, dtype="float32")
with Grad() as grad:
grad.wrt(out0, callback=cb)
random.shuffle(out0)
grad(out0, F.ones_like(out0))
m1 = RNG(seed=111, device="xpu0")
m2 = RNG(seed=111, device="xpu1")
m3 = RNG(seed=222, device="xpu0")
out1 = Tensor(arr, dtype="float32", device="xpu0")
out2 = Tensor(arr, dtype="float32", device="xpu1")
out3 = Tensor(arr, dtype="float32", device="xpu0")
m1.shuffle(out1)
m2.shuffle(out2)
m3.shuffle(out3)
np.testing.assert_allclose(out1.numpy(), out2.numpy(), atol=1e-6)
assert out1.device == "xpu0" and out2.device == "xpu1"
assert not (out1.numpy() == out3.numpy()).all()
out = Tensor(arr, dtype="float32").reshape(n, m)
m1.shuffle(out)
out_shp = out.shape
if isinstance(out_shp, tuple):
assert out_shp == (n, m)
else:
assert all(out.shape.numpy() == np.array([n, m]))
def test_PermutationRNG(symbolic):
m1 = RNG(seed=111, device="xpu0")
m2 = RNG(seed=111, device="xpu1")
m3 = RNG(seed=222, device="xpu0")
out1 = m1.permutation(1000)
out1_ = m1.uniform(size=(1000, ))
out2 = m2.permutation(1000)
out3 = m3.permutation(1000)
np.testing.assert_allclose(out1.numpy(), out2.numpy(), atol=1e-6)
assert out1.device == "xpu0" and out2.device == "xpu1"
assert not (out1.numpy() == out3.numpy()).all()
assert not (out1.numpy() == out1_.numpy()).all()
out = m1.permutation(1000)
out_shp = out.shape
if isinstance(out_shp, tuple):
assert out_shp == (1000, )
else:
assert all(out.shape.numpy() == np.array([1000]))
def sum_result(res, fun):
return sum(
[1 if i == v else 0 for i, v in enumerate(fun(res.numpy()))])
assert sum_result(out, lambda x: x) < 500
assert sum_result(out, np.sort) == 1000
def func():
out = m1.permutation(Tensor(7))
out_shp = out.shape
if isinstance(out_shp, tuple):
assert out_shp == (1, )
else:
assert all(out.shape.numpy() == np.array([1]))
n, m = 6, 3
out = m1.permutation(
Tensor(np.arange(n * m), dtype="float32").reshape(n, m))
out_shp = out.shape
if isinstance(out_shp, tuple):
assert out_shp == (n, m)
else:
assert all(out.shape.numpy() == np.array([n, m]))
func = trace(symbolic=symbolic)(func)
func()
def test_BetaRNG():
m1 = RNG(seed=111, device="xpu0")
m2 = RNG(seed=111, device="xpu1")
m3 = RNG(seed=222, device="xpu0")
out1 = m1.beta(2, 1, size=(100, ))
out1_ = m1.uniform(size=(100, ))
out2 = m2.beta(2, 1, size=(100, ))
out3 = m3.beta(2, 1, size=(100, ))
np.testing.assert_allclose(out1.numpy(), out2.numpy(), atol=1e-6)
assert out1.device == "xpu0" and out2.device == "xpu1"
assert not (out1.numpy() == out3.numpy()).all()
assert not (out1.numpy() == out1_.numpy()).all()
alpha = Tensor([[2, 3, 4], [9, 10, 11]], dtype=np.float32, device="xpu0")
beta = Tensor([0.5, 1, 1.5], dtype=np.float32, device="xpu0")
expected_mean = (alpha / (alpha + beta)).numpy()
expected_std = (F.sqrt(alpha * beta / (F.pow(alpha + beta, 2) *
(alpha + beta + 1)))).numpy()
out = m1.beta(alpha=alpha, beta=beta, size=(20, 30))
out_shp = out.shape
if isinstance(out_shp, tuple):
assert out_shp == (20, 30, 2, 3)
else:
assert all(out.shape.numpy() == np.array([20, 30, 2, 3]))
assert (np.abs(out.mean(axis=(0, 1)).numpy() - expected_mean) /
expected_std).mean() < 0.1
assert (np.abs(np.std(out.numpy(), axis=(0, 1)) -
expected_std)).mean() < 0.1
def test_GammaRNG():
m1 = RNG(seed=111, device="xpu0")
m2 = RNG(seed=111, device="xpu1")
m3 = RNG(seed=222, device="xpu0")
out1 = m1.gamma(2, size=(100, ))
out1_ = m1.uniform(size=(100, ))
out2 = m2.gamma(2, size=(100, ))
out3 = m3.gamma(2, size=(100, ))
np.testing.assert_allclose(out1.numpy(), out2.numpy(), atol=1e-6)
assert out1.device == "xpu0" and out2.device == "xpu1"
assert not (out1.numpy() == out3.numpy()).all()
assert not (out1.numpy() == out1_.numpy()).all()
shape = Tensor([[2, 3, 4], [9, 10, 11]], dtype=np.float32, device="xpu0")
scale = Tensor([0.5, 1, 1.5], dtype=np.float32, device="xpu0")
expected_mean = (shape * scale).numpy()
expected_std = (F.sqrt(shape) * scale).numpy()
out = m1.gamma(shape=shape, scale=scale, size=(20, 30, 40))
out_shp = out.shape
if isinstance(out_shp, tuple):
assert out_shp == (20, 30, 40, 2, 3)
else:
assert all(out.shape.numpy() == np.array([20, 30, 40, 2, 3]))
assert (np.abs(out.mean(axis=(0, 1)).numpy() - expected_mean) /
expected_std).mean() < 0.1
assert (np.abs(np.std(out.numpy(), axis=(0, 1)) -
expected_std)).mean() < 0.1
def test_NormalRNG():
m1 = RNG(seed=111, device="xpu0")
m2 = RNG(seed=111, device="xpu1")
m3 = RNG(seed=222, device="xpu0")
out1 = m1.normal(size=(100, ))
out1_ = m1.uniform(size=(100, ))
out2 = m2.normal(size=(100, ))
out3 = m3.normal(size=(100, ))
np.testing.assert_allclose(out1.numpy(), out2.numpy(), atol=1e-6)
assert out1.device == "xpu0" and out2.device == "xpu1"
assert not (out1.numpy() == out3.numpy()).all()
assert not (out1.numpy() == out1_.numpy()).all()
mean = -1
std = 2
out = m1.normal(mean=mean, std=std, size=(20, 30, 40))
out_shp = out.shape
if isinstance(out_shp, tuple):
assert out_shp == (20, 30, 40)
else:
assert all(out.shape.numpy() == np.array([20, 30, 40]))
assert np.abs(out.mean().numpy() - mean) / std < 0.1
assert np.abs(np.std(out.numpy()) - std) < 0.1
def test_UniformRNG():
m1 = RNG(seed=111, device="xpu0")
m2 = RNG(seed=111, device="xpu1")
m3 = RNG(seed=222, device="xpu0")
out1 = m1.uniform(size=(100, ))
out1_ = m1.uniform(size=(100, ))
out2 = m2.uniform(size=(100, ))
out3 = m3.uniform(size=(100, ))
np.testing.assert_allclose(out1.numpy(), out2.numpy(), atol=1e-6)
assert out1.device == "xpu0" and out2.device == "xpu1"
assert not (out1.numpy() == out3.numpy()).all()
assert not (out1.numpy() == out1_.numpy()).all()
low = -234
high = 123
out = m1.uniform(low=low, high=high, size=(20, 30, 40))
out_shp = out.shape
if isinstance(out_shp, tuple):
assert out_shp == (20, 30, 40)
else:
assert all(out.shape.numpy() == np.array([20, 30, 40]))
assert np.abs(out.mean().numpy() - ((low + high) / 2)) / (high - low) < 0.1