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

use crate::engine::{Array, DenseLinear, ModelTensors, Result, Stream};

#[derive(Debug)]
pub(super) struct VisionMlp {
    input: DenseLinear,
    output: DenseLinear,
}

impl VisionMlp {
    pub(super) fn load(tensors: &ModelTensors, prefix: &str, stream: &Stream) -> Result<Self> {
        Ok(Self {
            input: DenseLinear::load(tensors, &format!("{prefix}.linear_fc1"), stream)?,
            output: DenseLinear::load(tensors, &format!("{prefix}.linear_fc2"), stream)?,
        })
    }

    pub(super) fn forward(&self, input: &Array, stream: &Stream) -> Result<Array> {
        self.output
            .forward(&self.input.forward(input, stream)?.gelu_tanh(stream)?, stream)
    }
}