Add support onnx size (tracel-ai#3301)

* feat: support onnx size

* test: remove unnecessary args for test

Author: tye-singwa

crates/burn-import/SUPPORTED-ONNX-OPS.md

@@ -177,7 +177,7 @@ functionality.
 | [Sign][163]                      | ✅             | ✅           |
 | [Sin][164]                       | ✅             | ✅           |
 | [Sinh][165]                      | ✅             | ✅           |
-| [Size][166]                      | ❌             | ✅           |
+| [Size][166]                      | ✅             | ✅           |
 | [Slice][167]                     | ✅             | ✅           |
 | [Softmax][168]                   | ✅             | ✅           |
 | [SoftmaxCrossEntropyLoss][169]   | ❌             | ❌           |

crates/burn-import/onnx-tests/build.rs

@@ -120,6 +120,7 @@ fn main() {
         .input("tests/sign/sign.onnx")
         .input("tests/sin/sin.onnx")
         .input("tests/sinh/sinh.onnx")
+        .input("tests/size/size.onnx")
         .input("tests/slice/slice.onnx")
         .input("tests/slice/slice_shape.onnx")
         .input("tests/softmax/softmax.onnx")

crates/burn-import/onnx-tests/tests/size/mod.rs

@@ -0,0 +1,24 @@
+// Import the shared macro
+use crate::include_models;
+include_models!(size);
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use burn::tensor::{Tensor, TensorData};
+
+    type Backend = burn_ndarray::NdArray<f32>;
+
+    #[test]
+    fn size() {
+        let model: size::Model<Backend> = size::Model::default();
+        let device = Default::default();
+
+        let input =
+            Tensor::<Backend, 1>::arange(0..(1 * 2 * 3 * 4 * 5), &device).reshape([1, 2, 3, 4, 5]);
+        let output = model.forward(input);
+        let expected = TensorData::from([120]);
+
+        output.to_data().assert_eq(&expected, true);
+    }
+}

crates/burn-import/onnx-tests/tests/size/size.onnx

@@ -0,0 +1,64 @@
+#!/usr/bin/env python3
+
+# used to generate model: size.onnx
+
+import numpy as np
+import onnx
+import onnx.helper
+from onnx import TensorProto
+from onnx.reference import ReferenceEvaluator
+
+
+def build_model():
+    # Define the graph inputs and outputs
+    input = onnx.helper.make_tensor_value_info(
+        'input', TensorProto.FLOAT, [2, 6, 2, 3])
+    output = onnx.helper.make_tensor_value_info(
+        'output', TensorProto.FLOAT, [1])
+
+    # Create the Size node
+    size = onnx.helper.make_node(
+        "Size",
+        inputs=["input"],
+        outputs=["output"],
+        name="SizeNode",
+    )
+
+    # Create the graph
+    graph = onnx.helper.make_graph(
+        [size],
+        'SizeModel',
+        [input],
+        [output],
+    )
+
+    # Create the model
+    model = onnx.helper.make_model(
+        opset_imports=[onnx.helper.make_operatorsetid("", 16)],
+        graph=graph,
+        producer_name='ONNX_Generator',
+    )
+
+    return model
+
+
+if __name__ == "__main__":
+    # Set seed and precision
+    np.random.seed(42)
+    np.set_printoptions(precision=8)
+
+    # Build model
+    test_input = np.arange(1*2*3*4*5).reshape(1, 2, 3, 4, 5)
+    onnx_model = build_model()
+    file_name = "size.onnx"
+
+    # Ensure valid ONNX and save
+    onnx.checker.check_model(onnx_model)
+    onnx.save(onnx_model, file_name)
+    print(f"Finished exporting model to {file_name}")
+
+    # Output some test data for use in the test
+    print(f"Test input data shape: {test_input.shape}")
+    session = ReferenceEvaluator(file_name, verbose=1)
+    test_output, = session.run(None, {"input": test_input})
+    print(f"Test output: {repr(test_output)}")

crates/burn-import/onnx-tests/tests/size/size.py

@@ -51,6 +51,7 @@ pub enum UnaryNodeKind {
     Tanh,
     Transpose,
     Sign,
+    Size,
 }
 
 impl UnaryNodeKind {
@@ -86,6 +87,7 @@ impl UnaryNodeKind {
             Self::Tanh => "tanh",
             Self::Transpose => "transpose",
             Self::Sign => "sign",
+            Self::Size => "size",
         }
     }
 }
@@ -513,6 +515,11 @@ impl UnaryNode {
         let function = move |input| quote! { #input.sign()};
         Self::new(input, output, UnaryNodeKind::Sign, Rc::new(function))
     }
+
+    pub(crate) fn size(input: Type, output: Type) -> Self {
+        let function = move |input| quote! { #input.shape.num_elements()};
+        Self::new(input, output, UnaryNodeKind::Size, Rc::new(function))
+    }
 }
 
 #[cfg(test)]
@@ -1214,4 +1221,23 @@ mod tests {
             vec!["tensor2".to_string()],
         );
     }
