Method const_func

python/tests/test_vmap.py:541–542

(a, b)

        b = mx.random.uniform(shape=(4, 3))

        def const_func(a, b):
            return mx.array(2)

        out = mx.vmap(const_func, in_axes=(0, None))(a, b)
        self.assertTrue(mx.array_equal(mx.full((2,), 2), out))