| 34 | get_device_count("xpu") <= 2, reason="xpu counts need > 2", |
| 35 | ) |
| 36 | def test_gaussian_op(): |
| 37 | set_global_seed(1024) |
| 38 | shape = ( |
| 39 | 8, |
| 40 | 9, |
| 41 | 11, |
| 42 | 12, |
| 43 | ) |
| 44 | shape = Tensor(shape, dtype="int32") |
| 45 | op = GaussianRNG(seed=get_global_rng_seed(), mean=1.0, std=3.0, dtype="float32") |
| 46 | (output,) = apply(op, shape) |
| 47 | assert np.fabs(output.numpy().mean() - 1.0) < 1e-1 |
| 48 | assert np.fabs(np.sqrt(output.numpy().var()) - 3.0) < 1e-1 |
| 49 | assert str(output.device) == str(CompNode("xpux")) |
| 50 | assert output.dtype == np.float32 |
| 51 | |
| 52 | cn = CompNode("xpu2") |
| 53 | seed = 233333 |
| 54 | h = new_rng_handle(cn, seed) |
| 55 | op = GaussianRNG(seed=seed, mean=3.0, std=1.0, dtype="float32", handle=h) |
| 56 | (output,) = apply(op, shape) |
| 57 | delete_rng_handle(h) |
| 58 | assert np.fabs(output.numpy().mean() - 3.0) < 1e-1 |
| 59 | assert np.fabs(np.sqrt(output.numpy().var()) - 1.0) < 1e-1 |
| 60 | assert str(output.device) == str(cn) |
| 61 | assert output.dtype == np.float32 |
| 62 | |
| 63 | |
| 64 | @pytest.mark.skipif( |