| 414 | get_device_count("xpu") <= 1, reason="xpu counts need > 1", |
| 415 | ) |
| 416 | def test_GammaRNG(): |
| 417 | m1 = RNG(seed=111, device="xpu0") |
| 418 | m2 = RNG(seed=111, device="xpu1") |
| 419 | m3 = RNG(seed=222, device="xpu0") |
| 420 | out1 = m1.gamma(2, size=(100,)) |
| 421 | out1_ = m1.uniform(size=(100,)) |
| 422 | out2 = m2.gamma(2, size=(100,)) |
| 423 | out3 = m3.gamma(2, size=(100,)) |
| 424 | |
| 425 | np.testing.assert_allclose(out1.numpy(), out2.numpy(), atol=1e-6) |
| 426 | assert out1.device == "xpu0" and out2.device == "xpu1" |
| 427 | assert not (out1.numpy() == out3.numpy()).all() |
| 428 | assert not (out1.numpy() == out1_.numpy()).all() |
| 429 | |
| 430 | shape = Tensor([[2, 3, 4], [9, 10, 11]], dtype=np.float32, device="xpu0") |
| 431 | scale = Tensor([0.5, 1, 1.5], dtype=np.float32, device="xpu0") |
| 432 | expected_mean = (shape * scale).numpy() |
| 433 | expected_std = (F.sqrt(shape) * scale).numpy() |
| 434 | out = m1.gamma(shape=shape, scale=scale, size=(20, 30, 40)) |
| 435 | out_shp = out.shape |
| 436 | if isinstance(out_shp, tuple): |
| 437 | assert out_shp == (20, 30, 40, 2, 3) |
| 438 | else: |
| 439 | assert all(out.shape.numpy() == np.array([20, 30, 40, 2, 3])) |
| 440 | assert ( |
| 441 | np.abs(out.mean(axis=(0, 1)).numpy() - expected_mean) / expected_std |
| 442 | ).mean() < 0.1 |
| 443 | assert (np.abs(np.std(out.numpy(), axis=(0, 1)) - expected_std)).mean() < 0.1 |
| 444 | |
| 445 | |
| 446 | @pytest.mark.skipif( |