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

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

#include <cstring>

#include "mlx/c/closure.h"
#include "mlx/c/object.h"
#include "mlx/c/private/closure.h"
#include "mlx/c/private/string.h"
#include "mlx/c/private/utils.h"

mlx_string_* mlx_closure_::tostring() {
  RETURN_MLX_C_STRING("mlx_vector_array closure(mlx_vector_array, void*)");
}

extern "C" mlx_closure mlx_closure_new(
    mlx_vector_array (*fun)(const mlx_vector_array)) {
  MLX_TRY_CATCH(
      auto cpp_closure =
          [fun](const std::vector<mlx::core::array>& input) {
            auto c_input = new mlx_vector_array_(input);
            auto c_res = fun(c_input);
            mlx_free(c_input);
            auto res = c_res->ctx;
            mlx_free(c_res);
            return res;
          };
      return new mlx_closure_(cpp_closure), return nullptr);
}

extern "C" mlx_closure mlx_closure_new_unary(
    mlx_array (*fun)(const mlx_array)) {
  MLX_TRY_CATCH(
      auto cpp_closure =
          [fun](const std::vector<mlx::core::array>& input) {
            if (input.size() != 1) {
              throw std::runtime_error("closure: expected unary input");
            }
            auto c_input = new mlx_array_(input[0]);
            auto c_res = fun(c_input);
            mlx_free(c_input);
            std::vector<mlx::core::array> res = {c_res->ctx};
            mlx_free(c_res);
            return res;
          };
      return new mlx_closure_(cpp_closure), return nullptr);
}

extern "C" mlx_closure mlx_closure_new_with_payload(
    mlx_vector_array (*fun)(const mlx_vector_array, void*),
    void* payload,
    void (*dtor)(void*)) {
  auto cpp_payload = std::shared_ptr<void>(payload, dtor);
  auto cpp_closure =
      [fun, cpp_payload, dtor](const std::vector<mlx::core::array>& input) {
        auto c_input = new mlx_vector_array_(input);
        auto c_res = fun(c_input, cpp_payload.get());
        mlx_free(c_input);
        auto res = c_res->ctx;
        mlx_free(c_res);
        return res;
      };
  MLX_TRY_CATCH(return new mlx_closure_(cpp_closure), return nullptr);
}

extern "C" mlx_vector_array mlx_closure_apply(
    mlx_closure cls,
    const mlx_vector_array inputs) {
  MLX_TRY_CATCH(auto cpp_res = cls->ctx(inputs->ctx);
                return new mlx_vector_array_(cpp_res), return nullptr);
}

extern "C" mlx_vector_vector_array mlx_closure_value_and_grad_apply(
    mlx_closure_value_and_grad cls,
    const mlx_vector_array inputs) {
  MLX_TRY_CATCH(auto cpp_res = cls->ctx(inputs->ctx);
                return new mlx_vector_vector_array_(cpp_res), return nullptr);
}

mlx_string_* mlx_closure_value_and_grad_::tostring() {
  MLX_TRY_CATCH(
      return new mlx_string_(
                 "mlx_vector_vector_array closure(mlx_vector_array, void*)"),
             return nullptr);
}