quantize_rs/onnx_utils/

graph_builder.rs

//! Graph-level operations for quantized ONNX models.
//!
//! Three responsibilities:
//!   1. **QDQ transform** — replace FP32 initializers with INT8 + DequantizeLinear
//!   2. **Connectivity validation** — walk the graph and verify every edge resolves
//!   3. **Opset management** — ensure the model declares a sufficient opset and
//!      upgrades deprecated op attributes when bumping across breaking opset boundaries

use crate::errors::{QuantizeError, Result};
use crate::onnx_proto::{
    attribute_proto, tensor_proto, AttributeProto, GraphProto, ModelProto, OperatorSetIdProto,
    TensorProto,
};
use std::collections::{HashMap, HashSet};

use super::quantization_nodes::{
    build_dequantize_linear_node, build_quantized_weight_tensor, build_scale_tensor,
    build_zero_point_tensor, DequantLinearNames, StorageFormat,
};

// ===========================================================================
// Public types
// ===========================================================================

/// One weight to convert: FP32 initializer → INT8 + DequantizeLinear block.
///
/// Fields are public for ergonomic construction in callers; new fields are
/// added only on a minor version bump and downstream `..Default::default()`
/// callers continue to compile.  Always construct via
/// `QdqWeightInput { /* fields */, ..Default::default() }`.
#[derive(Debug, Clone, Default)]
pub struct QdqWeightInput {
    /// Original initializer name (e.g., `"conv1.weight"`)
    pub original_name: String,
    /// Quantized values as i8.  For INT4 these are in [-8, 7]; for INT8 in [-128, 127].
    /// Always unpacked — one value per element.
    pub quantized_values: Vec<i8>,
    /// Quantization scales (FP32).
    /// Length 1 for per-tensor; one per channel for per-channel.
    pub scales: Vec<f32>,
    /// Zero points (INT8).
    /// Same length as `scales`.
    pub zero_points: Vec<i8>,
    /// Original bit-width (4 or 8).  Informational only — ONNX storage is always INT8.
    /// Persisted in model metadata so `load_quantized_info` can recover it.
    pub bits: u8,
    /// Per-channel quantization axis, or `None` for per-tensor.
    pub axis: Option<usize>,
}

/// Options controlling how a quantized model is written to disk.
///
/// Marked `#[non_exhaustive]` so future save options can be added without a
/// breaking change.  Construct with [`SaveOptions::default()`] and the
/// `.with_*()` builders.
#[derive(Debug, Clone, Copy, Default)]
#[non_exhaustive]
pub struct SaveOptions {
    /// Use native ONNX INT4 storage (opset 21, `DataType::Int4`) instead of
    /// widening INT4 values into INT8 bytes.
    ///
    /// Native storage halves the on-disk size of INT4-quantized weights
    /// (true 8× compression from FP32) but raises the required opset to 21.
    /// Older runtimes without opset 21 support will refuse to load the model.
    /// Defaults to `false` for backward compatibility.
    ///
    /// Only affects weights where [`QdqWeightInput::bits`] is 4; INT8 weights
    /// are stored as INT8 regardless of this flag.
    pub native_int4: bool,
}

impl SaveOptions {
    /// Opt in to native INT4 storage (requires opset 21).
    pub fn with_native_int4(mut self, enabled: bool) -> Self {
        self.native_int4 = enabled;
        self
    }
}

/// Result of a graph-connectivity check.
#[derive(Debug)]
#[must_use]
pub struct ConnectivityReport {
    /// `true` if every node input resolves to a known tensor.
    pub valid: bool,
    /// Human-readable description of every dangling reference.  Empty when valid.
    pub broken_refs: Vec<String>,
}

impl ConnectivityReport {
    /// Render the report as a printable string (useful for CLI `validate` output).
    pub fn summary(&self) -> String {
        if self.valid {
            "  Graph connectivity: OK\n".to_string()
        } else {
            let mut s = format!(
                "  Graph connectivity: BROKEN ({} dangling reference{})\n",
                self.broken_refs.len(),
                if self.broken_refs.len() == 1 { "" } else { "s" }
            );
            for (i, r) in self.broken_refs.iter().enumerate() {
                s.push_str(&format!("    {}. {}\n", i + 1, r));
            }
            s
        }
    }
}

// ===========================================================================
// Connectivity validation
// ===========================================================================

/// Walk the graph and verify every node input resolves to *something*.
///
/// A valid input is exactly one of:
///   • a declared graph input (`graph.input`)
///   • an initializer name (`graph.initializer`)
///   • the output of a node that appears **earlier** in `graph.node`
///
/// This is the check ONNX Runtime performs on load — and the check that
/// v0.2.0's `validate` command skipped, letting the rename bug through.
pub(crate) fn validate_graph_connectivity(graph: &GraphProto) -> ConnectivityReport {
    let mut known: HashSet<String> = HashSet::new();

    // Seed: graph inputs + initializers are always available
    for inp in &graph.input {
        known.insert(inp.name.clone());
    }
    for init in &graph.initializer {
        known.insert(init.name.clone());
    }

    let mut broken = Vec::new();

    // Walk nodes in serialized order; each node's outputs become known afterwards
    for node in &graph.node {
        for name in &node.input {
            if name.is_empty() {
                continue; // optional input slot — empty string is valid
            }
            if !known.contains(name.as_str()) {
                broken.push(format!(
                    "Node '{}' (op={}) → unknown input '{}'",
                    node.name, node.op_type, name
                ));
            }
        }
        // Register outputs so later nodes can consume them
        for name in &node.output {
            if !name.is_empty() {
                known.insert(name.clone());
            }
        }
    }

    ConnectivityReport {
        valid: broken.is_empty(),
        broken_refs: broken,
    }
}

