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
60
61
62
63
64
65

use models::weights::{TensorBinding, TensorStorage};

use super::{Error, ModelTensors, NormWeight, QuantizedEmbedding, QuantizedLinear, Result, Stream};

pub(super) fn affine_linear(
    tensors: &ModelTensors,
    binding: &TensorBinding,
) -> Result<QuantizedLinear> {
    let TensorStorage::AffineQuantized {
        scales,
        biases: Some(biases),
        output_bias,
        group_size: Some(group_size),
        ..
    } = &binding.storage
    else {
        return Err(invalid("linear", binding));
    };
    QuantizedLinear::load_names(
        tensors,
        &binding.source,
        scales,
        biases,
        output_bias.as_deref(),
        i32::try_from(*group_size)?,
    )
}

pub(super) fn affine_embedding(
    tensors: &ModelTensors,
    binding: &TensorBinding,
) -> Result<QuantizedEmbedding> {
    let TensorStorage::AffineQuantized {
        scales,
        biases: Some(biases),
        group_size: Some(group_size),
        ..
    } = &binding.storage
    else {
        return Err(invalid("embedding", binding));
    };
    QuantizedEmbedding::load_names(
        tensors,
        &binding.source,
        scales,
        biases,
        i32::try_from(*group_size)?,
    )
}

pub(super) fn adjusted_norm(
    tensors: &ModelTensors,
    binding: &TensorBinding,
    shift: f32,
    stream: &Stream,
) -> Result<NormWeight> {
    NormWeight::load_name_adjusted(tensors, &binding.source, shift, stream)
}

fn invalid(kind: &str, binding: &TensorBinding) -> Error {
    Error::InvalidQuantization(format!(
        "{kind} requires a complete affine binding: {}",
        binding.source
    ))
}