| 750 | get_device_count("xpu") <= 1, reason="xpu counts need > 1", |
| 751 | ) |
| 752 | def test_ExponentialRNG(): |
| 753 | m1 = RNG(seed=111, device="xpu0") |
| 754 | m2 = RNG(seed=111, device="xpu1") |
| 755 | m3 = RNG(seed=222, device="xpu0") |
| 756 | rate = Tensor([[2, 3, 4], [9, 10, 11]], dtype=np.float32) |
| 757 | out1 = m1.exponential(rate.to("xpu0"), size=(100,)) |
| 758 | out2 = m2.exponential(rate.to("xpu1"), size=(100,)) |
| 759 | out3 = m3.exponential(rate.to("xpu0"), size=(100,)) |
| 760 | |
| 761 | np.testing.assert_allclose(out1.numpy(), out2.numpy(), atol=1e-6) |
| 762 | assert out1.device == "xpu0" and out2.device == "xpu1" |
| 763 | assert not (out1.numpy() == out3.numpy()).all() |
| 764 | |
| 765 | out = m1.exponential(rate.to("xpu0"), size=(20, 30)) |
| 766 | out_shp = out.shape |
| 767 | expected_shape = (20, 30) + rate._tuple_shape |
| 768 | if isinstance(out_shp, tuple): |
| 769 | assert out_shp == expected_shape |
| 770 | else: |
| 771 | assert all(out.shape.numpy() == np.array(expected_shape)) |
| 772 | rate = rate.numpy() |
| 773 | |
| 774 | expected_mean = 1.0 / rate |
| 775 | expected_std = np.sqrt(1.0 / (rate * rate)) |
| 776 | assert ( |
| 777 | np.abs(out.mean(axis=(0, 1)).numpy() - expected_mean) / expected_std |
| 778 | ).mean() < 0.1 |
| 779 | assert np.abs(np.std(out.numpy(), axis=(0, 1)) - expected_std).mean() < 0.1 |
| 780 | |
| 781 | |
| 782 | def test_seed(): |