// ===========================================================================
// Opset version management
// ===========================================================================

/// Ensure the default ONNX domain opset is at least `min_version`.
///
/// DequantizeLinear requires opset ≥ 10 (per-tensor) or ≥ 13 (per-channel axis).
///
/// When bumping the opset past a breaking boundary, this function also
/// upgrades deprecated op attributes:
///   - **opset 9**: `BatchNormalization.spatial` removed (was always 1)
///   - **opset 12**: `Dropout.ratio` migrated from attribute to input
pub(crate) fn ensure_opset_version(model: &mut ModelProto, min_version: i64) {
    let old_version = get_opset_version(model);

    // Update or insert the default-domain opset entry
    let mut found = false;
    for opset in model.opset_import.iter_mut() {
        if opset.domain.is_empty() {
            if opset.version < min_version {
                opset.version = min_version;
            }
            found = true;
            break;
        }
    }
    if !found {
        model.opset_import.push(OperatorSetIdProto {
            domain: String::new(),
            version: min_version,
        });
    }

    // Strip deprecated attributes when crossing breaking opset boundaries
    if old_version < min_version {
        if let Some(graph) = model.graph.as_mut() {
            // Warn about ops we *can't* auto-migrate before rewriting the ones
            // we can.  (`&mut graph` reborrows as `&graph` for the read.)
            let unhandled = unhandled_opset_migrations(graph, old_version, min_version);
            if !unhandled.is_empty() {
                warn_unhandled_opset_migrations(old_version, min_version, &unhandled);
            }
            upgrade_deprecated_ops(graph, old_version, min_version);
        }
    }

    // Keep `ir_version` consistent with the (possibly bumped) opset.  Native
    // INT4 emits opset 21, whose `INT4`/`UINT4` tensor types require IR ≥ 10; a
    // model that advertises opset 21 but an older IR is out of spec (current
    // ONNX Runtime tolerates it, stricter/older/future ones may not).  Only ever
    // raise — never lower an already-newer IR — and base it on the *effective*
    // opset so a model already above `min_version` is still covered.
    let effective_opset = old_version.max(min_version);
    let min_ir = min_ir_version_for_opset(effective_opset);
    if model.ir_version < min_ir {
        model.ir_version = min_ir;
    }
}

/// Get the current default-domain opset version (0 if not present).
fn get_opset_version(model: &ModelProto) -> i64 {
    model
        .opset_import
        .iter()
        .find(|o| o.domain.is_empty())
        .map_or(0, |o| o.version)
}

/// Minimum ONNX IR version required to declare a given default-domain opset,
/// per the ONNX spec's opset↔IR mapping (`docs/Versioning.md`).
///
/// This matters because bumping the opset alone is not enough: a model that
/// advertises opset 21 (which is where the `INT4`/`UINT4` tensor types and
/// native-INT4 `DequantizeLinear` were introduced) but keeps an older
/// `ir_version` is out of spec.  Lenient runtimes (current ONNX Runtime)
/// accept it, but stricter or future ones may not.  [`ensure_opset_version`]
/// uses this to raise `ir_version` in lock-step with the opset it sets.
fn min_ir_version_for_opset(opset: i64) -> i64 {
    match opset {
        ..=8 => 3,
        9 => 4,
        10 => 5,
        11 => 6,
        12..=14 => 7,
        15..=18 => 8,
        19..=20 => 9,
        21..=22 => 10,
        _ => 11, // opset 23+ → IR 11 (safe floor for anything newer)
    }
}

/// Operators whose schema changed in a backward-incompatible way at the given
/// opset (an attribute became an input, or the op was removed/renamed) and that
/// [`upgrade_deprecated_ops`] does **not** rewrite.
///
/// When a save bumps the declared opset across one of these boundaries, the node
/// keeps its old-opset form but the model now advertises the new opset — which
/// ONNX Runtime may reject on load.  Auto-migrating them correctly needs ONNX's
/// full version converter, so quantize-rs warns instead of silently emitting a
/// model that might not load.
///
/// `BatchNormalization`, `Dropout`, `Softmax`, and `LogSoftmax` are intentionally
/// absent — [`upgrade_deprecated_ops`] already handles those.
const UNHANDLED_OPSET_MIGRATIONS: &[(&str, i64)] = &[
    ("Upsample", 10),  // removed in opset 10 (replaced by Resize)
    ("Slice", 10),     // starts/ends/axes/steps: attribute → input
    ("TopK", 10),      // K: attribute → input
    ("Resize", 11),    // roi/scales/sizes signature change
    ("Pad", 11),       // pads/value: attribute → input
    ("Clip", 11),      // min/max: attribute → input
    ("Scatter", 11),   // deprecated → ScatterElements
    ("Split", 13),     // split: attribute → input
    ("Squeeze", 13),   // axes: attribute → input
    ("Unsqueeze", 13), // axes: attribute → input
    ("ReduceSum", 13), // axes: attribute → input
];

