Example MLX bindings to Zig

This is an example/experimental repo to show how Zig bindings MLX can be written based on the zig-build-mlx and zig-build-mlx-c repos. For supported features and more details on the build process please visit the zig-build-mlx repo.

Examples

Since Zig has very good C interop, you can use the C API directly. Consider this example (link to original example) from the original MLX-C repo:

int main() {
  mlx_stream stream = mlx_default_gpu_stream_new();
  float data[] = {1, 2, 3, 4, 5, 6};
  int shape[] = {2, 3};
  mlx_array arr = mlx_array_new_data(data, shape, 2, MLX_FLOAT32);

  mlx_array two = mlx_array_new_int(2);
  mlx_divide(&arr, arr, two, stream);
  print_array("divive by 2!", arr);

  mlx_arange(&arr, 0, 3, 0.5, MLX_FLOAT32, stream);
  print_array("arange", arr);

  mlx_array_free(arr);
  mlx_array_free(two);
  mlx_stream_free(stream);
  return 0;
}

Can be ported to Zig like this (full code in src/no_bindings_exaple/main.zig):

pub fn main() !void {
    const stream = c.mlx_default_cpu_stream_new();
    defer _ = c.mlx_stream_free(stream);

    const data = [_]f32{ 1, 2, 3, 4, 5, 6 };
    const shape = [_]c_int{ 2, 3 };

    var arr = c.mlx_array_new_data(@ptrCast(&data), &shape, 2, c.MLX_FLOAT32);
    defer _ = c.mlx_array_free(arr);

    const two = c.mlx_array_new_int(2);
    defer _ = c.mlx_array_free(two);

    _ = c.mlx_divide(&arr, arr, two, stream);
    print_array("divided by 2\n", arr);

    _ = c.mlx_arange(&arr, 0, 3, 0.5, c.MLX_FLOAT32, stream);
    print_array("arange\n", arr);
}

But to get more of the benefits that Zig offers it's probably beneficial to create some thin wrappers and utilities around the C API. I've started tinkering with this in the src/bindings_example. Something like so:

pub const String = struct {
    inner: c.mlx_string,

    pub fn init() String {
        return .{ .inner = c.mlx_string_new() };
    }

    // etc. etc.

But this is all just experimental and playing round for now.

Usage

⚠️ Caution: I've only tested this with Zig v 0.13.0.

if you have Zig installed you can build and run the example with this command. The first run will compile MLX from ground up which can take a few minutes, after that it's fast.

zig build run