burn-mlx 0.1.2

MLX backend for Burn deep learning framework - native Apple Silicon GPU acceleration
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

//! Module operations (linear, etc.) for MLX backend.

use crate::element::MlxElement;
use crate::tensor::MlxTensor;

/// Module operations on MLX tensors.
impl<E: MlxElement> MlxTensor<E> {
    /// Linear layer (matmul + optional bias).
    /// input: [batch, in_features]
    /// weight: [out_features, in_features]
    /// output: [batch, out_features]
    pub fn linear(&self, weight: &MlxTensor<E>, bias: Option<&MlxTensor<E>>) -> MlxTensor<E> {
        let weight_t = mlx_rs::ops::transpose(&weight.array, None)
            .expect("Failed to transpose weight");
        let out = self.array.matmul(&weight_t)
            .expect("Failed to matmul");

        let result = if let Some(b) = bias {
            mlx_rs::ops::add(&out, &b.array)
                .expect("Failed to add bias")
        } else {
            out
        };

        MlxTensor::new(result, self.device)
    }

    /// Layer normalization.
    pub fn layer_norm(&self, weight: Option<&MlxTensor<E>>, bias: Option<&MlxTensor<E>>, epsilon: f32) -> MlxTensor<E> {
        // Normalize over the last dimension
        let mean = mlx_rs::ops::mean(&self.array, &[-1], true)
            .expect("Failed to compute mean");
        let var = mlx_rs::ops::var(&self.array, &[-1], true, 0)
            .expect("Failed to compute var");

        let eps = mlx_rs::Array::from_f32(epsilon);
        let var_eps = mlx_rs::ops::add(&var, &eps).expect("Failed to add eps");
        let std = mlx_rs::ops::sqrt(&var_eps).expect("Failed to sqrt");

        let centered = mlx_rs::ops::subtract(&self.array, &mean).expect("Failed to subtract");
        let normalized = mlx_rs::ops::divide(&centered, &std).expect("Failed to divide");

        let scaled = if let Some(w) = weight {
            mlx_rs::ops::multiply(&normalized, &w.array).expect("Failed to scale")
        } else {
            normalized
        };

        let result = if let Some(b) = bias {
            mlx_rs::ops::add(&scaled, &b.array).expect("Failed to add bias")
        } else {
            scaled
        };

        MlxTensor::new(result, self.device)
    }

    /// Embedding lookup.
    pub fn embedding(&self, indices: &MlxTensor<i32>) -> MlxTensor<E> {
        let array = mlx_rs::ops::indexing::take(&self.array, &indices.array, 0)
            .expect("Failed to embedding lookup");
        MlxTensor::new(array, self.device)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::device::MlxDevice;

    #[test]
    fn test_linear() {
        let input = MlxTensor::<f32>::ones(&[2, 4], MlxDevice::Gpu);
        let weight = MlxTensor::<f32>::ones(&[3, 4], MlxDevice::Gpu);
        let output = input.linear(&weight, None);
        assert_eq!(output.shape(), vec![2, 3]);
    }
}