+
+    #[test]
+    fn test_unary_codegen_size() {
+        one_node_graph(
+            UnaryNode::size(
+                Type::Tensor(TensorType::new_float("tensor1", 4)),
+                Type::Scalar(ScalarType::new("scalar1", ScalarKind::Int64)),
+            ),
+            quote! {
+                pub fn forward(&self, tensor1: Tensor<B, 4>) -> i64 {
+                    let scalar1 = tensor1.shape.num_elements();
+
+                    scalar1
+                }
+            },
+            vec!["tensor1".to_string()],
+            vec!["scalar1".to_string()],
+        );
+    }
 }

crates/burn-import/src/burn/node/unary.rs

@@ -368,6 +368,7 @@ impl ParsedOnnxGraph {
                 NodeType::Sigmoid => graph.register(Self::sigmoid_conversion(node)),
                 NodeType::Sin => graph.register(Self::sin_conversion(node)),
                 NodeType::Sinh => graph.register(Self::sinh_conversion(node)),
+                NodeType::Size => graph.register(Self::size_conversion(node)),
                 NodeType::Slice => graph.register(Self::slice_conversion(node)),
                 NodeType::SpaceToDepth => graph.register(Self::space_to_depth_conversion(node)),
                 NodeType::Sum => graph.register(Self::sum_conversion(node)),
@@ -908,6 +909,13 @@ impl ParsedOnnxGraph {
         UnaryNode::sinh(input, output)
     }
 
+    fn size_conversion(node: Node) -> UnaryNode {
+        let input = Type::from(node.inputs.first().unwrap());
+        let output = Type::from(node.outputs.first().unwrap());
+
+        UnaryNode::size(input, output)
+    }
+
     fn slice_conversion(node: Node) -> SliceNode {
         let input = Type::from(node.inputs.first().unwrap());
         let output = Type::from(node.outputs.first().unwrap());

crates/burn-import/src/onnx/to_burn.rs

@@ -57,6 +57,7 @@ pub mod reduce_sum;
 pub mod reshape;
 pub mod resize;
 pub mod shape;
+pub mod size;
 pub mod slice;
 pub mod softmax;
 pub mod space_to_depth;

crates/onnx-ir/src/node/mod.rs

@@ -0,0 +1,40 @@
+use crate::ir::{ArgType, ElementType, Node};
+
+/// Update output type for Size (always scalar).
+pub fn size_update_outputs(node: &mut Node) {
+    log::debug!("Size rank inference for node {}", node.name);
+
+    assert_eq!(
+        node.inputs.len(),
+        1,
+        "Size: expected 1 input, found {}",
+        node.inputs.len()
+    );
+
+    node.outputs[0].ty = ArgType::Scalar(ElementType::Int64);
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::ir::NodeType;
+    use crate::node::test_utils::NodeBuilder;
+
+    fn create_test_node(rank: usize) -> Node {
+        let builder = NodeBuilder::new(NodeType::Size, "test_size")
+            .input_tensor_f32("data", rank, None)
+            .output_scalar_i64("size");
+
+        builder.build()
+    }
+
+    #[test]
+    fn test_size_update_outputs() {
+        let mut node = create_test_node(4);
+        size_update_outputs(&mut node);
+        assert!(matches!(
+            &node.outputs[0].ty,
+            ArgType::Scalar(ElementType::Int64)
+        ));
+    }
+}

crates/onnx-ir/src/node/size.rs

@@ -12,7 +12,7 @@ use crate::{
         reduce_max::reduce_max_update_outputs, reduce_mean::reduce_mean_update_outputs,
         reduce_min::reduce_min_update_outputs, reduce_prod::reduce_prod_update_outputs,
         reduce_sum::reduce_sum_update_outputs, reshape::reshape_update_outputs,
-        shape::shape_update_outputs, slice::slice_update_output_rank,
+        shape::shape_update_outputs, size::size_update_outputs, slice::slice_update_output_rank,
         space_to_depth::space_to_depth_update_outputs, split::split_update_outputs,
         squeeze::squeeze_update_output, topk::top_k_update_output,
         unsqueeze::unsqueeze_update_output, where_op::where_update_outputs,
@@ -104,6 +104,7 @@ pub fn rank_inference(node: &mut Node) {
         NodeType::Sign => same_as_input(node),
         NodeType::Sin => same_as_input(node),
         NodeType::Sinh => same_as_input(node),
+        NodeType::Size => size_update_outputs(node),
         NodeType::Slice => slice_update_output_rank(node),
         NodeType::Softmax => same_as_input(node),
         NodeType::SpaceToDepth => space_to_depth_update_outputs(node),