| 482 | get_device_count("xpu") <= 1, reason="xpu counts need > 1", |
| 483 | ) |
| 484 | def test_PoissonRNG(): |
| 485 | m1 = RNG(seed=111, device="xpu0") |
| 486 | m2 = RNG(seed=111, device="xpu1") |
| 487 | m3 = RNG(seed=222, device="xpu0") |
| 488 | lam = Tensor([[2, 3, 4], [9, 10, 11]], dtype=np.float32) |
| 489 | out1 = m1.poisson(lam.to("xpu0"), size=(100,)) |
| 490 | out2 = m2.poisson(lam.to("xpu1"), size=(100,)) |
| 491 | out3 = m3.poisson(lam.to("xpu0"), size=(100,)) |
| 492 | |
| 493 | np.testing.assert_allclose(out1.numpy(), out2.numpy(), atol=1e-6) |
| 494 | assert out1.device == "xpu0" and out2.device == "xpu1" |
| 495 | assert not (out1.numpy() == out3.numpy()).all() |
| 496 | |
| 497 | out = m1.poisson(lam.to("xpu0"), size=(20, 30)) |
| 498 | out_shp = out.shape |
| 499 | expected_shape = (20, 30) + lam._tuple_shape |
| 500 | if isinstance(out_shp, tuple): |
| 501 | assert out_shp == expected_shape |
| 502 | else: |
| 503 | assert all(out.shape.numpy() == np.array(expected_shape)) |
| 504 | lam = lam.numpy() |
| 505 | |
| 506 | assert (np.abs(out.mean(axis=(0, 1)).numpy() - lam) / np.sqrt(lam)).mean() < 0.1 |
| 507 | assert np.abs(np.std(out.numpy(), axis=(0, 1)) - np.sqrt(lam)).mean() < 0.1 |
| 508 | |
| 509 | |
| 510 | @pytest.mark.skipif( |