mlx.core.compile

compile(fun: Callable, inputs: Optional[object] = None, outputs: Optional[object] = None, shapeless: bool = False) → object

Returns a compiled function which produces the same output as fun.

Parameters:

• fun (callable) – A function which takes a variable number of array or trees of array and returns a variable number of array or trees of array.

• inputs (list or dict, optional) – These inputs will be captured during the function compilation along with the inputs to fun. The inputs can be a list or a dict containing arbitrarily nested lists, dictionaries, or arrays. Leaf nodes that are not array are ignored. Default: None

• outputs (list or dict, optional) – These outputs will be captured and updated in a compiled function. The outputs can be a list or a dict containing arbitrarily nested lists, dictionaries, or arrays. Leaf nodes that are not array are ignored. Default: None

• shapeless (bool, optional) – A function compiled with the shapeless option enabled will not be recompiled when the input shape changes. Not all functions can be compiled with shapeless enabled. Attempting to compile such functions with shapeless enabled will throw. Note, changing the number of dimensions or type of any input will result in a recompilation even with shapeless set to True. Default: False

Returns:

A compiled function which has the same input arguments as fun and returns the the same output(s).

Return type:

callable