mlx-c-rust 0.32.4-alpha.3

A complete, low-level, policy-free Rust bridge to the MLX-C API of Apple's MLX for Apple silicon, tracking MLX releases. Version = the verified MLX release (0.32.4; carries MLX-C 0.7 forward).
Documentation 
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
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150

/* The mlx-c-rust gap shim implementation. See gap_shim.h for the contract
 * and the retirement policy. The handle casts below inline what MLX-C's
 * private headers do: the public C struct carries the C++ object as a
 * void* ctx (that layout is the published C ABI), so casting ctx is both
 * self-contained — this file compiles against any installed MLX-C, no
 * private headers — and exactly what the vendored MLX-C itself does.
 */

#include "mlxrs/gap_shim.h"

#include <optional>
#include <utility>
#include <vector>

#include "mlx/c/error.h"
#include "mlx/mlx.h"
#include "mlx/ops.h"

namespace {

inline mlx::core::array& shim_array_get(const mlx_array handle) {
  if (!handle.ctx) {
    throw std::runtime_error("expected a non-empty mlx_array");
  }
  return *static_cast<mlx::core::array*>(handle.ctx);
}

inline void shim_array_set(mlx_array& destination, mlx::core::array&& value) {
  if (destination.ctx) {
    *static_cast<mlx::core::array*>(destination.ctx) = std::move(value);
  } else {
    destination.ctx = new mlx::core::array(std::move(value));
  }
}

inline std::vector<mlx::core::array>& shim_vector_array_get(
    const mlx_vector_array handle) {
  if (!handle.ctx) {
    throw std::runtime_error("expected a non-empty mlx_vector_array");
  }
  return *static_cast<std::vector<mlx::core::array>*>(handle.ctx);
}

inline void shim_vector_array_set(
    mlx_vector_array& destination,
    std::vector<mlx::core::array>&& value) {
  if (destination.ctx) {
    *static_cast<std::vector<mlx::core::array>*>(destination.ctx) =
        std::move(value);
  } else {
    destination.ctx = new std::vector<mlx::core::array>(std::move(value));
  }
}

inline mlx::core::Stream& shim_stream_get(const mlx_stream handle) {
  if (!handle.ctx) {
    throw std::runtime_error("expected a non-empty mlx_stream");
  }
  return *static_cast<mlx::core::Stream*>(handle.ctx);
}

} // namespace

extern "C" int mlxrs_minmax_axes(
    mlx_array* res_0,
    mlx_array* res_1,
    const mlx_array a,
    const int* axes,
    size_t axes_num,
    bool keepdims,
    const mlx_stream s) {
  try {
    auto [minima, maxima] = mlx::core::minmax(
        shim_array_get(a),
        std::vector<int>(axes, axes + axes_num),
        keepdims,
        shim_stream_get(s));
    shim_array_set(*res_0, std::move(minima));
    shim_array_set(*res_1, std::move(maxima));
  } catch (std::exception& error) {
    mlx_error("%s", error.what());
    return 1;
  }
  return 0;
}

extern "C" int mlxrs_minmax_axis(
    mlx_array* res_0,
    mlx_array* res_1,
    const mlx_array a,
    int axis,
    bool keepdims,
    const mlx_stream s) {
  try {
    auto [minima, maxima] =
        mlx::core::minmax(shim_array_get(a), axis, keepdims, shim_stream_get(s));
    shim_array_set(*res_0, std::move(minima));
    shim_array_set(*res_1, std::move(maxima));
  } catch (std::exception& error) {
    mlx_error("%s", error.what());
    return 1;
  }
  return 0;
}

extern "C" int mlxrs_matrix_transpose(
    mlx_array* res,
    const mlx_array a,
    const mlx_stream s) {
  try {
    shim_array_set(
        *res,
        mlx::core::matrix_transpose(shim_array_get(a), shim_stream_get(s)));
  } catch (std::exception& error) {
    mlx_error("%s", error.what());
    return 1;
  }
  return 0;
}

extern "C" int mlxrs_unique(
    mlx_vector_array* res,
    const mlx_array a,
    int size,
    bool return_index,
    bool return_inverse,
    bool return_counts,
    const mlx_array fill_value,
    const mlx_stream s) {
  try {
    std::optional<mlx::core::array> fill;
    if (fill_value.ctx) {
      fill = shim_array_get(fill_value);
    }
    shim_vector_array_set(
        *res,
        mlx::core::unique(
            shim_array_get(a),
            size,
            return_index,
            return_inverse,
            return_counts,
            fill,
            shim_stream_get(s)));
  } catch (std::exception& error) {
    mlx_error("%s", error.what());
    return 1;
  }
  return 0;
}