| 94 | get_device_count("xpu") <= 2, reason="xpu counts need > 2", |
| 95 | ) |
| 96 | def test_gamma_op(): |
| 97 | set_global_seed(1024) |
| 98 | _shape, _scale = 2, 0.8 |
| 99 | _expected_mean, _expected_std = _shape * _scale, np.sqrt(_shape) * _scale |
| 100 | |
| 101 | shape = F.full([8, 9, 11, 12], value=_shape, dtype="float32") |
| 102 | scale = F.full([8, 9, 11, 12], value=_scale, dtype="float32") |
| 103 | op = GammaRNG(seed=get_global_rng_seed(), handle=0) |
| 104 | (output,) = apply(op, shape, scale) |
| 105 | assert np.fabs(output.numpy().mean() - _expected_mean) < 1e-1 |
| 106 | assert np.fabs(np.sqrt(output.numpy().var()) - _expected_std) < 1e-1 |
| 107 | assert str(output.device) == str(CompNode("xpux")) |
| 108 | |
| 109 | cn = CompNode("xpu2") |
| 110 | seed = 233333 |
| 111 | h = new_rng_handle(cn, seed) |
| 112 | shape = F.full([8, 9, 11, 12], value=_shape, dtype="float32", device="xpu2") |
| 113 | scale = F.full([8, 9, 11, 12], value=_scale, dtype="float32", device="xpu2") |
| 114 | op = GammaRNG(seed=seed, handle=h) |
| 115 | (output,) = apply(op, shape, scale) |
| 116 | delete_rng_handle(h) |
| 117 | assert np.fabs(output.numpy().mean() - _expected_mean) < 1e-1 |
| 118 | assert np.fabs(np.sqrt(output.numpy().var()) - _expected_std) < 1e-1 |
| 119 | assert str(output.device) == str(cn) |
| 120 | |
| 121 | |
| 122 | @pytest.mark.skipif( |