mlx-c-rust 0.32.4-alpha.2

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

use std::os::raw::c_int;

use crate::compile::graph;
use crate::{MlxStream, raw};

pub(crate) unsafe extern "C" fn build_sparse_shared_expert_combination_graph(
    output_vector: *mut raw::mlx_vector_array,
    input_vector: raw::mlx_vector_array,
) -> c_int {
    if output_vector.is_null() || unsafe { raw::mlx_vector_array_size(input_vector) } != 3 {
        return 1;
    }
    let sparse_expert_output = match graph::array_from_vector(input_vector, 0) {
        Ok(sparse_expert_output) => sparse_expert_output,
        Err(get_status) => return get_status,
    };
    let shared_expert_output = match graph::array_from_vector(input_vector, 1) {
        Ok(shared_expert_output) => shared_expert_output,
        Err(get_status) => return get_status,
    };
    let shared_expert_gate_logits = match graph::array_from_vector(input_vector, 2) {
        Ok(shared_expert_gate_logits) => shared_expert_gate_logits,
        Err(get_status) => return get_status,
    };
    let gpu_stream = match MlxStream::default_gpu() {
        Ok(gpu_stream) => gpu_stream,
        Err(_) => return 1,
    };
    let shared_expert_gate_weights = match graph::graph_output_array(|output_array| {
        // SAFETY: The input and stream are live, and the output is uniquely writable.
        unsafe {
            raw::mlx_sigmoid(
                output_array,
                shared_expert_gate_logits.raw(),
                gpu_stream.raw(),
            )
        }
    }) {
        Ok(shared_expert_gate_weights) => shared_expert_gate_weights,
        Err(build_status) => return build_status,
    };
    let gated_shared_expert_output = match graph::graph_output_array(|output_array| {
        // SAFETY: Inputs and stream are live, and the output is uniquely writable.
        unsafe {
            raw::mlx_multiply(
                output_array,
                shared_expert_output.raw(),
                shared_expert_gate_weights.raw(),
                gpu_stream.raw(),
            )
        }
    }) {
        Ok(gated_shared_expert_output) => gated_shared_expert_output,
        Err(build_status) => return build_status,
    };
    let combined_expert_output = match graph::graph_output_array(|output_array| {
        // SAFETY: Inputs and stream are live, and the output is uniquely writable.
        unsafe {
            raw::mlx_add(
                output_array,
                sparse_expert_output.raw(),
                gated_shared_expert_output.raw(),
                gpu_stream.raw(),
            )
        }
    }) {
        Ok(combined_expert_output) => combined_expert_output,
        Err(build_status) => return build_status,
    };
    // SAFETY: The output vector is unique and live for this callback.
    unsafe { graph::set_graph_output(output_vector, &combined_expert_output) }
}