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
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

//! Validated axis-specific padding for lazy MLX tensors.

use crate::MlxBindingsContext;
use crate::MlxCError;
use crate::{MlxArray, MlxDtype, raw};

impl MlxBindingsContext {
    /// Applies axis-specific constant padding, including asymmetric image padding.
    pub fn pad(
        &self,
        input: &MlxArray,
        axes: &[i32],
        low_padding: &[i32],
        high_padding: &[i32],
        pad_value: f32,
    ) -> Result<MlxArray, MlxCError> {
        const OPERATION: &str = "pad an MLX array";
        validate_padding_arguments(input, axes, low_padding, high_padding, pad_value)?;
        let float_pad_value = self.array_from_f32(&[pad_value], &[])?;
        let typed_pad_value = self.astype(&float_pad_value, input.dtype())?;
        self.output_array(OPERATION, |output_array, stream| {
            // SAFETY: Input, scalar, and stream are live; axis and padding
            // vectors were validated and remain borrowed for this graph call.
            unsafe {
                raw::mlx_pad(
                    output_array,
                    input.raw(),
                    axes.as_ptr(),
                    axes.len(),
                    low_padding.as_ptr(),
                    low_padding.len(),
                    high_padding.as_ptr(),
                    high_padding.len(),
                    typed_pad_value.raw(),
                    c"constant".as_ptr(),
                    stream,
                )
            }
        })
    }
}

fn validate_padding_arguments(
    input: &MlxArray,
    axes: &[i32],
    low_padding: &[i32],
    high_padding: &[i32],
    pad_value: f32,
) -> Result<(), MlxCError> {
    const OPERATION: &str = "pad an MLX array";
    if axes.is_empty() || axes.len() != low_padding.len() || axes.len() != high_padding.len() {
        return Err(runtime_operation_error(
            OPERATION,
            "padding axes and low/high sizes must have one matching nonempty length",
        ));
    }
    if !matches!(
        input.dtype(),
        MlxDtype::Float16 | MlxDtype::Float32 | MlxDtype::BFloat16
    ) || !pad_value.is_finite()
    {
        return Err(runtime_operation_error(
            OPERATION,
            "padding requires a supported floating array and finite pad value",
        ));
    }
    if low_padding.iter().chain(high_padding).any(|size| *size < 0) {
        return Err(runtime_operation_error(
            OPERATION,
            "padding sizes must be nonnegative",
        ));
    }
    let input_shape = input.shape();
    let rank = i32::try_from(input_shape.len())
        .map_err(|_| runtime_operation_error(OPERATION, "input rank exceeds i32 range"))?;
    let mut normalized_axes = Vec::with_capacity(axes.len());
    for axis in axes {
        if *axis < -rank || *axis >= rank {
            return Err(runtime_operation_error(
                OPERATION,
                "padding axes must refer to existing input dimensions",
            ));
        }
        let normalized_axis = if *axis < 0 { *axis + rank } else { *axis };
        if normalized_axes.contains(&normalized_axis) {
            return Err(runtime_operation_error(
                OPERATION,
                "padding axes must be unique",
            ));
        }
        normalized_axes.push(normalized_axis);
    }
    for ((normalized_axis, low_size), high_size) in
        normalized_axes.iter().zip(low_padding).zip(high_padding)
    {
        let axis_index = usize::try_from(*normalized_axis)
            .map_err(|_| runtime_operation_error(OPERATION, "padding axis conversion failed"))?;
        input_shape[axis_index]
            .checked_add(*low_size)
            .and_then(|size| size.checked_add(*high_size))
            .ok_or_else(|| {
                runtime_operation_error(OPERATION, "padded dimension exceeds i32 range")
            })?;
    }
    Ok(())
}

fn runtime_operation_error(operation: &'static str, description: &'static str) -> MlxCError {
    MlxCError {
        operation,
        description: description.to_owned(),
    }
}