| 220 | self.assertTrue(mx.array_equal(out, expected)) |
| 221 | |
| 222 | def test_vmap_reduce(self): |
| 223 | a = mx.ones((5, 5), mx.int32) |
| 224 | out = mx.vmap(lambda x: x.sum())(a) |
| 225 | self.assertTrue(mx.array_equal(out, mx.full((5,), 5))) |
| 226 | |
| 227 | out = mx.vmap(lambda x: x.sum(keepdims=True))(a) |
| 228 | self.assertTrue(mx.array_equal(out, mx.full((5, 1), 5))) |
| 229 | |
| 230 | out = mx.vmap(lambda x: x.sum(axis=0))(a) |
| 231 | self.assertTrue(mx.array_equal(out, mx.full((5,), 5))) |
| 232 | |
| 233 | a = mx.ones((5, 3, 2), mx.int32) |
| 234 | out = mx.vmap(lambda x: x.sum(axis=(0, 1)))(a) |
| 235 | self.assertTrue(mx.array_equal(out, mx.full((5,), 6))) |
| 236 | |
| 237 | a = mx.ones((5, 3, 2), mx.int32) |
| 238 | out = mx.vmap(lambda x: x.sum(axis=(0, 1)), in_axes=(1,))(a) |
| 239 | self.assertTrue(mx.array_equal(out, mx.full((3,), 10))) |
| 240 | |
| 241 | a = mx.ones((5, 3, 2), mx.int32) |
| 242 | out = mx.vmap(lambda x: x.sum(axis=(0, 1)), in_axes=(2,))(a) |
| 243 | self.assertTrue(mx.array_equal(out, mx.full((2,), 15))) |
| 244 | |
| 245 | def test_vmap_argreduce(self): |
| 246 | a = mx.array([[1, 2, 3], [2, 3, 1]]) |