| 710 | get_device_count("xpu") <= 1, reason="xpu counts need > 1", |
| 711 | ) |
| 712 | def test_ShuffleRNG(): |
| 713 | g = [] |
| 714 | |
| 715 | def cb(grad): |
| 716 | g.append(grad) |
| 717 | |
| 718 | n, m = 6, 3 |
| 719 | arr = np.arange(n * m) |
| 720 | out0 = Tensor(arr, dtype="float32") |
| 721 | with Grad() as grad: |
| 722 | grad.wrt(out0, callback=cb) |
| 723 | random.shuffle(out0) |
| 724 | grad(out0, F.ones_like(out0)) |
| 725 | m1 = RNG(seed=111, device="xpu0") |
| 726 | m2 = RNG(seed=111, device="xpu1") |
| 727 | m3 = RNG(seed=222, device="xpu0") |
| 728 | out1 = Tensor(arr, dtype="float32", device="xpu0") |
| 729 | out2 = Tensor(arr, dtype="float32", device="xpu1") |
| 730 | out3 = Tensor(arr, dtype="float32", device="xpu0") |
| 731 | m1.shuffle(out1) |
| 732 | m2.shuffle(out2) |
| 733 | m3.shuffle(out3) |
| 734 | |
| 735 | np.testing.assert_allclose(out1.numpy(), out2.numpy(), atol=1e-6) |
| 736 | assert out1.device == "xpu0" and out2.device == "xpu1" |
| 737 | assert not (out1.numpy() == out3.numpy()).all() |
| 738 | |
| 739 | out = Tensor(arr, dtype="float32").reshape(n, m) |
| 740 | m1.shuffle(out) |
| 741 | |
| 742 | out_shp = out.shape |
| 743 | if isinstance(out_shp, tuple): |
| 744 | assert out_shp == (n, m) |
| 745 | else: |
| 746 | assert all(out.shape.numpy() == np.array([n, m])) |
| 747 | |
| 748 | |
| 749 | @pytest.mark.skipif( |