| 78 | self.assertTrue(mx.all(y == 1)) |
| 79 | |
| 80 | def test_send_recv(self): |
| 81 | world = mx.distributed.init() |
| 82 | dtypes = [ |
| 83 | mx.int8, |
| 84 | mx.uint8, |
| 85 | mx.int16, |
| 86 | mx.uint16, |
| 87 | mx.int32, |
| 88 | mx.uint32, |
| 89 | mx.float32, |
| 90 | mx.float16, |
| 91 | mx.bfloat16, |
| 92 | mx.complex64, |
| 93 | ] |
| 94 | sizes = [ |
| 95 | (7,), |
| 96 | (10,), |
| 97 | (1024,), |
| 98 | (1024, 1024), |
| 99 | ] |
| 100 | key = mx.random.key(0) |
| 101 | right = (world.rank() + 1) % world.size() |
| 102 | left = (world.rank() + world.size() - 1) % world.size() |
| 103 | for dt in dtypes: |
| 104 | for sh in sizes: |
| 105 | x = ( |
| 106 | mx.random.uniform(shape=(world.size(),) + sh, key=key) * 10 |
| 107 | ).astype(dt) |
| 108 | if world.rank() % 2 == 0: |
| 109 | y = mx.distributed.send(x[world.rank()], right) |
| 110 | z = mx.distributed.recv_like(y, left) |
| 111 | mx.eval(y, z) |
| 112 | else: |
| 113 | z = mx.distributed.recv_like(x[world.rank()], left) |
| 114 | y = mx.distributed.send(x[world.rank()], right) |
| 115 | mx.eval(z, y) |
| 116 | self.assertTrue(mx.all(y == x[world.rank()])) |
| 117 | self.assertTrue(mx.all(z == x[left])) |
| 118 | |
| 119 | def test_all_gather_vjp(self): |
| 120 | def fun(x): |