Method fun

python/tests/test_vmap.py:581–582

(a, idx)

        idx = mx.zeros((2, 4, 1), mx.int32)

        def fun(a, idx):
            return mx.take_along_axis(a, idx, axis=0)

        out = mx.vmap(fun, in_axes=(0, 1))(a, idx)
        self.assertEqual(out.shape, (4, 2, 1))