Overview

MLX C wraps and follows as closely as possible the C++ API of MLX.

C Object

MLX C relies on several C struct which are all compatible with a private object struct (that we will refer as a MLX C object), which implements basic memory management. When using MLX C, it is important to note that:

• All returned C objects come with a reference count of 1, which needs to be decreased with a matching mlx_free() call.

• When the reference count of an object gets to 0, the object is freed.

• Reference count may be increased via mlx_retain(), in which case it should be matched at some point with a mlx_free() to avoid any memory leak. Increasing the reference count may be necessary when storing a MLX C object into a custom C structure.

In addition, MLX C objects come with a convenience mlx_tostring() function which returns a string describing the contents of the object (which is particularly interesting for array objects).

Array

The most important object in MLX C is certainly the array (mlx_array), which holds the data on which computations are performed. As MLX is lazy, the contents of the array obtained via the mlx_array_data_*() functions are valid only if mlx_eval() as been called (see transforms).

Vector of Arrays, and Vector of Vector of Arrays

Vector of arrays (mlx_vector_array) can hold multiple arrays, and vector of vector arrays (mlx_vector_vector_array) can hold multiple vector of arrays.

An array added to a mlx_vector_array will stay alive until the vector of arrays is destroyed (via mlx_free()). Arrays returned by mlx_vector_array_get() will need to be matched with a corresponding mlx_free().

Same idea applies to mlx_vector_vector_array.

Device and Stream

In MLX, arrays are not tied to a device. Instead, operations on arrays are scheduled on a stream, which is associated to a particular device.

MLX C provides MLX_CPU_STREAM and MLX_GPU_STREAM, which point to the default CPU and GPU streams. See the basic MLX C example.

String and Maps

MLX C has a mlx_string which encapsulates a C char pointer. Just like other MLX C objects, it must be freed with mlx_free().

MLX C also has a string-to-array map named mlx_map_string_to_array.

Futures

In MLX C, some async operations (see for example transforms) may return a future. The only operation one can currently perform on futures is wait(), which ensure resources associated with the future are then materialized.

Array Operations

Many array operations are available, with additional support for random number generation, and FFTs. Advanced linear algebra operations are in their early stages.

IO Operations

MLX C wraps a number of array IO operations, which save and load arrays in several common formats. See IO utils for specific MLX objects defined for IO purposes.

Function Transformations

MLX supports the concept of function transforms.

These are also available in MLX C through the use of closures that contain a C function pointer and optional payloads. Closures obey the same memory management rules as other MLX C objects and must be release with a matching mlx_free() call.

MLX C transforms will are applied on closures and may return closures.

For more details, see the example using closures.

Compilation

When using the same function multiple times, compilation may be beneficial. Compiling functions makes them more efficient by reducing redundant work, fusing kernels, and reducing overhead. Compilation operations are function transformations which take a closure and return a new closure (which is the compiled version of the given closure).

Fast Custom Ops

To maximize performance MLX has fast custom implementations for some common operations.

Metal backend-specific functions

MLX C exposes some useful functions related to the MLX Metal backend.