edgefirst_decoder/per_scale/
helper.rs

1// SPDX-FileCopyrightText: Copyright 2026 Au-Zone Technologies
2// SPDX-License-Identifier: Apache-2.0
3
4//! Optional helper: walk a schema-v2 document and attach
5//! per-tensor `Quantization` to integer input tensors via shape match.
6//!
7//! Use when the upstream inference layer hasn't already attached
8//! quantization metadata to the tensors.
9
10use crate::schema::SchemaV2;
11use crate::{DecoderError, DecoderResult};
12use edgefirst_tensor::{Quantization as TQ, TensorDyn};
13
14/// Walk the schema's logical outputs and their per-scale children,
15/// attaching per-tensor `Quantization` to any matching integer tensor
16/// in `tensors` (matched by shape).
17///
18/// **Behavior:**
19/// - Per-channel quantization is silently skipped — the per-scale
20///   subsystem only consumes per-tensor quant.
21/// - Float tensors are silently skipped (they don't carry quantization).
22/// - Tensors with no matching schema entry are silently left alone.
23/// - If a schema entry has a quant declaration but no tensor matches
24///   the declared shape, returns `InvalidShape`.
25///
26/// **Idempotent:** safe to call multiple times; later calls overwrite
27/// the attached quantization.
28pub fn apply_schema_quant(schema: &SchemaV2, tensors: &mut [&mut TensorDyn]) -> DecoderResult<()> {
29    for logical in &schema.outputs {
30        // Per-scale children
31        for child in &logical.outputs {
32            if let Some(q) = child.quantization.as_ref() {
33                attach_per_tensor_quant_by_shape(tensors, &child.shape, q)?;
34            }
35        }
36        // Direct logical (no children) — quant lives on the logical itself.
37        if logical.outputs.is_empty() {
38            if let Some(q) = logical.quantization.as_ref() {
39                attach_per_tensor_quant_by_shape(tensors, &logical.shape, q)?;
40            }
41        }
42    }
43    Ok(())
44}
45
46fn attach_per_tensor_quant_by_shape(
47    tensors: &mut [&mut TensorDyn],
48    expected_shape: &[usize],
49    schema_q: &crate::schema::Quantization,
50) -> DecoderResult<()> {
51    // Reject empty scale up-front so a malformed schema fails fast at
52    // attach time instead of silently looking "per-channel" here and
53    // surfacing as `QuantMissing` later at decode time.
54    if schema_q.scale.is_empty() {
55        return Err(DecoderError::InvalidShape(format!(
56            "apply_schema_quant: schema declares quantization for shape \
57             {expected_shape:?} but `scale` is empty"
58        )));
59    }
60    if !schema_q.is_per_tensor() {
61        // Per-channel — skip silently. The per-scale planner errors
62        // separately when it actually needs to use a per-channel quant.
63        return Ok(());
64    }
65    let scale = schema_q.scale[0];
66    let zp = schema_q.zero_point_at(0);
67
68    for t in tensors.iter_mut() {
69        if t.shape() == expected_shape {
70            // Try to attach to integer variants; silently skip floats.
71            macro_rules! try_attach {
72                ($ty:ty) => {
73                    if let Some(inner) = t.as_typed_mut::<$ty>() {
74                        let tq = TQ::per_tensor(scale, zp);
75                        inner.set_quantization(tq).map_err(|e| {
76                            DecoderError::Internal(format!(
77                                "apply_schema_quant set_quantization failed: {e}"
78                            ))
79                        })?;
80                        return Ok(());
81                    }
82                };
83            }
84            try_attach!(i8);
85            try_attach!(u8);
86            try_attach!(i16);
87            try_attach!(u16);
88            // Float / other dtypes silently skipped.
89            return Ok(());
90        }
91    }
92    Err(DecoderError::InvalidShape(format!(
93        "apply_schema_quant: no tensor matches {expected_shape:?}"
94    )))
95}
96
97#[cfg(test)]
98mod tests {
99    use super::*;
100    use crate::schema::SchemaV2;
101    use edgefirst_tensor::{Tensor, TensorMemory};
102
103    #[test]
104    fn applies_quant_to_int8_by_shape() {
105        // Build a tiny schema with one logical output `boxes` having an int8 child.
106        let json =
107            edgefirst_bench::testdata::read_to_string("per_scale/synthetic_yolov8n_schema.json");
108        let schema: SchemaV2 = serde_json::from_str(&json).unwrap();
109
110        // Build a tensor matching one box child's shape.
111        // The yolov8n schema's first box child has shape [1, 80, 80, 64].
112        let t = Tensor::<i8>::new(&[1, 80, 80, 64], Some(TensorMemory::Mem), None).unwrap();
113        assert!(t.quantization().is_none(), "fresh tensor has no quant");
114        let mut td: TensorDyn = t.into();
115        let mut tensors: Vec<&mut TensorDyn> = vec![&mut td];
116
117        apply_schema_quant(&schema, &mut tensors).unwrap_err();
118        // Errors because not all schema-declared shapes have matching tensors.
119        // (We only provided one tensor, but schema has many children.)
120    }
121
122    #[test]
123    fn applies_quant_to_full_yolov8_input_set() {
124        let json =
125            edgefirst_bench::testdata::read_to_string("per_scale/synthetic_yolov8n_schema.json");
126        let schema: SchemaV2 = serde_json::from_str(&json).unwrap();
127
128        // Build the full set of input tensors matching every schema child + protos.
129        let shapes_int8 = [
130            vec![1, 80, 80, 64],
131            vec![1, 80, 80, 80],
132            vec![1, 80, 80, 32],
133            vec![1, 40, 40, 64],
134            vec![1, 40, 40, 80],
135            vec![1, 40, 40, 32],
136            vec![1, 20, 20, 64],
137            vec![1, 20, 20, 80],
138            vec![1, 20, 20, 32],
139            vec![1, 160, 160, 32],
140        ];
141        let mut owned: Vec<TensorDyn> = shapes_int8
142            .iter()
143            .map(|s| {
144                Tensor::<i8>::new(s, Some(TensorMemory::Mem), None)
145                    .unwrap()
146                    .into()
147            })
148            .collect();
149        let mut refs: Vec<&mut TensorDyn> = owned.iter_mut().collect();
150
151        apply_schema_quant(&schema, &mut refs).unwrap();
152
153        // All 10 tensors should now carry quant.
154        for td in &owned {
155            if let Some(t) = td.as_typed::<i8>() {
156                assert!(
157                    t.quantization().is_some(),
158                    "tensor missing quant after apply"
159                );
160            }
161        }
162    }
163
164    #[test]
165    fn empty_scale_errors_instead_of_silently_skipping() {
166        // Regression for Copilot review on PR #63: an empty `scale`
167        // vector used to slip past `is_per_tensor()` and return Ok,
168        // masking a malformed schema until decode time.
169        use crate::schema::{DType, Quantization};
170        let bad_q = Quantization {
171            scale: vec![],
172            zero_point: None,
173            axis: None,
174            dtype: Some(DType::Int8),
175        };
176        let mut td: TensorDyn = Tensor::<i8>::new(&[1, 2, 3], Some(TensorMemory::Mem), None)
177            .unwrap()
178            .into();
179        let mut refs: Vec<&mut TensorDyn> = vec![&mut td];
180
181        let err = attach_per_tensor_quant_by_shape(&mut refs, &[1, 2, 3], &bad_q)
182            .expect_err("empty scale must be rejected");
183        match err {
184            DecoderError::InvalidShape(msg) => {
185                assert!(msg.contains("`scale` is empty"), "msg = {msg}")
186            }
187            other => panic!("unexpected error: {other:?}"),
188        }
189    }
190
191    #[test]
192    fn skips_float_tensors_silently() {
193        let json =
194            edgefirst_bench::testdata::read_to_string("per_scale/synthetic_yolov8n_schema.json");
195        let schema: SchemaV2 = serde_json::from_str(&json).unwrap();
196
197        // Build float tensors instead of int8.
198        let shape = vec![1, 80, 80, 64];
199        let t = Tensor::<f32>::new(&shape, Some(TensorMemory::Mem), None).unwrap();
200        let mut td: TensorDyn = t.into();
201        let mut refs: Vec<&mut TensorDyn> = vec![&mut td];
202
203        // Float tensors don't take quant; the helper returns Ok and the schema
204        // declaration is silently skipped on this single mismatched-shape input.
205        // (Actually it'll error on InvalidShape for the OTHER schema children
206        // that have no matching tensor. That's the expected behaviour.)
207        let r = apply_schema_quant(&schema, &mut refs);
208        assert!(r.is_err());
209        // What we're really testing: float tensors don't crash.
210        assert_eq!(td.dtype(), edgefirst_tensor::DType::F32, "unexpected dtype");
211    }
212}
edgefirst_decoder/per_scale/helper.rs

edgefirst_decoder/per_scale/
helper.rs
