(self)
| 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]]) |
| 247 | out = mx.vmap(lambda x: mx.argmin(x))(a) |
| 248 | expected = mx.array([0, 2]) |
| 249 | self.assertTrue(mx.array_equal(out, expected)) |
| 250 | |
| 251 | out = mx.vmap(lambda x: mx.argmax(x))(a) |
| 252 | expected = mx.array([2, 1]) |
| 253 | self.assertTrue(mx.array_equal(out, expected)) |
| 254 | |
| 255 | def test_vmap_mean(self): |
| 256 | a = mx.arange(8).reshape(2, 4) |