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

//! # Size
//!
//! Returns the total number of elements in the input tensor as a scalar int64 value.
//!
//! **ONNX Spec**: <https://onnx.ai/onnx/operators/onnx__Size.html>
//!
//! ## Opset Versions
//! - **Since version 23**: Current version
//! - **Since version 1**: Initial implementation

use onnx_ir_derive::NodeBuilder;

use crate::ir::{ArgType, Argument, DType, Node, RawNode};
use crate::processor::{
    InputSpec, NodeProcessor, NodeSpec, OutputPreferences, OutputSpec, ProcessError,
};

/// Node representation for Size operation
#[derive(Debug, Clone, NodeBuilder)]
pub struct SizeNode {
    pub name: String,
    pub inputs: Vec<Argument>,
    pub outputs: Vec<Argument>,
}

pub(crate) struct SizeProcessor;

impl NodeProcessor for SizeProcessor {
    type Config = ();

    fn spec(&self) -> NodeSpec {
        NodeSpec {
            min_opset: 1,
            max_opset: None,
            inputs: InputSpec::Exact(1),
            outputs: OutputSpec::Exact(1),
        }
    }

    fn infer_types(
        &self,
        node: &mut RawNode,
        _opset: usize,
        _output_preferences: &OutputPreferences,
    ) -> Result<(), ProcessError> {
        node.outputs[0].ty = ArgType::ScalarNative(DType::I64);

        Ok(())
    }

    fn build_node(&self, builder: RawNode, _opset: usize) -> Node {
        Node::Size(SizeNode {
            name: builder.name,
            inputs: builder.inputs,
            outputs: builder.outputs,
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::ir::NodeType;
    use crate::node::test_utils::TestNodeBuilder;

    fn create_test_node(rank: usize) -> RawNode {
        let builder = TestNodeBuilder::new(NodeType::Size, "test_size")
            .input_tensor_f32("data", rank, None)
            .output_scalar_i64("size");

        builder.build()
    }

    // TODO: Missing test for zero-size tensors - what is the size of a tensor with shape [0, 5, 3]?
    // Should be 0 (0 * 5 * 3 = 0) but this edge case is not tested.

    // TODO: Missing test for scalar (rank-0) tensors - what is the size of a scalar?
    // Should be 1 according to ONNX spec but not validated.

    // TODO: Missing test for very large tensors - size can overflow i64 for extremely large tensors.
    // Need to verify behavior when product of dimensions exceeds i64::MAX.

    #[test]
    fn test_size_update_outputs() {
        let mut node = create_test_node(4);

        let processor = SizeProcessor;
        let prefs = OutputPreferences::new();
        processor.infer_types(&mut node, 16, &prefs).unwrap();

        assert!(matches!(
            &node.outputs[0].ty,
            ArgType::ScalarNative(DType::I64)
        ));
    }
}