/// Operators present in `graph` whose unhandled migration boundary falls inside
/// `(old_opset, new_opset]`.  Returns `(op_type, boundary_opset)` pairs, each op
/// at most once.  Pure and testable; the user-facing warning is emitted by
/// [`warn_unhandled_opset_migrations`].
fn unhandled_opset_migrations(
    graph: &GraphProto,
    old_opset: i64,
    new_opset: i64,
) -> Vec<(&'static str, i64)> {
    let mut hits: Vec<(&'static str, i64)> = Vec::new();
    for &(op, boundary) in UNHANDLED_OPSET_MIGRATIONS {
        if old_opset < boundary
            && new_opset >= boundary
            && graph.node.iter().any(|n| n.op_type == op)
            && !hits.iter().any(|(o, _)| *o == op)
        {
            hits.push((op, boundary));
        }
    }
    hits
}

/// Emit a single stderr warning describing operators that may break when the
/// declared opset is bumped (see [`unhandled_opset_migrations`]).
fn warn_unhandled_opset_migrations(old: i64, new: i64, hits: &[(&str, i64)]) {
    let list = hits
        .iter()
        .map(|&(op, b)| format!("{op} (changed at opset {b})"))
        .collect::<Vec<_>>()
        .join(", ");
    log::warn!(
        "bumping the model opset from {old} to {new} to emit DequantizeLinear, but the graph \
         contains operator(s) whose schema changed at an opset in that range and which \
         quantize-rs does not rewrite: {list}.  The saved model will declare opset {new} while \
         keeping the older node form, so ONNX Runtime may reject it.  If it does, run the model \
         through onnx.version_converter (or re-export at opset {new}+) before quantizing."
    );
}

/// Upgrade graph nodes whose attribute semantics changed between `old_opset`
/// and `new_opset`.
///
/// - **BatchNormalization** (opset 9): `spatial` attribute removed (was always 1)
/// - **Dropout** (opset 12): `ratio` attribute → 2nd input
/// - **Softmax / LogSoftmax** (opset 13): default axis changed from 1 to -1
fn upgrade_deprecated_ops(graph: &mut GraphProto, old_opset: i64, new_opset: i64) {
    let mut new_initializers: Vec<TensorProto> = Vec::new();

    for (node_idx, node) in graph.node.iter_mut().enumerate() {
        // Opset 9: BatchNormalization removed the `spatial` attribute.
        // It was always 1 (the only valid value) and had no effect.
        if node.op_type == "BatchNormalization" && old_opset < 9 && new_opset >= 9 {
            node.attribute.retain(|a| a.name != "spatial");
        }

        // Opset 12: Dropout `ratio` moved from attribute to 2nd input.
        // Extract the value, remove the attribute, and wire in a constant.
        if node.op_type == "Dropout" && old_opset < 12 && new_opset >= 12 {
            let ratio = node
                .attribute
                .iter()
                .find(|a| a.name == "ratio")
                .map(|a| a.f)
                .unwrap_or(0.5);
            node.attribute.retain(|a| a.name != "ratio");

            // Include the node index so two Dropouts with the same
            // (or empty) first output don't collide on the initializer name.
            let init_name = format!(
                "_quantize_rs_dropout_ratio_{}_{}",
                node_idx,
                node.output.first().map_or("", |s| s.as_str()),
            );
            new_initializers.push(TensorProto {
                name: init_name.clone(),
                data_type: tensor_proto::DataType::Float as i32,
                float_data: vec![ratio],
                ..Default::default()
            });

            if node.input.len() < 2 {
                node.input.push(init_name);
            } else {
                node.input[1] = init_name;
            }
        }

        // Opset 13: Softmax/LogSoftmax default axis changed from 1 to -1.
        // If the node has no explicit axis attribute, add axis=1 to preserve
        // the old behavior.
        if (node.op_type == "Softmax" || node.op_type == "LogSoftmax")
            && old_opset < 13
            && new_opset >= 13
        {
            let has_axis = node.attribute.iter().any(|a| a.name == "axis");
            if !has_axis {
                node.attribute.push(AttributeProto {
                    name: "axis".to_string(),
                    r#type: attribute_proto::AttributeType::Int as i32,
                    i: 1, // old default
                    ..Default::default()
                });
            }
        }
    }

    graph.initializer.extend(new_initializers);
}

// ===========================================================================
// QDQ transform
// ===========================================================================

/// Replace FP32 weight initializers with INT8 quantized equivalents +
/// DequantizeLinear nodes.
///
/// ### What happens per weight in `inputs`:
///
/// **Removed:**
///   - Initializer `"{name}"` (the original FP32 weight data)
///
/// **Added (initializers):**
///   - `"{name}_quantized"` — INT8, same shape as original
///   - `"{name}_scale"`     — FP32 scalar
///   - `"{name}_zp"`        — INT8 scalar
///
/// **Added (node, prepended before all existing nodes):**
///   - `DequantizeLinear` with output = `"{name}"`
///
/// Because the DequantizeLinear output carries the **original** name, every
/// downstream node (Conv, MatMul, BatchNorm, …) remains completely unchanged.
/// Graph connectivity is preserved by construction.
///
/// ---
/// ### INT4 storage note
///
/// ONNX `DequantizeLinear` requires INT8 input in opsets < 21.  By default,
/// INT4-quantized values (range [-8, 7]) are widened to INT8 here — 4×
/// compression from FP32.  To get the full 8× compression, pass
/// [`SaveOptions::with_native_int4`]`(true)` to [`apply_qdq_transform_with_options`];
/// that emits native `DataType::Int4` (opset 21) with two values packed per byte.
#[cfg(test)]
pub(crate) fn apply_qdq_transform(graph: &mut GraphProto, inputs: &[QdqWeightInput]) -> Result<()> {
    apply_qdq_transform_with_options(graph, inputs, SaveOptions::default())
}

/// Same as [`apply_qdq_transform`] but configurable via [`SaveOptions`].
///
/// Prefer this entry point for any new code; the short-form wrapper exists
/// only for backward compatibility.
pub(crate) fn apply_qdq_transform_with_options(
    graph: &mut GraphProto,
    inputs: &[QdqWeightInput],
    options: SaveOptions,
) -> Result<()> {
    // -----------------------------------------------------------------------
    // 0.  Snapshot shapes before modifying the initializer list
    // -----------------------------------------------------------------------
    let shape_map: HashMap<String, Vec<i64>> = graph
        .initializer
        .iter()
        .map(|init| (init.name.clone(), init.dims.clone()))
        .collect();

    let quant_set: HashSet<&str> = inputs.iter().map(|i| i.original_name.as_str()).collect();

    // -----------------------------------------------------------------------
    // 0b.  Pre-flight: verify every requested weight still exists as an FP32
    //      initializer.  Failing here keeps the graph in a consistent state
    //      so the caller can retry; without this, a missing input would
    //      surface only after some initializers were already removed.
    //      Also catches the "already quantized — second save call" case
    //      with a helpful diagnostic.
    // -----------------------------------------------------------------------
    for inp in inputs {
        if !shape_map.contains_key(&inp.original_name) {
            let already_quantized = shape_map
                .keys()
                .any(|n| n == &format!("{}_quantized", inp.original_name));
            return Err(QuantizeError::GraphTransform {
                reason: if already_quantized {
                    format!(
                        "Weight '{}' has already been quantized in this graph \
                         (found '{}_quantized' initializer); apply_qdq_transform \
                         is not idempotent — load a fresh OnnxModel before retrying",
                        inp.original_name, inp.original_name
                    )
                } else {
                    format!(
                        "Weight '{}' not found in model initializers — \
                         verify the name matches exactly",
                        inp.original_name
                    )
                },
            });
        }
    }

    // -----------------------------------------------------------------------
    // 1.  Remove the original FP32 initializers for every weight we're replacing
    // -----------------------------------------------------------------------
    graph
        .initializer
        .retain(|init| !quant_set.contains(init.name.as_str()));

    // -----------------------------------------------------------------------
    // 1b. Also remove weights from graph.input (critical fix for "Duplicate definition")
    // -----------------------------------------------------------------------
    // Some ONNX models list weights as both initializers AND graph inputs.
    // This is valid ONNX, but when DequantizeLinear outputs reuse the original
    // weight names, ONNX Runtime sees two definitions of the same tensor.
    graph
        .input
        .retain(|inp| !quant_set.contains(inp.name.as_str()));

    // -----------------------------------------------------------------------
    // 2.  Add quantized initializer triples + build DequantizeLinear nodes
    // -----------------------------------------------------------------------
    let mut dq_nodes = Vec::new();

    for inp in inputs {
        // The pre-flight loop above already verified every original_name is
        // present, and shape_map is not mutated in between — so this is
        // unreachable in practice. Returning a clean error rather than
        // panicking means a future refactor that breaks the invariant degrades
        // to a recoverable `GraphTransform` error instead of a crash.
        let shape =
            shape_map
                .get(&inp.original_name)
                .ok_or_else(|| QuantizeError::GraphTransform {
                    reason: format!(
                        "internal invariant violated: weight '{}' missing from shape_map \
                     after pre-flight validation",
                        inp.original_name
                    ),
                })?;

        let expected_len: i64 = shape.iter().product();
        if inp.quantized_values.len() as i64 != expected_len {
            return Err(QuantizeError::GraphTransform {
                reason: format!(
                    "Weight '{}': quantized_values has {} elements but shape {:?} expects {}",
                    inp.original_name,
                    inp.quantized_values.len(),
                    shape,
                    expected_len
                ),
            });
        }

        let names = DequantLinearNames::from_original(&inp.original_name);

        let format = if options.native_int4 && inp.bits == 4 {
            StorageFormat::NativeInt4
        } else {
            StorageFormat::Int8Widened
        };

        graph.initializer.push(build_quantized_weight_tensor(
            &names,
            &inp.quantized_values,
            shape,
            format,
        ));
        graph
            .initializer
            .push(build_scale_tensor(&names, &inp.scales));
        graph
            .initializer
            .push(build_zero_point_tensor(&names, &inp.zero_points, format));

        dq_nodes.push(build_dequantize_linear_node(&names, inp.axis));
    }

    // -----------------------------------------------------------------------
    // 3.  Prepend DequantizeLinear nodes before all existing computation nodes.
    //     They must appear first so their outputs are "known" when the validator
    //     (or ONNX Runtime) walks the node list in order.
    // -----------------------------------------------------------------------
    let existing_nodes = std::mem::take(&mut graph.node);
    graph.node = dq_nodes;
    graph.node.extend(existing_nodes);

    Ok(())
}

// ===========================================================================
// Tests
// ===========================================================================

#[cfg(test)]
mod tests {
    use super::*;
    use crate::onnx_proto::{
        tensor_proto, GraphProto, ModelProto, NodeProto, OperatorSetIdProto, TensorProto,
        ValueInfoProto,
    };

    // -----------------------------------------------------------------------
    // Test helpers
    // -----------------------------------------------------------------------

    /// Minimal graph: one graph input "input", one FP32 initializer "w" (shape [2,2]),
    /// one Conv node consuming both, producing "out".
    fn make_simple_graph() -> GraphProto {
        GraphProto {
            input: vec![ValueInfoProto {
                name: "input".to_string(),
                ..Default::default()
            }],
            initializer: vec![TensorProto {
                name: "w".to_string(),
                data_type: tensor_proto::DataType::Float as i32,
                dims: vec![2, 2],
                float_data: vec![1.0, 2.0, 3.0, 4.0],
                ..Default::default()
            }],
            node: vec![NodeProto {
                op_type: "Conv".to_string(),
                name: "conv0".to_string(),
                input: vec!["input".to_string(), "w".to_string()],
                output: vec!["out".to_string()],
                ..Default::default()
            }],
            ..Default::default()
        }
    }

    /// Two-weight graph: "w1" and "w2", two Conv nodes chained.
    fn make_two_weight_graph() -> GraphProto {
        GraphProto {
            input: vec![ValueInfoProto {
                name: "input".to_string(),
                ..Default::default()
            }],
            initializer: vec![
                TensorProto {
                    name: "w1".to_string(),
                    data_type: tensor_proto::DataType::Float as i32,
                    dims: vec![2, 2],
                    float_data: vec![1.0, 2.0, 3.0, 4.0],
                    ..Default::default()
                },
                TensorProto {
                    name: "w2".to_string(),
                    data_type: tensor_proto::DataType::Float as i32,
                    dims: vec![2, 2],
                    float_data: vec![5.0, 6.0, 7.0, 8.0],
                    ..Default::default()
                },
            ],
            node: vec![
                NodeProto {
                    op_type: "Conv".to_string(),
                    name: "conv1".to_string(),
                    input: vec!["input".to_string(), "w1".to_string()],
                    output: vec!["mid".to_string()],
                    ..Default::default()
                },
                NodeProto {
                    op_type: "Conv".to_string(),
                    name: "conv2".to_string(),
                    input: vec!["mid".to_string(), "w2".to_string()],
                    output: vec!["out".to_string()],
                    ..Default::default()
                },
            ],
            ..Default::default()
        }
    }

    // -----------------------------------------------------------------------
    // Connectivity validation tests
    // -----------------------------------------------------------------------

    #[test]
    fn test_connectivity_passes_on_valid_graph() {
        let graph = make_simple_graph();
        let report = validate_graph_connectivity(&graph);
        assert!(
            report.valid,
            "original graph should be valid; broken: {:?}",
            report.broken_refs
        );
    }

    #[test]
    fn test_connectivity_detects_renamed_initializer() {
        // Simulate the exact v0.2.0 bug: rename "w" in the initializer list
        // without updating the Conv node that references it.
        let mut graph = make_simple_graph();

        for init in graph.initializer.iter_mut() {
            if init.name == "w" {
                init.name = "w__qINT8_s0.00392_z-3_len4".to_string();
            }
        }

        let report = validate_graph_connectivity(&graph);
        assert!(!report.valid, "should detect broken reference to 'w'");
        assert_eq!(report.broken_refs.len(), 1);
        assert!(
            report.broken_refs[0].contains("'w'"),
            "error should mention 'w': {}",
            report.broken_refs[0]
        );
    }

    #[test]
    fn test_connectivity_detects_multiple_broken_refs() {
        let mut graph = make_two_weight_graph();

        for init in graph.initializer.iter_mut() {
            if init.name == "w1" {
                init.name = "w1_broken".to_string();
            } else if init.name == "w2" {
                init.name = "w2_broken".to_string();
            }
        }

        let report = validate_graph_connectivity(&graph);
        assert!(!report.valid);
        assert_eq!(report.broken_refs.len(), 2);
    }

    #[test]
    fn test_connectivity_summary_formatting() {
        let valid = ConnectivityReport {
            valid: true,
            broken_refs: vec![],
        };
        assert!(valid.summary().contains("OK"));

        let broken = ConnectivityReport {
            valid: false,
            broken_refs: vec!["Node 'x' → unknown input 'y'".to_string()],
        };
        let s = broken.summary();
        assert!(s.contains("BROKEN"));
        assert!(s.contains("1 dangling reference"));
        assert!(s.contains("unknown input 'y'"));
    }

    // -----------------------------------------------------------------------
    // Opset version tests
    // -----------------------------------------------------------------------

    #[test]
    fn test_ensure_opset_bumps_low_version() {
        let mut model = ModelProto {
            opset_import: vec![OperatorSetIdProto {
                domain: String::new(),
                version: 10,
            }],
            ..Default::default()
        };

        ensure_opset_version(&mut model, 13);

        assert_eq!(model.opset_import[0].version, 13);
    }

    #[test]
    fn test_ensure_opset_leaves_sufficient_version() {
        let mut model = ModelProto {
            opset_import: vec![OperatorSetIdProto {
                domain: String::new(),
                version: 17,
            }],
            ..Default::default()
        };

        ensure_opset_version(&mut model, 13);

        assert_eq!(model.opset_import[0].version, 17, "should not downgrade");
    }

    #[test]
    fn test_ensure_opset_adds_missing_default_domain() {
        let mut model = ModelProto::default();
        // No opset_import at all
        ensure_opset_version(&mut model, 13);

        assert_eq!(model.opset_import.len(), 1);
        assert!(model.opset_import[0].domain.is_empty());
        assert_eq!(model.opset_import[0].version, 13);
    }

    // -----------------------------------------------------------------------
    // QDQ transform tests
    // -----------------------------------------------------------------------

    #[test]
    fn test_qdq_single_weight_produces_valid_graph() {
        let mut graph = make_simple_graph();

        let inputs = vec![QdqWeightInput {
            original_name: "w".to_string(),
            quantized_values: vec![25, 51, 76, 102],
            scales: vec![0.039_215_686], // ≈ 1/25.5
            zero_points: vec![0],
            bits: 8,
            axis: None,
        }];

        apply_qdq_transform(&mut graph, &inputs).expect("QDQ transform failed");

        let report = validate_graph_connectivity(&graph);
        assert!(
            report.valid,
            "graph after QDQ must be valid; broken: {:?}",
            report.broken_refs
        );
    }

    #[test]
    fn test_qdq_adds_correct_initializers() {
        let mut graph = make_simple_graph();

        let inputs = vec![QdqWeightInput {
            original_name: "w".to_string(),
            quantized_values: vec![10, 20, 30, 40],
            scales: vec![0.1],
            zero_points: vec![-5],
            bits: 8,
            axis: None,
        }];

        apply_qdq_transform(&mut graph, &inputs).expect("QDQ transform failed");

        let init_names: Vec<&str> = graph.initializer.iter().map(|i| i.name.as_str()).collect();

        assert!(init_names.contains(&"w_quantized"), "missing w_quantized");
        assert!(init_names.contains(&"w_scale"), "missing w_scale");
        assert!(init_names.contains(&"w_zp"), "missing w_zp");
        assert!(
            !init_names.contains(&"w"),
            "original FP32 'w' should be removed"
        );
    }

    #[test]
    fn test_qdq_node_order_dequant_first() {
        let mut graph = make_simple_graph();

        let inputs = vec![QdqWeightInput {
            original_name: "w".to_string(),
            quantized_values: vec![10, 20, 30, 40],
            scales: vec![0.1],
            zero_points: vec![0],
            bits: 8,
            axis: None,
        }];

        apply_qdq_transform(&mut graph, &inputs).expect("QDQ transform failed");

        let ops: Vec<&str> = graph.node.iter().map(|n| n.op_type.as_str()).collect();

        assert_eq!(ops.len(), 2);
        assert_eq!(ops[0], "DequantizeLinear");
        assert_eq!(ops[1], "Conv");
    }

    #[test]
    fn test_qdq_dequant_output_is_original_name() {
        let mut graph = make_simple_graph();

        let inputs = vec![QdqWeightInput {
            original_name: "w".to_string(),
            quantized_values: vec![1, 2, 3, 4],
            scales: vec![1.0],
            zero_points: vec![0],
            bits: 8,
            axis: None,
        }];

        apply_qdq_transform(&mut graph, &inputs).expect("QDQ transform failed");

        let dq = &graph.node[0]; // first node = DequantizeLinear
        assert_eq!(
            dq.output[0], "w",
            "DequantizeLinear output must be original name"
        );
    }

    #[test]
    fn test_qdq_two_weights_both_transformed() {
        let mut graph = make_two_weight_graph();

        let inputs = vec![
            QdqWeightInput {
                original_name: "w1".to_string(),
                quantized_values: vec![10, 20, 30, 40],
                scales: vec![0.1],
                zero_points: vec![0],
                bits: 8,
                axis: None,
            },
            QdqWeightInput {
                original_name: "w2".to_string(),
                quantized_values: vec![50, 60, 70, 80],
                scales: vec![0.2],
                zero_points: vec![-1],
                bits: 8,
                axis: None,
            },
        ];

        apply_qdq_transform(&mut graph, &inputs).expect("QDQ transform failed");

        // Connectivity must still be valid
        let report = validate_graph_connectivity(&graph);
        assert!(
            report.valid,
            "two-weight graph broken: {:?}",
            report.broken_refs
        );

        // Should have 2 DequantizeLinear + 2 Conv = 4 nodes
        assert_eq!(graph.node.len(), 4);

        // First two nodes are DequantizeLinear
        assert_eq!(graph.node[0].op_type, "DequantizeLinear");
        assert_eq!(graph.node[1].op_type, "DequantizeLinear");

        // Their outputs are the original weight names
        let dq_outputs: Vec<&str> = graph
            .node
            .iter()
            .take(2)
            .map(|n| n.output[0].as_str())
            .collect();
        assert!(dq_outputs.contains(&"w1"));
        assert!(dq_outputs.contains(&"w2"));
    }

    #[test]
    fn test_qdq_int4_values_stored_as_int8() {
        let mut graph = make_simple_graph();

        // INT4 range [-8, 7] — these arrive as i8 from ensure_unpacked()
        let inputs = vec![QdqWeightInput {
            original_name: "w".to_string(),
            quantized_values: vec![-8, -1, 0, 7],
            scales: vec![0.5],
            zero_points: vec![0],
            bits: 4, // flag says INT4; storage must still be INT8
            axis: None,
        }];

        apply_qdq_transform(&mut graph, &inputs).expect("QDQ transform failed");

        let quant_init = graph
            .initializer
            .iter()
            .find(|i| i.name == "w_quantized")
            .expect("w_quantized not found");

        // Data type must be INT8 (ONNX DequantizeLinear requirement)
        assert_eq!(quant_init.data_type, tensor_proto::DataType::Int8 as i32);

        // Byte-level round-trip must be exact
        let recovered: Vec<i8> = quant_init.raw_data.iter().map(|&b| b as i8).collect();
        assert_eq!(recovered, vec![-8, -1, 0, 7]);
    }

    #[test]
    fn test_qdq_unknown_weight_returns_error() {
        let mut graph = make_simple_graph();

        let inputs = vec![QdqWeightInput {
            original_name: "does_not_exist".to_string(),
            quantized_values: vec![1, 2, 3],
            scales: vec![1.0],
            zero_points: vec![0],
            bits: 8,
            axis: None,
        }];

        let result = apply_qdq_transform(&mut graph, &inputs);
        assert!(result.is_err());
        assert!(
            result.unwrap_err().to_string().contains("does_not_exist"),
            "error should name the missing weight"
        );
    }

    #[test]
    fn test_qdq_non_quantized_initializers_preserved() {
        // Add an extra initializer "bias" that is NOT being quantized.
        // It must survive the transform untouched.
        let mut graph = make_simple_graph();

        graph.initializer.push(TensorProto {
            name: "bias".to_string(),
            data_type: tensor_proto::DataType::Float as i32,
            dims: vec![2],
            float_data: vec![0.1, 0.2],
            ..Default::default()
        });

        // Also add "bias" as a Conv input so connectivity stays valid
        graph.node[0].input.push("bias".to_string());

        let inputs = vec![QdqWeightInput {
            original_name: "w".to_string(),
            quantized_values: vec![10, 20, 30, 40],
            scales: vec![0.1],
            zero_points: vec![0],
            bits: 8,
            axis: None,
        }];

        apply_qdq_transform(&mut graph, &inputs).expect("QDQ transform failed");

        // "bias" must still be present and untouched
        let bias_init = graph.initializer.iter().find(|i| i.name == "bias");

        assert!(
            bias_init.is_some(),
            "non-quantized 'bias' initializer must be preserved"
        );
        assert!((bias_init.unwrap().float_data[0] - 0.1).abs() < 1e-6);

        // Full connectivity check
        let report = validate_graph_connectivity(&graph);
        assert!(report.valid, "broken: {:?}", report.broken_refs);
    }

    #[test]
    fn test_ensure_opset_strips_deprecated_attrs() {
        use crate::onnx_proto::NodeProto;

        let mut model = ModelProto {
            opset_import: vec![OperatorSetIdProto {
                domain: String::new(),
                version: 7,
            }],
            graph: Some(GraphProto {
                node: vec![
                    // BatchNormalization with deprecated `spatial` attr (removed opset 9)
                    NodeProto {
                        op_type: "BatchNormalization".to_string(),
                        input: vec!["x".into(), "s".into(), "b".into(), "m".into(), "v".into()],
                        output: vec!["bn_out".into()],
                        attribute: vec![
                            AttributeProto {
                                name: "epsilon".to_string(),
                                r#type: attribute_proto::AttributeType::Float as i32,
                                f: 1e-5,
                                ..Default::default()
                            },
                            AttributeProto {
                                name: "spatial".to_string(),
                                r#type: attribute_proto::AttributeType::Int as i32,
                                i: 1,
                                ..Default::default()
                            },
                        ],
                        ..Default::default()
                    },
                    // Dropout with `ratio` attribute — should be migrated
                    // from attribute to 2nd input.
                    NodeProto {
                        op_type: "Dropout".to_string(),
                        input: vec!["bn_out".into()],
                        output: vec!["drop_out".into(), "drop_mask".into()],
                        attribute: vec![AttributeProto {
                            name: "ratio".to_string(),
                            r#type: attribute_proto::AttributeType::Float as i32,
                            f: 0.3,
                            ..Default::default()
                        }],
                        ..Default::default()
                    },
                    // Softmax with NO axis attribute — old default is 1,
                    // opset 13 changes default to -1, so axis=1 must be added.
                    NodeProto {
                        op_type: "Softmax".to_string(),
                        input: vec!["drop_out".into()],
                        output: vec!["sm_out".into()],
                        attribute: vec![],
                        ..Default::default()
                    },
                ],
                ..Default::default()
            }),
            ..Default::default()
        };

        ensure_opset_version(&mut model, 13);

        // Opset should be 13
        let opset = model
            .opset_import
            .iter()
            .find(|o| o.domain.is_empty())
            .unwrap();
        assert_eq!(opset.version, 13);

        let graph = model.graph.as_ref().unwrap();

        // BatchNormalization: `spatial` must be removed, `epsilon` kept
        let bn = &graph.node[0];
        assert!(
            !bn.attribute.iter().any(|a| a.name == "spatial"),
            "BatchNormalization.spatial should be stripped"
        );
        assert!(
            bn.attribute.iter().any(|a| a.name == "epsilon"),
            "BatchNormalization.epsilon should be preserved"
        );

        // Dropout: `ratio` attribute must be removed and moved to 2nd input
        let drop = &graph.node[1];
        assert!(
            !drop.attribute.iter().any(|a| a.name == "ratio"),
            "Dropout.ratio attribute should be removed"
        );
        assert_eq!(drop.input.len(), 2, "Dropout should now have 2 inputs");
        let ratio_init_name = &drop.input[1];

        // The ratio value should be stored as an initializer
        let ratio_init = graph
            .initializer
            .iter()
            .find(|i| &i.name == ratio_init_name)
            .expect("Dropout ratio initializer should exist");
        assert_eq!(ratio_init.data_type, tensor_proto::DataType::Float as i32);
        assert!(
            (ratio_init.float_data[0] - 0.3).abs() < 1e-6,
            "ratio should be 0.3"
        );

        // Softmax: must have explicit axis=1 added (old default, since opset 13 changed it to -1)
        let sm = &graph.node[2];
        assert_eq!(sm.op_type, "Softmax");
        let axis_attr = sm
            .attribute
            .iter()
            .find(|a| a.name == "axis")
            .expect("Softmax should have axis attribute added");
        assert_eq!(axis_attr.i, 1, "Softmax axis should be 1 (old default)");
    }

    #[test]
    fn test_ensure_opset_no_downgrade() {
        let mut model = ModelProto {
            opset_import: vec![OperatorSetIdProto {
                domain: String::new(),
                version: 15,
            }],
            graph: Some(GraphProto::default()),
            ..Default::default()
        };

        // Requesting 10 should NOT downgrade from 15
        ensure_opset_version(&mut model, 10);
        let opset = model
            .opset_import
            .iter()
            .find(|o| o.domain.is_empty())
            .unwrap();
        assert_eq!(opset.version, 15);
    }

    #[test]
    fn test_unhandled_opset_migration_detection() {
        let slice_graph = GraphProto {
            node: vec![NodeProto {
                op_type: "Slice".to_string(),
                ..Default::default()
            }],
            ..Default::default()
        };

        // Crossing 9 → 13 passes the opset-10 Slice boundary → flagged.
        assert_eq!(
            unhandled_opset_migrations(&slice_graph, 9, 13),
            vec![("Slice", 10)]
        );

        // Already at/above the boundary (11 → 13) → not flagged.
        assert!(unhandled_opset_migrations(&slice_graph, 11, 13).is_empty());

        // Not crossing far enough (8 → 9, boundary is 10) → not flagged.
        assert!(unhandled_opset_migrations(&slice_graph, 8, 9).is_empty());

        // Ops that `upgrade_deprecated_ops` handles are never reported.
        let softmax_graph = GraphProto {
            node: vec![NodeProto {
                op_type: "Softmax".to_string(),
                ..Default::default()
            }],
            ..Default::default()
        };
        assert!(unhandled_opset_migrations(&softmax_graph, 7, 13).is_empty());

        // Each affected op is reported once even with multiple boundaries crossed.
        let multi = GraphProto {
            node: vec![
                NodeProto {
                    op_type: "Pad".to_string(),
                    ..Default::default()
                },
                NodeProto {
                    op_type: "Unsqueeze".to_string(),
                    ..Default::default()
                },
            ],
            ..Default::default()
        };
        let hits = unhandled_opset_migrations(&multi, 7, 21);
        assert_eq!(hits.len(), 2);
        assert!(hits.contains(&("Pad", 11)));
        assert!(hits.contains(&("Unsqueeze", 13)));
    }

    #[test]
    fn test_ensure_opset_bumps_ir_version_for_native_int4() {
        // Native INT4 emits opset 21, whose INT4/UINT4 tensor types require
        // IR >= 10.  Bumping the opset must drag ir_version along.
        let mut model = ModelProto {
            ir_version: 8,
            opset_import: vec![OperatorSetIdProto {
                domain: String::new(),
                version: 13,
            }],
            graph: Some(GraphProto::default()),
            ..Default::default()
        };
        ensure_opset_version(&mut model, 21);
        assert_eq!(model.opset_import[0].version, 21);
        assert!(
            model.ir_version >= 10,
            "ir_version must be raised to >= 10 for opset 21, got {}",
            model.ir_version
        );
    }

    #[test]
    fn test_ensure_opset_never_lowers_ir_version() {
        let mut model = ModelProto {
            ir_version: 10,
            opset_import: vec![OperatorSetIdProto {
                domain: String::new(),
                version: 17,
            }],
            graph: Some(GraphProto::default()),
            ..Default::default()
        };
        // Requesting a lower opset must not touch the already-newer IR.
        ensure_opset_version(&mut model, 13);
        assert_eq!(model.ir_version, 10, "must not lower ir_version");
    }
}