| .. _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: |
|
|
