mlx.core.vmap

vmap(fun: Callable, in_axes: object = 0, out_axes: object = 0) → Callable

返回函数 fun 的向量化版本。

参数：

• fun (Callable) – 一个函数，它接受可变数量的 array 或 array 树，并返回可变数量的 array 或 array 树。

• in_axes (int, 可选) – 一个整数或输入到 fun 的有效前缀树，其中每个节点指定了向量化轴。如果值为 None，则相应的输入不会被向量化。默认为 0。

• out_axes (int, 可选) – 一个整数或 fun 输出的有效前缀树，其中每个节点指定了向量化轴。如果值为 None，则相应的输出不会被向量化。默认为 0。

返回值：

向量化函数。

返回类型：

Callable