mlx-sys 0.0.8

Low-level interface and binding generation for the mlx library
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112

/* Copyright © 2023-2024 Apple Inc. */

#include <stdio.h>
#include <stdlib.h>

#include "mlx/c/mlx.h"

void print_array(const char* msg, mlx_array arr) {
  mlx_string str;
  str = mlx_tostring(arr);
  printf("%s\n%s\n", msg, mlx_string_data(str));
  mlx_free(str);
}

mlx_array inc_fun(mlx_array in) {
  mlx_stream stream = mlx_gpu_stream();
  mlx_array y = mlx_array_from_float(1.0);
  mlx_array res = mlx_add(in, y, stream);
  mlx_free(y);
  mlx_free(stream);
  return res;
}

mlx_vector_array inc_fun_value(mlx_vector_array in, void* payload) {
  mlx_stream stream = mlx_gpu_stream();
  if (mlx_vector_array_size(in) != 1) {
    fprintf(stderr, "inc_func_value: expected 1 argument");
    exit(EXIT_FAILURE);
  }
  mlx_array value = (mlx_array)payload;
  mlx_array res = mlx_add(mlx_vector_array_get(in, 0), value, stream);
  mlx_vector_array vres = mlx_vector_array_from_array(res);
  mlx_free(res);
  mlx_free(stream);
  return vres;
}

int main() {
  mlx_array x = mlx_array_from_float(1.0);
  mlx_array y = mlx_array_from_float(1.0);
  mlx_closure cls = mlx_closure_new_unary(inc_fun);
  mlx_closure cls_with_value =
      mlx_closure_new_with_payload(inc_fun_value, y, mlx_free);

  // jvp
  {
    printf("jvp:\n");
    mlx_array one = mlx_array_from_float(1.0);
    mlx_vector_array primals = mlx_vector_array_from_array(x);
    mlx_vector_array tangents = mlx_vector_array_from_array(one);
    mlx_vector_vector_array res = mlx_jvp(cls, primals, tangents);
    mlx_array out = mlx_vector_vector_array_get2d(res, 0, 0);
    mlx_array dout = mlx_vector_vector_array_get2d(res, 1, 0);

    print_array("out", out);
    print_array("dout", dout);

    mlx_free(dout);
    mlx_free(out);
    mlx_free(res);
    mlx_free(tangents);
    mlx_free(primals);
    mlx_free(one);
  }

  // value_and_grad
  {
    printf("value_and_grad:\n");
    int garg = 0;
    mlx_closure_value_and_grad vag = mlx_value_and_grad(cls, &garg, 1);
    mlx_vector_array inputs = mlx_vector_array_from_array(x);
    mlx_vector_vector_array res = mlx_closure_value_and_grad_apply(vag, inputs);
    mlx_array out = mlx_vector_vector_array_get2d(res, 0, 0);
    mlx_array dout = mlx_vector_vector_array_get2d(res, 1, 0);

    print_array("out", out);
    print_array("dout", dout);

    mlx_free(dout);
    mlx_free(out);
    mlx_free(inputs);
    mlx_free(res);
    mlx_free(vag);
  }

  // value_and_grad with payload
  {
    printf("value_and_grad with payload:\n");
    int garg = 0;
    mlx_closure_value_and_grad vag =
        mlx_value_and_grad(cls_with_value, &garg, 1);
    mlx_vector_array inputs = mlx_vector_array_from_array(x);
    mlx_vector_vector_array res = mlx_closure_value_and_grad_apply(vag, inputs);
    mlx_array out = mlx_vector_vector_array_get2d(res, 0, 0);
    mlx_array dout = mlx_vector_vector_array_get2d(res, 1, 0);

    print_array("out", out);
    print_array("dout", dout);

    mlx_free(dout);
    mlx_free(out);
    mlx_free(inputs);
    mlx_free(res);
    mlx_free(vag);
  }

  mlx_free(cls_with_value);
  mlx_free(cls);
  mlx_free(x);

  return 0;
}