feat: implement shape inference for Flatten operator

Author: pixelspark

README.md

@@ -253,7 +253,7 @@ fn test_matmul_square_matrix() {
 |<a href="https://github.com/onnx/onnx/blob/main/docs/Operators.md#Exp">Exp</a>|<a href="https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Exp-13">13</a>, <a href="https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Exp-6">6</a>, <a href="https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Exp-1">1</a>|✅|✅|
 |<a href="https://github.com/onnx/onnx/blob/main/docs/Operators.md#Expand">Expand</a>|<a href="https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Expand-13">13</a>, <a href="https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Expand-8">8</a>|
 |<a href="https://github.com/onnx/onnx/blob/main/docs/Operators.md#EyeLike">EyeLike</a>|<a href="https://github.com/onnx/onnx/blob/main/docs/Changelog.md#EyeLike-9">9</a>|
-|<a href="https://github.com/onnx/onnx/blob/main/docs/Operators.md#Flatten">Flatten</a>|<a href="https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Flatten-13">13</a>, <a href="https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Flatten-11">11</a>, <a href="https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Flatten-9">9</a>, <a href="https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Flatten-1">1</a>|✅|
+|<a href="https://github.com/onnx/onnx/blob/main/docs/Operators.md#Flatten">Flatten</a>|<a href="https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Flatten-13">13</a>, <a href="https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Flatten-11">11</a>, <a href="https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Flatten-9">9</a>, <a href="https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Flatten-1">1</a>|✅|✅|
 |<a href="https://github.com/onnx/onnx/blob/main/docs/Operators.md#Floor">Floor</a>|<a href="https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Floor-13">13</a>, <a href="https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Floor-6">6</a>, <a href="https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Floor-1">1</a>|✅|✅|
 |<a href="https://github.com/onnx/onnx/blob/main/docs/Operators.md#GRU">GRU</a>|<a href="https://github.com/onnx/onnx/blob/main/docs/Changelog.md#GRU-14">14</a>, <a href="https://github.com/onnx/onnx/blob/main/docs/Changelog.md#GRU-7">7</a>, <a href="https://github.com/onnx/onnx/blob/main/docs/Changelog.md#GRU-3">3</a>, <a href="https://github.com/onnx/onnx/blob/main/docs/Changelog.md#GRU-1">1</a>|
 |<a href="https://github.com/onnx/onnx/blob/main/docs/Operators.md#Gather">Gather</a>|<a href="https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Gather-13">13</a>, <a href="https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Gather-11">11</a>, <a href="https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Gather-1">1</a>|✅ (axis=0)|✅|

wonnx-preprocessing/src/shape_inference.rs

@@ -506,6 +506,33 @@ pub(crate) fn infer_output_shapes(
             Ok(vec![output_shape])
         }
 
+        ("Flatten", 1, 1) => {
+            let axis: usize = {
+                let a = node.get_attribute_value("axis", Some(1)).unwrap();
+                if a < 0 {
+                    (a + input_shapes[0].rank() as i64) as usize
+                } else {
+                    a as usize
+                }
+            };
+            if axis > input_shapes[0].rank() {
+                return Err(ShapeInferenceError::InvalidNode(
+                    node.get_name().to_string(),
+                    format!("Flatten axis attribute ({axis}) should be less than or equal to rank of input ({})",input_shapes[0].rank()),
+                ));
+            }
+            let input_dims = &input_shapes[0].dims;
+            let outer_dim = if axis == 0 {
+                1
+            } else {
+                input_dims[0..=(axis - 1)].iter().product::<u64>() as i64
+            };
+            let inner_dim = input_dims[axis..].iter().product::<u64>() as i64;
+
+            let new_dims = vec![outer_dim, inner_dim];
+            Ok(vec![Shape::from(input_shapes[0].data_type, &new_dims)])
+        }
+
         ("GlobalAveragePool", 1, 1) => {
             let mut output_shape = input_shapes[0].clone();
             if output_shape.rank() < 2 {