(self)
| 253 | self.assertTrue(mx.array_equal(out, expected)) |
| 254 | |
| 255 | def test_vmap_mean(self): |
| 256 | a = mx.arange(8).reshape(2, 4) |
| 257 | out = mx.vmap(mx.mean)(a) |
| 258 | expected = mx.mean(a, axis=1) |
| 259 | self.assertTrue(mx.allclose(out, expected)) |
| 260 | |
| 261 | a = mx.arange(16).reshape(2, 2, 4) |
| 262 | out = mx.vmap(mx.vmap(mx.mean))(a) |
| 263 | expected = mx.mean(a, axis=2) |
| 264 | self.assertTrue(mx.allclose(out, expected)) |
| 265 | |
| 266 | def test_mismatch_input_sizes(self): |
| 267 | a = mx.ones((10, 1)) |