| 447 | get_device_count("xpu") <= 1, reason="xpu counts need > 1", |
| 448 | ) |
| 449 | def test_BetaRNG(): |
| 450 | m1 = RNG(seed=111, device="xpu0") |
| 451 | m2 = RNG(seed=111, device="xpu1") |
| 452 | m3 = RNG(seed=222, device="xpu0") |
| 453 | out1 = m1.beta(2, 1, size=(100,)) |
| 454 | out1_ = m1.uniform(size=(100,)) |
| 455 | out2 = m2.beta(2, 1, size=(100,)) |
| 456 | out3 = m3.beta(2, 1, size=(100,)) |
| 457 | |
| 458 | np.testing.assert_allclose(out1.numpy(), out2.numpy(), atol=1e-6) |
| 459 | assert out1.device == "xpu0" and out2.device == "xpu1" |
| 460 | assert not (out1.numpy() == out3.numpy()).all() |
| 461 | assert not (out1.numpy() == out1_.numpy()).all() |
| 462 | |
| 463 | alpha = Tensor([[2, 3, 4], [9, 10, 11]], dtype=np.float32, device="xpu0") |
| 464 | beta = Tensor([0.5, 1, 1.5], dtype=np.float32, device="xpu0") |
| 465 | expected_mean = (alpha / (alpha + beta)).numpy() |
| 466 | expected_std = ( |
| 467 | F.sqrt(alpha * beta / (F.pow(alpha + beta, 2) * (alpha + beta + 1))) |
| 468 | ).numpy() |
| 469 | out = m1.beta(alpha=alpha, beta=beta, size=(20, 30)) |
| 470 | out_shp = out.shape |
| 471 | if isinstance(out_shp, tuple): |
| 472 | assert out_shp == (20, 30, 2, 3) |
| 473 | else: |
| 474 | assert all(out.shape.numpy() == np.array([20, 30, 2, 3])) |
| 475 | assert ( |
| 476 | np.abs(out.mean(axis=(0, 1)).numpy() - expected_mean) / expected_std |
| 477 | ).mean() < 0.1 |
| 478 | assert (np.abs(np.std(out.numpy(), axis=(0, 1)) - expected_std)).mean() < 0.1 |
| 479 | |
| 480 | |
| 481 | @pytest.mark.skipif( |