mlx.core.export_function

mlx.core.export_function#

export_function(file: str, fun: Callable, *args, shapeless: bool = False, **kwargs) → None#

将函数导出到文件。

必须提供示例输入数组来导出函数。示例输入可以是可变参数 *args 和 **kwargs，或者是一个数组元组和/或一个字符串键对应数组值的字典。

警告

这是实验性 API 的一部分，未来版本的 MLX 中可能会发生变化。使用旧版本 MLX 导出的函数可能与未来版本不兼容。

参数:

  • file (str) – 导出函数的文件路径。

  • fun (Callable) – 一个函数，接受零个或多个 array 作为输入，并返回一个或多个 array。

  • *args (array) – 函数的示例数组输入。

  • shapeless (bool, optional) – 函数是否允许输入具有可变形状。默认值: False。

  • **kwargs (array) – 函数的额外示例关键字数组输入。

示例

def fun(x, y):
    return x + y

x = mx.array(1)
y = mx.array([1, 2, 3])
mx.export_function("fun.mlxfn", fun, x, y=y)
目录