| | .. _export_usage: |
| |
|
| | Exporting Functions |
| | =================== |
| |
|
| | .. currentmodule:: mlx.core |
| |
|
| | MLX has an API to export and import functions to and from a file. This lets you |
| | run computations written in one MLX front-end (e.g. Python) in another MLX |
| | front-end (e.g. C++). |
| |
|
| | This guide walks through the basics of the MLX export API with some examples. |
| | To see the full list of functions check-out the :ref:`API documentation |
| | <export>`. |
| |
|
| | Basics of Exporting |
| | ------------------- |
| |
|
| | Let's start with a simple example: |
| |
|
| | .. code-block:: python |
| |
|
| | def fun(x, y): |
| | return x + y |
| |
|
| | x = mx.array(1.0) |
| | y = mx.array(1.0) |
| | mx.export_function("add.mlxfn", fun, x, y) |
| |
|
| | To export a function, provide sample input arrays that the function |
| | can be called with. The data doesn't matter, but the shapes and types of the |
| | arrays do. In the above example we exported ``fun`` with two ``float32`` |
| | scalar arrays. We can then import the function and run it: |
| |
|
| | .. code-block:: python |
| |
|
| | add_fun = mx.import_function("add.mlxfn") |
| |
|
| | out, = add_fun(mx.array(1.0), mx.array(2.0)) |
| | # Prints: array(3, dtype=float32) |
| | print(out) |
| |
|
| | out, = add_fun(mx.array(1.0), mx.array(3.0)) |
| | # Prints: array(4, dtype=float32) |
| | print(out) |
| |
|
| | # Raises an exception |
| | add_fun(mx.array(1), mx.array(3.0)) |
| |
|
| | # Raises an exception |
| | add_fun(mx.array([1.0, 2.0]), mx.array(3.0)) |
| |
|
| | Notice the third and fourth calls to ``add_fun`` raise exceptions because the |
| | shapes and types of the inputs are different than the shapes and types of the |
| | example inputs we exported the function with. |
| |
|
| | Also notice that even though the original ``fun`` returns a single output |
| | array, the imported function always returns a tuple of one or more arrays. |
| |
|
| | The inputs to :func:`export_function` and to an imported function can be |
| | specified as variable positional arguments or as a tuple of arrays: |
| |
|
| | .. code-block:: python |
| |
|
| | def fun(x, y): |
| | return x + y |
| |
|
| | x = mx.array(1.0) |
| | y = mx.array(1.0) |
| |
|
| | # Both arguments to fun are positional |
| | mx.export_function("add.mlxfn", fun, x, y) |
| |
|
| | # Same as above |
| | mx.export_function("add.mlxfn", fun, (x, y)) |
| |
|
| | imported_fun = mx.import_function("add.mlxfn") |
| |
|
| | # Ok |
| | out, = imported_fun(x, y) |
| |
|
| | # Also ok |
| | out, = imported_fun((x, y)) |
| |
|
| | You can pass example inputs to functions as positional or keyword arguments. If |
| | you use keyword arguments to export the function, then you have to use the same |
| | keyword arguments when calling the imported function. |
| |
|
| | .. code-block:: python |
| |
|
| | def fun(x, y): |
| | return x + y |
| |
|
| | # One argument to fun is positional, the other is a kwarg |
| | mx.export_function("add.mlxfn", fun, x, y=y) |
| |
|
| | imported_fun = mx.import_function("add.mlxfn") |
| |
|
| | # Ok |
| | out, = imported_fun(x, y=y) |
| |
|
| | # Also ok |
| | out, = imported_fun((x,), {"y": y}) |
| |
|
| | # Raises since the keyword argument is missing |
| | out, = imported_fun(x, y) |
| |
|
| | # Raises since the keyword argument has the wrong key |
| | out, = imported_fun(x, z=y) |
| |
|
| |
|
| | Exporting Modules |
| | ----------------- |
| |
|
| | An :obj:`mlx.nn.Module` can be exported with or without the parameters included |
| | in the exported function. Here's an example: |
| |
|
| | .. code-block:: python |
| |
|
| | model = nn.Linear(4, 4) |
| | mx.eval(model.parameters()) |
| |
|
| | def call(x): |
| | return model(x) |
| |
|
| | mx.export_function("model.mlxfn", call, mx.zeros(4)) |
| |
|
| | In the above example, the :obj:`mlx.nn.Linear` module is exported. Its |
| | parameters are also saved to the ``model.mlxfn`` file. |
| |
|
| | .. note:: |
| |
|
| | For enclosed arrays inside an exported function, be extra careful to ensure |
| | they are evaluated. The computation graph that gets exported will include |
| | the computation that produces enclosed inputs. |
| |
|
| | If the above example was missing ``mx.eval(model.parameters()``, the |
| | exported function would include the random initialization of the |
| | :obj:`mlx.nn.Module` parameters. |
| |
|
| | If you only want to export the ``Module.__call__`` function without the |
| | parameters, pass them as inputs to the ``call`` wrapper: |
| |
|
| | .. code-block:: python |
| |
|
| | model = nn.Linear(4, 4) |
| | mx.eval(model.parameters()) |
| |
|
| | def call(x, **params): |
| | # Set the model's parameters to the input parameters |
| | model.update(tree_unflatten(list(params.items()))) |
| | return model(x) |
| |
|
| | params = tree_flatten(model.parameters(), destination={}) |
| | mx.export_function("model.mlxfn", call, (mx.zeros(4),), params) |
| |
|
| |
|
| | Shapeless Exports |
| | ----------------- |
| |
|
| | Just like :func:`compile`, functions can also be exported for dynamically shaped |
| | inputs. Pass ``shapeless=True`` to :func:`export_function` or :func:`exporter` |
| | to export a function which can be used for inputs with variable shapes: |
| |
|
| | .. code-block:: python |
| |
|
| | mx.export_function("fun.mlxfn", mx.abs, mx.array([0.0]), shapeless=True) |
| | imported_abs = mx.import_function("fun.mlxfn") |
| |
|
| | # Ok |
| | out, = imported_abs(mx.array([-1.0])) |
| |
|
| | # Also ok |
| | out, = imported_abs(mx.array([-1.0, -2.0])) |
| |
|
| | With ``shapeless=False`` (which is the default), the second call to |
| | ``imported_abs`` would raise an exception with a shape mismatch. |
| |
|
| | Shapeless exporting works the same as shapeless compilation and should be |
| | used carefully. See the :ref:`documentation on shapeless compilation |
| | <shapeless_compile>` for more information. |
| |
|
| | Exporting Multiple Traces |
| | ------------------------- |
| |
|
| | In some cases, functions build different computation graphs for different |
| | input arguments. A simple way to manage this is to export to a new file with |
| | each set of inputs. This is a fine option in many cases. But it can be |
| | suboptimal if the exported functions have a large amount of duplicate constant |
| | data (for example the parameters of a :obj:`mlx.nn.Module`). |
| |
|
| | The export API in MLX lets you export multiple traces of the same function to |
| | a single file by creating an exporting context manager with :func:`exporter`: |
| |
|
| | .. code-block:: python |
| |
|
| | def fun(x, y=None): |
| | constant = mx.array(3.0) |
| | if y is not None: |
| | x += y |
| | return x + constant |
| |
|
| | with mx.exporter("fun.mlxfn", fun) as exporter: |
| | exporter(mx.array(1.0)) |
| | exporter(mx.array(1.0), y=mx.array(0.0)) |
| |
|
| | imported_function = mx.import_function("fun.mlxfn") |
| |
|
| | # Call the function with y=None |
| | out, = imported_function(mx.array(1.0)) |
| | print(out) |
| |
|
| | # Call the function with y specified |
| | out, = imported_function(mx.array(1.0), y=mx.array(1.0)) |
| | print(out) |
| |
|
| | In the above example the function constant data, (i.e. ``constant``), is only |
| | saved once. |
| |
|
| | Transformations with Imported Functions |
| | --------------------------------------- |
| |
|
| | Function transformations like :func:`grad`, :func:`vmap`, and :func:`compile` work |
| | on imported functions just like regular Python functions: |
| |
|
| | .. code-block:: python |
| |
|
| | def fun(x): |
| | return mx.sin(x) |
| |
|
| | x = mx.array(0.0) |
| | mx.export_function("sine.mlxfn", fun, x) |
| |
|
| | imported_fun = mx.import_function("sine.mlxfn") |
| |
|
| | # Take the derivative of the imported function |
| | dfdx = mx.grad(lambda x: imported_fun(x)[0]) |
| | # Prints: array(1, dtype=float32) |
| | print(dfdx(x)) |
| |
|
| | # Compile the imported function |
| | mx.compile(imported_fun) |
| | # Prints: array(0, dtype=float32) |
| | print(compiled_fun(x)[0]) |
| |
|
| |
|
| | Importing Functions in C++ |
| | -------------------------- |
| |
|
| | Importing and running functions in C++ is basically the same as importing and |
| | running them in Python. First, follow the :ref:`instructions <mlx_in_cpp>` to |
| | setup a simple C++ project that uses MLX as a library. |
| |
|
| | Next, export a simple function from Python: |
| |
|
| | .. code-block:: python |
| |
|
| | def fun(x, y): |
| | return mx.exp(x + y) |
| |
|
| | x = mx.array(1.0) |
| | y = mx.array(1.0) |
| | mx.export_function("fun.mlxfn", fun, x, y) |
| |
|
| |
|
| | Import and run the function in C++ with only a few lines of code: |
| |
|
| | .. code-block:: c++ |
| |
|
| | auto fun = mx::import_function("fun.mlxfn"); |
| |
|
| | auto inputs = {mx::array(1.0), mx::array(1.0)}; |
| | auto outputs = fun(inputs); |
| |
|
| | |
| | std::cout << outputs[0] << std::endl; |
| |
|
| | Imported functions can be transformed in C++ just like in Python. Use |
| | ``std::vector<mx::array>`` for positional arguments and ``std::map<std::string, |
| | mx::array>`` for keyword arguments when calling imported functions in C++. |
| |
|
| | More Examples |
| | ------------- |
| |
|
| | Here are a few more complete examples exporting more complex functions from |
| | Python and importing and running them in C++: |
| | |
| | * `Inference and training a multi-layer perceptron <https: |
| |
|
