| 535 | self.assertTrue(mx.allclose(out, expected)) |
| 536 | |
| 537 | def test_vmap_const_func(self): |
| 538 | a = mx.random.uniform(shape=(2, 3, 4)) |
| 539 | b = mx.random.uniform(shape=(4, 3)) |
| 540 | |
| 541 | def const_func(a, b): |
| 542 | return mx.array(2) |
| 543 | |
| 544 | out = mx.vmap(const_func, in_axes=(0, None))(a, b) |
| 545 | self.assertTrue(mx.array_equal(mx.full((2,), 2), out)) |
| 546 | out = mx.vmap(const_func, in_axes=(None, 0))(a, b) |
| 547 | self.assertTrue(mx.array_equal(mx.full((4,), 2), out)) |
| 548 | out = mx.vmap(const_func, in_axes=(1, 1))(a, b) |
| 549 | self.assertTrue(mx.array_equal(mx.full((3,), 2), out)) |
| 550 | |
| 551 | with self.assertRaises(ValueError): |
| 552 | out = mx.vmap(const_func, in_axes=(None, None))(a, b) |
| 553 | |
| 554 | with self.assertRaises(ValueError): |
| 555 | out = mx.vmap(const_func, in_axes=(0, 0))(a, b) |
| 556 | |
| 557 | def test_vmap_concatenate(self): |
| 558 | x = mx.random.uniform(shape=(2, 2, 2)) |