tract-linalg 0.23.0-dev.4

Tiny, no-nonsense, self contained, TensorFlow and ONNX inference
Documentation 
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
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121

use std::fmt;
use std::sync::Arc;

use tract_data::internal::*;

use super::BlockQuant;
use super::BlockQuantFact;

/// Concrete tensor storage for block-quantized weights.
///
/// Stores a single contiguous `Arc<Blob>` of quantized data along with the
/// block-quant format. Shape lives on the tensor, not here.
#[derive(Clone, PartialEq, Eq)]
pub struct BlockQuantStorage {
    format: Box<dyn BlockQuant>,
    data: Arc<Blob>,
}

impl BlockQuantStorage {
    fn expected_bytes(format: &dyn BlockQuant, m: usize, k: usize) -> usize {
        m * k / format.block_len() * format.block_bytes()
    }

    pub fn new(
        format: Box<dyn BlockQuant>,
        m: usize,
        k: usize,
        data: Arc<Blob>,
    ) -> TractResult<Self> {
        let expected = Self::expected_bytes(&*format, m, k);
        ensure!(
            data.len() == expected,
            "BlockQuantStorage::new: blob length {} does not match expected {} (m={}, k={}, format={})",
            data.len(),
            expected,
            m,
            k,
            format,
        );
        Ok(Self { format, data })
    }

    pub fn format(&self) -> &dyn BlockQuant {
        &*self.format
    }

    /// Returns the single contiguous blob.
    pub fn value(&self) -> &Arc<Blob> {
        &self.data
    }

    /// Converts this storage into a `Tensor` with the given shape.
    ///
    /// `dt` is the logical element type (e.g. f32, f16) — the type these
    /// weights represent when dequantized.
    pub fn into_tensor_with_shape(self, dt: DatumType, shape: &[usize]) -> Tensor {
        Tensor::from_storage(dt, shape, self)
    }
}

/// Returns a byte slice for a single group within contiguous block-quant data.
pub fn block_quant_slice<'a>(
    data: &'a [u8],
    format: &dyn BlockQuant,
    m_per_group: usize,
    k: usize,
    g: usize,
) -> &'a [u8] {
    let row_bytes = k / format.block_len() * format.block_bytes();
    let group_bytes = m_per_group * row_bytes;
    let start = g * group_bytes;
    &data[start..start + group_bytes]
}

impl fmt::Debug for BlockQuantStorage {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "BlockQuantStorage({}, bytes={})", self.format, self.data.len())
    }
}

impl fmt::Display for BlockQuantStorage {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "BlockQuantStorage({}, bytes={})", self.format, self.data.len())
    }
}

impl TensorStorage for BlockQuantStorage {
    fn byte_len(&self) -> usize {
        self.data.len()
    }

    fn is_empty(&self) -> bool {
        self.data.is_empty()
    }

    fn deep_clone(&self) -> Box<dyn TensorStorage> {
        Box::new(self.clone())
    }

    fn as_plain(&self) -> Option<&PlainStorage> {
        None
    }

    fn as_plain_mut(&mut self) -> Option<&mut PlainStorage> {
        None
    }

    fn into_plain(self: Box<Self>) -> Option<PlainStorage> {
        None
    }

    fn dyn_hash(&self, state: &mut dyn std::hash::Hasher) {
        state.write_u8(1);
        self.format.dyn_hash(state);
        state.write(self.data.as_bytes());
    }

    fn exotic_fact(&self, shape: &[usize]) -> TractResult<Option<Box<dyn ExoticFact>>> {
        Ok(Some(Box::new(BlockQuantFact::new(dyn_clone::clone_box(&*self.format), shape.into()))))
    }
}