mlx.core.export_function

export_function(file_or_callback: str | Callable, fun: Callable, *args, shapeless: bool = False, metadata: str | None = None, **kwargs) → None

Export an MLX function.

Example input arrays must be provided to export a function. The example inputs can be variable *args and **kwargs or a tuple of arrays and/or dictionary of string keys with array values.

Warning

This is part of an experimental API which is likely to change in future versions of MLX. Functions exported with older versions of MLX may not be compatible with future versions.

Parameters:

• file_or_callback (str or Callable) – Either a file path to export the function to or a callback.

• fun (Callable) – A function which takes as input zero or more array and returns one or more array.

• *args (array) – Example array inputs to the function.

• shapeless (bool, optional) – Whether or not the function allows inputs with variable shapes. Default: False.

• metadata (str, optional) – A string to save alongside the function, for example a JSON encoded model configuration. Only supported when exporting to a file. Read it back with import_function(). Default: None.

• **kwargs (array) – Additional example keyword array inputs to the function.

Raises:

ValueError – If metadata is given when exporting with a callback.

Example

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)