onnx-ir 0.21.0

ONNX-IR is a pure Rust library for parsing ONNX models into an intermediate representation that can be used to generate code for various ML/DL frameworks
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
122
123
124
125
126
127
128
129
130

#![allow(dead_code)]

/// Test utilities for ONNX-IR integration tests
///
/// Provides helper functions for loading ONNX models and common test assertions.
use std::path::PathBuf;

/// Load an ONNX model from the tests directory (without simplification)
///
/// # Arguments
/// * `model_name` - Name of the ONNX file (e.g., "basic_model.onnx")
///
/// # Returns
/// Parsed OnnxGraph
///
/// # Panics
/// Panics if the model file doesn't exist or parsing fails
pub fn load_onnx(model_name: &str) -> onnx_ir::ir::OnnxGraph {
    load_onnx_with_simplify(model_name, false)
}

/// Load an ONNX model with simplification enabled
pub fn load_onnx_simplified(model_name: &str) -> onnx_ir::ir::OnnxGraph {
    load_onnx_with_simplify(model_name, true)
}

fn load_onnx_with_simplify(model_name: &str, simplify: bool) -> onnx_ir::ir::OnnxGraph {
    let model_path = get_model_path(model_name);
    onnx_ir::OnnxGraphBuilder::new()
        .simplify(simplify)
        .parse_file(&model_path)
        .unwrap_or_else(|e| {
            panic!(
                "Failed to parse ONNX model '{}': {}",
                model_path.display(),
                e
            )
        })
}

/// Get the path to an ONNX model in the tests/fixtures directory
///
/// # Arguments
/// * `model_name` - Name of the ONNX file
///
/// # Returns
/// PathBuf to the model file
pub fn get_model_path(model_name: &str) -> PathBuf {
    PathBuf::from(env!("CARGO_MANIFEST_DIR"))
        .join("tests")
        .join("fixtures")
        .join(model_name)
}

/// Count nodes matching a predicate
///
/// # Arguments
/// * `graph` - The OnnxGraph to search
/// * `predicate` - Function to test each node
///
/// # Returns
/// Number of nodes matching the predicate
///
/// # Example
/// ```ignore
/// count_nodes(&graph, |n| matches!(n, Node::PRelu { .. }))
/// ```
pub fn count_nodes<F>(graph: &onnx_ir::ir::OnnxGraph, predicate: F) -> usize
where
    F: Fn(&onnx_ir::ir::Node) -> bool,
{
    graph.nodes.iter().filter(|n| predicate(n)).count()
}

/// Count Constant nodes in the graph
pub fn count_constant_nodes(graph: &onnx_ir::ir::OnnxGraph) -> usize {
    count_nodes(graph, |n| matches!(n, onnx_ir::ir::Node::Constant { .. }))
}

/// Count operation nodes (non-Constant nodes) in the graph
pub fn count_operation_nodes(graph: &onnx_ir::ir::OnnxGraph) -> usize {
    graph
        .nodes
        .iter()
        .filter(|n| !matches!(n, onnx_ir::ir::Node::Constant { .. }))
        .count()
}

/// Check if a graph contains nodes matching a predicate
///
/// # Example
/// ```ignore
/// has_node_type(&graph, |n| matches!(n, Node::MatMul { .. }))
/// ```
pub fn has_node_type<F>(graph: &onnx_ir::ir::OnnxGraph, predicate: F) -> bool
where
    F: Fn(&onnx_ir::ir::Node) -> bool,
{
    graph.nodes.iter().any(predicate)
}

/// Get all unique data types from graph inputs
pub fn get_input_dtypes(graph: &onnx_ir::ir::OnnxGraph) -> Vec<burn_tensor::DType> {
    graph.inputs.iter().map(|inp| inp.ty.elem_type()).collect()
}

/// Count outputs by ArgType variant
pub fn count_outputs_by_type(graph: &onnx_ir::ir::OnnxGraph) -> (usize, usize, usize) {
    use onnx_ir::ir::ArgType;

    let scalars = graph
        .outputs
        .iter()
        .filter(|out| matches!(out.ty, ArgType::ScalarTensor(_) | ArgType::ScalarNative(_)))
        .count();

    let shapes = graph
        .outputs
        .iter()
        .filter(|out| matches!(out.ty, ArgType::Shape(_)))
        .count();

    let tensors = graph
        .outputs
        .iter()
        .filter(|out| matches!(out.ty, ArgType::Tensor(_)))
        .count();

    (scalars, shapes, tensors)
}