onnx-ir 0.21.0

/// Tests for ONNX external data support
///
/// These tests verify that models with tensor data stored in external files
/// (used for models >2GB) are parsed correctly.
///
/// Test fixtures are generated by `scripts/external_data.py`
mod test_utils;

use test_utils::*;

/// Test loading a model with all weights stored externally
#[test]
fn test_external_data_basic() {
    // This model has weight [4x4] and bias [4] stored in external_data.bin
    // The model should parse without errors
    let graph = load_onnx("external_data.onnx");

    // Verify graph structure
    assert_eq!(graph.inputs.len(), 1, "Expected 1 input");
    assert_eq!(graph.outputs.len(), 1, "Expected 1 output");

    // The graph should have nodes (may be fused/optimized)
    assert!(!graph.nodes.is_empty(), "Graph should have nodes");
}

/// Test that external data tensors have correct values
#[test]
fn test_external_data_values() {
    let graph = load_onnx("external_data.onnx");

    // Find the constant nodes and verify their values
    for node in &graph.nodes {
        if let onnx_ir::ir::Node::Constant(const_node) = node {
            // Get the tensor data from the constant's output
            if let Some(output) = const_node.outputs.first()
                && let Some(data) = output.value()
            {
                let shape = data.shape.clone();

                // Weight tensor should be [4, 4] with diagonal values [1, 2, 3, 4]
                if shape == burn_tensor::Shape::new([4, 4]) {
                    let values: Vec<f32> = data.as_slice().unwrap().to_vec();
                    // Check diagonal elements
                    assert_eq!(values[0], 1.0, "weight[0,0] should be 1.0");
                    assert_eq!(values[5], 2.0, "weight[1,1] should be 2.0");
                    assert_eq!(values[10], 3.0, "weight[2,2] should be 3.0");
                    assert_eq!(values[15], 4.0, "weight[3,3] should be 4.0");
                }

                // Bias tensor should be [4] with values [0.1, 0.2, 0.3, 0.4]
                if shape == burn_tensor::Shape::new([4]) {
                    let values: Vec<f32> = data.as_slice().unwrap().to_vec();
                    assert!((values[0] - 0.1).abs() < 1e-6, "bias[0] should be 0.1");
                    assert!((values[1] - 0.2).abs() < 1e-6, "bias[1] should be 0.2");
                    assert!((values[2] - 0.3).abs() < 1e-6, "bias[2] should be 0.3");
                    assert!((values[3] - 0.4).abs() < 1e-6, "bias[3] should be 0.4");
                }
            }
        }
    }
}

/// Test loading a model with external data at non-zero offset
#[test]
fn test_external_data_with_offset() {
    let graph = load_onnx("external_data_offset.onnx");

    // Verify graph structure
    assert_eq!(graph.inputs.len(), 1, "Expected 1 input");
    assert_eq!(graph.outputs.len(), 1, "Expected 1 output");

    // Should have Add operation
    assert!(has_node_type(&graph, |n| matches!(
        n,
        onnx_ir::ir::Node::Add { .. }
    )));

    // Should have 1 constant node
    let const_count = count_constant_nodes(&graph);
    assert_eq!(const_count, 1, "Expected 1 constant node");

    // Verify the constant values
    for node in &graph.nodes {
        if let onnx_ir::ir::Node::Constant(const_node) = node
            && let Some(output) = const_node.outputs.first()
            && let Some(data) = output.value()
        {
            assert_eq!(
                data.shape,
                burn_tensor::Shape::new([2, 3]),
                "Expected shape [2, 3]"
            );
            let values: Vec<f32> = data.as_slice().unwrap().to_vec();
            assert_eq!(values, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
        }
    }
}

/// Test loading a model with mixed embedded and external data
#[test]
fn test_mixed_embedded_and_external_data() {
    // This model has:
    // - Large weight tensor (64x64) stored externally
    // - Small bias tensor (64) embedded in the ONNX file
    // The model should parse without errors
    let graph = load_onnx("mixed_data.onnx");

    // Verify graph structure
    assert_eq!(graph.inputs.len(), 1, "Expected 1 input");
    assert_eq!(graph.outputs.len(), 1, "Expected 1 output");

    // The graph should have nodes (may be fused/optimized)
    assert!(!graph.nodes.is_empty(), "Graph should have nodes");
}

/// Test that parsing from bytes (without file path) fails for external data
///
/// Models with external data cannot be parsed from raw bytes because
/// we need the file path to resolve external data file locations.
#[test]
#[should_panic(expected = "external data")]
fn test_external_data_from_bytes_panics() {
    use std::fs;

    let model_path = get_model_path("external_data.onnx");
    let bytes = fs::read(&model_path).expect("Failed to read model file");

    // This should panic because external data requires a file path context
    let _ = onnx_ir::OnnxGraphBuilder::new()
        .simplify(false)
        .parse_bytes(&bytes);
}

/// Test loading a model with tensors stored in multiple external files
///
/// This tests the ONNX spec requirement that different tensors can reference
/// different external files via separate "location" values.
#[test]
fn test_multiple_external_files() {
    // This model has:
    // - weight [4x4] stored in "multi_external_weights.bin"
    // - bias [4] stored in "multi_external_bias.bin"
    let graph = load_onnx("multi_external_files.onnx");

    // Verify graph structure
    assert_eq!(graph.inputs.len(), 1, "Expected 1 input");
    assert_eq!(graph.outputs.len(), 1, "Expected 1 output");

    // The graph should have nodes (may be fused/optimized to Linear)
    assert!(!graph.nodes.is_empty(), "Graph should have nodes");

    // Verify we can access the tensor values from different external files
    let mut found_weight = false;
    let mut found_bias = false;

    for node in &graph.nodes {
        if let onnx_ir::ir::Node::Constant(const_node) = node
            && let Some(output) = const_node.outputs.first()
            && let Some(data) = output.value()
        {
            let shape = data.shape.clone();

            // Weight tensor should be [4, 4] with diagonal values [1, 2, 3, 4]
            if shape == burn_tensor::Shape::new([4, 4]) {
                let values: Vec<f32> = data.as_slice().unwrap().to_vec();
                assert_eq!(values[0], 1.0, "weight[0,0] should be 1.0");
                assert_eq!(values[5], 2.0, "weight[1,1] should be 2.0");
                assert_eq!(values[10], 3.0, "weight[2,2] should be 3.0");
                assert_eq!(values[15], 4.0, "weight[3,3] should be 4.0");
                found_weight = true;
            }

            // Bias tensor should be [4] with values [0.5, 0.5, 0.5, 0.5]
            if shape == burn_tensor::Shape::new([4]) {
                let values: Vec<f32> = data.as_slice().unwrap().to_vec();
                assert!((values[0] - 0.5).abs() < 1e-6, "bias[0] should be 0.5");
                assert!((values[1] - 0.5).abs() < 1e-6, "bias[1] should be 0.5");
                assert!((values[2] - 0.5).abs() < 1e-6, "bias[2] should be 0.5");
                assert!((values[3] - 0.5).abs() < 1e-6, "bias[3] should be 0.5");
                found_bias = true;
            }
        }
    }

    // At least verify the graph parsed successfully
    // (constants may be fused into other nodes like Linear)
    assert!(
        found_weight || found_bias || !graph.nodes.is_empty(),
        "Should successfully parse model with multiple external files"
    );
}