libmir-metal 0.2.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

use super::dimension;
use crate::engine::{Array, DenseLinear, Error, LayerNorm, ModelTensors, Result, Stream};

#[derive(Debug)]
pub(super) struct PatchMerger {
    norm: LayerNorm,
    input: DenseLinear,
    output: DenseLinear,
    hidden_size: usize,
    merge: usize,
}

impl PatchMerger {
    pub(super) fn load(
        tensors: &ModelTensors,
        hidden_size: usize,
        merge: usize,
        prefix: &str,
        stream: &Stream,
    ) -> Result<Self> {
        Ok(Self {
            norm: LayerNorm::load(tensors, &format!("{prefix}.merger.norm"), 1.0e-6)?,
            input: DenseLinear::load(tensors, &format!("{prefix}.merger.linear_fc1"), stream)?,
            output: DenseLinear::load(tensors, &format!("{prefix}.merger.linear_fc2"), stream)?,
            hidden_size,
            merge,
        })
    }

    pub(super) fn forward(&self, input: &Array, stream: &Stream) -> Result<Array> {
        let shape = input.shape()?;
        let sequence = usize::try_from(*shape.get(1).ok_or(Error::ShapeOverflow)?)?;
        let unit = self.merge * self.merge;
        if !sequence.is_multiple_of(unit) {
            return Err(Error::InvalidModel(
                "spatial-merge vision patch sequence is incompatible with merger".into(),
            ));
        }
        let normalized = self.norm.forward(input, stream)?.reshape(
            &[
                dimension(sequence / unit, "merged sequence")?,
                dimension(self.hidden_size * unit, "merged hidden size")?,
            ],
            stream,
        )?;
        let output = self
            .output
            .forward(&self.input.forward(&normalized, stream)?.gelu(stream)?, stream)?;
        let output_shape = output.shape()?;
        output.reshape(
            &[
                1,
                dimension(sequence / unit, "merged sequence")?,
                *output_shape.get(1).ok_or(Error::ShapeOverflow)?,
            ],
            stream,
        )
    }
}