teeny_quant/quant/
granularity.rs

1/*
2 * Copyright (c) 2026 Teenygrad.
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 *   http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17//! How many independent scale/zero-point pairs a quantized tensor gets.
18
19/// Quantization granularity: how a tensor's elements are partitioned into groups that each get
20/// their own scale/zero-point.
21#[derive(Debug, Clone, Copy, PartialEq, Eq)]
22pub enum Granularity {
23    /// One scale/zero-point for the whole tensor.
24    PerTensor,
25
26    /// One scale/zero-point per index along `axis` (e.g. per output channel of a `Linear` or
27    /// `Conv` weight). Equivalent to [`Granularity::Group`] with `group_size: 1`.
28    PerChannel {
29        /// Axis to quantize independently along (typically the output-channel axis, `0`).
30        axis: usize,
31    },
32
33    /// One scale/zero-point per contiguous `group_size`-element block along `axis` (the
34    /// GPTQ/AWQ/compressed-tensors convention; `axis` is typically the reduction/input-channel
35    /// axis).
36    Group {
37        /// Axis the grouping runs along.
38        axis: usize,
39        /// Number of elements per group along `axis`.
40        group_size: usize,
41    },
42}
43
44impl Granularity {
45    /// The `(axis, group_size)` this granularity reduces to, or `None` for [`Granularity::PerTensor`].
46    pub(crate) fn axis_and_group_size(self) -> Option<(usize, usize)> {
47        match self {
48            Granularity::PerTensor => None,
49            Granularity::PerChannel { axis } => Some((axis, 1)),
50            Granularity::Group { axis, group_size } => Some((axis, group_size)),
51        }
52    }
53
54    /// A short, stable name for the compressed-tensors `strategy` field.
55    pub fn strategy_name(self) -> &'static str {
56        match self {
57            Granularity::PerTensor => "tensor",
58            Granularity::PerChannel { .. } => "channel",
59            Granularity::Group { .. } => "group",
60        }
61    }
62}
teeny_quant/quant/granularity.rs

teeny_quant/quant/
granularity.rs
