| 331 | get_device_count("xpu") <= 2, reason="xpu counts need > 2", |
| 332 | ) |
| 333 | def test_exponential_op(): |
| 334 | set_global_seed(1024) |
| 335 | rate = F.full([8, 9, 11, 12], value=2, dtype="float32") |
| 336 | op = ExponentialRNG(seed=get_global_rng_seed()) |
| 337 | (output,) = apply(op, rate) |
| 338 | expected_mean = 1.0 / rate |
| 339 | expected_std = np.sqrt(1.0 / rate ** 2) |
| 340 | assert np.fabs(output.numpy().mean() - expected_mean).max() < 1e-1 |
| 341 | assert np.fabs(np.sqrt(output.numpy().var()) - expected_std).max() < 1e-1 |
| 342 | assert str(output.device) == str(CompNode("xpux")) |
| 343 | |
| 344 | cn = CompNode("xpu2") |
| 345 | seed = 233333 |
| 346 | h = new_rng_handle(cn, seed) |
| 347 | rate = F.full([8, 9, 11, 12], value=2, dtype="float32", device=cn) |
| 348 | op = ExponentialRNG(seed=seed, handle=h) |
| 349 | (output,) = apply(op, rate) |
| 350 | delete_rng_handle(h) |
| 351 | assert np.fabs(output.numpy().mean() - expected_mean).max() < 1e-1 |
| 352 | assert np.fabs(np.sqrt(output.numpy().var()) - expected_std).max() < 1e-1 |
| 353 | assert str(output.device) == str(cn) |
| 354 | |
| 355 | |
| 356 | @pytest.mark.skipif( |