libmir-metal 0.1.0

Metal inference backend for libmir
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

use super::{Array, ModelTensors, Result, Stream};

#[derive(Debug)]
pub struct NormWeight {
    weight: Array,
}

impl NormWeight {
    #[cfg(test)]
    pub(in crate::engine) fn from_weight(weight: Array) -> Self {
        Self { weight }
    }

    pub(in crate::engine) fn load(tensors: &ModelTensors, prefix: &str) -> Result<Self> {
        Ok(Self {
            weight: tensors.get(&format!("{prefix}.weight"))?,
        })
    }

    pub(in crate::engine) fn load_shifted(
        tensors: &ModelTensors,
        prefix: &str,
        shift: f32,
        stream: &Stream,
    ) -> Result<Self> {
        let weight = tensors.get(&format!("{prefix}.weight"))?.add_scalar(shift, stream)?;
        Ok(Self { weight })
    }

    pub(in crate::engine) fn load_adjusted(
        tensors: &ModelTensors,
        prefix: &str,
        shift: f32,
        stream: &Stream,
    ) -> Result<Self> {
        if shift == 0.0 {
            return Self::load(tensors, prefix);
        }
        Self::load_shifted(tensors, prefix, shift, stream)
    }

    pub(in crate::engine) fn load_optional(
        tensors: &ModelTensors,
        prefix: &str,
    ) -> Result<Option<Self>> {
        tensors
            .get_optional(&format!("{prefix}.weight"))
            .map(|weight| weight.map(|weight| Self { weight }))
    }

    pub(in crate::engine) fn apply(
        &self,
        input: &Array,
        eps: f32,
        stream: &Stream,
    ) -> Result<Array> {
        input.rms_norm(&self.weight, eps, stream)
    }

    pub(super) fn native_clone(&self) -> mirtal::Array {
        self.weight.native().clone()
    }
}