Exporting Functions

Exporting Functions#

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. Ruby) 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 API documentation.

Basics of Exporting#

Let’s start with a simple example:

def fun(x, y)
  x + y
end

x = mx.array(1.0)
y = mx.array(1.0)
mx.export_function("add.mlxfn", method(: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:

add_fun = mx.import_function("add.mlxfn")

out = add_fun.call(mx.array(1.0), mx.array(2.0))
# Prints: array(3, dtype=float32)
puts out

out = add_fun.call(mx.array(1.0), mx.array(3.0))
# Prints: array(4, dtype=float32)
puts out

begin
  add_fun.call(mx.array(1), mx.array(3.0))
rescue RuntimeError => e
  puts e.message
end

begin
  add_fun.call(mx.array([1.0, 2.0]), mx.array(3.0))
rescue RuntimeError => e
  puts e.message
end

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.

For a single-output function, the imported function returns a single mlx.core.array.

The inputs to export_function() and an imported function are specified as positional arguments and optional keyword arguments:

def fun(x, y)
  x + y
end

x = mx.array(1.0)
y = mx.array(1.0)

# Both arguments to fun are positional
mx.export_function("add.mlxfn", method(:fun), x, y)

imported_fun = mx.import_function("add.mlxfn")

# Ok
out = imported_fun.call(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.

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

x = mx.array(1.0)
y = mx.array(1.0)

# One argument to fun is positional, the other is a kwarg
mx.export_function("add.mlxfn", method(:fun), x, y: y)

imported_fun = mx.import_function("add.mlxfn")

# Ok
out = imported_fun.call(x, y: y)

begin
  # Raises since the keyword argument is missing
  imported_fun.call(x, y)
rescue StandardError => e
  puts e.message
end

begin
  # Raises since the keyword argument has the wrong key
  imported_fun.call(x, z: y)
rescue StandardError => e
  puts e.message
end

Exporting Modules#

An mlx.nn.Module can be exported with or without the parameters included in the exported function. Here’s an example:

nn = MLX::NN
model = nn::Linear.new(4, 4)
mx.eval(model.parameters)

call = ->(x) { model.call(x) }

mx.export_function("model.mlxfn", call, mx.zeros([4]))

In the above example, the 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 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:

nn = MLX::NN
model = nn::Linear.new(4, 4)
mx.eval(model.parameters)

call = ->(x, weight:, bias:) do
  # Set the model's parameters to the input parameters
  model.update({"weight" => weight, "bias" => bias})
  model.call(x)
end

params = model.parameters
mx.export_function(
  "model.mlxfn",
  call,
  mx.zeros([4]),
  weight: params["weight"],
  bias: params["bias"]
)

Shapeless Exports#

Just like compile(), functions can also be exported for dynamically shaped inputs. Pass true as the final argument to export_function() or exporter() to export a function which can be used for inputs with variable shapes:

mx.export_function("fun.mlxfn", ->(x) { mx.abs(x) }, mx.array([0.0]), true)
imported_abs = mx.import_function("fun.mlxfn")

# Ok
out = imported_abs.call(mx.array([-1.0]))

# Also ok
out = imported_abs.call(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 documentation on shapeless compilation 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 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 exporter():

fun = ->(x, y: nil) do
  constant = mx.array(3.0)
  x = x + y unless y.nil?
  x + constant
end

exporter = mx.exporter("fun.mlxfn", fun)
exporter.call(mx.array(1.0), y: mx.array(0.0))
exporter.close

imported_function = mx.import_function("fun.mlxfn")

# Call the function with y specified
out = imported_function.call(mx.array(1.0), y: mx.array(1.0))
puts out

In the above example the function constant data, (i.e. constant), is only saved once.

Transformations with Imported Functions#

Function transformations like grad(), vmap(), and compile() work on imported functions just like regular Ruby functions:

def fun(x)
  mx.sin(x)
end

x = mx.array(0.0)
mx.export_function("sine.mlxfn", method(:fun), x)

imported_fun = mx.import_function("sine.mlxfn")

# Take the derivative of the imported function
dfdx = mx.grad(->(t) { imported_fun.call(t) })
# Prints: array(1, dtype=float32)
puts dfdx.call(x)

# Compile the imported function
compiled_fun = mx.compile(imported_fun)
# Prints: array(0, dtype=float32)
puts compiled_fun.call(x)

Importing Functions in C++#

Importing and running functions in C++ is basically the same as importing and running them in Ruby. First, follow the instructions to setup a simple C++ project that uses MLX as a library.

Next, export a simple function from Ruby:

def fun(x, y)
  mx.exp(x + y)
end

x = mx.array(1.0)
y = mx.array(1.0)
mx.export_function("fun.mlxfn", method(:fun), x, y)

Import and run the function in C++ with only a few lines of code:

auto fun = mx::import_function("fun.mlxfn");

auto inputs = {mx::array(1.0), mx::array(1.0)};
auto outputs = fun(inputs);

// Prints: array(2, dtype=float32)
std::cout << outputs[0] << std::endl;

Imported functions can be transformed in C++ just like in Ruby. 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 Ruby and importing and running them in C++:

Contents