onnx-ir 0.21.0

//! # Expand
//!
//! Broadcasts input tensor to a target shape using numpy-style broadcasting.
//!
//! **ONNX Spec**: <https://onnx.ai/onnx/operators/onnx__Expand.html>
//!
//! ## Opset Versions
//! - **Opset 8**: Initial version (replaces deprecated Tile for broadcasting)
//! - **Opset 13**: Extended type support (bfloat16)
use onnx_ir_derive::NodeBuilder;

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

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

/// Shape information for the Expand operation.
#[derive(Debug, Clone)]
// TODO rename ExpandConfig to ExpandConfig
pub enum ExpandConfig {
    /// Static shape information known at compile time.
    Static(Vec<i64>),
    /// Runtime shape determined during execution - references node.inputs\[input_index\].
    Runtime(RuntimeInputRef),
}

pub(crate) struct ExpandProcessor;

impl NodeProcessor for ExpandProcessor {
    type Config = ExpandConfig;

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

    fn lift_constants(&self, node: &mut RawNode, _opset: usize) -> Result<(), ProcessError> {
        // Only lift shape input (input[1]) if it has a static value
        // Runtime shapes should remain in the graph
        if node.inputs.len() > 1 && node.inputs[1].is_constant() {
            node.inputs[1].to_static()?;
        }

        Ok(())
    }

    fn infer_types(
        &self,
        node: &mut RawNode,
        opset: usize,
        _output_preferences: &OutputPreferences,
    ) -> Result<(), ProcessError> {
        // TODO: Validate no unexpected attributes - Expand has no attributes per spec - Missing attribute validation

        // Validate shape input type
        match &node.inputs[1].ty {
            ArgType::Tensor(tensor) => {
                if tensor.rank != 1 {
                    return Err(ProcessError::Custom(
                        "Expand: shape tensor must be 1D".to_string(),
                    ));
                }
                if !matches!(tensor.dtype, DType::I64) {
                    return Err(ProcessError::Custom(
                        "Expand: shape tensor must have element type int64".to_string(),
                    ));
                }
            }
            ArgType::Shape(_) => {
                // Shapes are always 1-D int64 data, so nothing to validate here
            }
            ArgType::ScalarTensor(_) | ArgType::ScalarNative(_) => {
                // Scalar shape means expanding to rank 0 (scalar output)
            }
        }

        // Get reference to config for type inference
        let config = self
            .extract_config(node, opset)
            .expect("Config extraction failed");

        // Get input element type - Expand should preserve the input's element type
        let input_elem_type = match &node.inputs[0].ty {
            ArgType::Tensor(tensor) => tensor.dtype,
            ArgType::ScalarTensor(dtype) | ArgType::ScalarNative(dtype) => *dtype,
            // Shape is a 1D int64 tensor of dimension sizes
            ArgType::Shape(_) => DType::I64,
        };

        // Determine output type based on config
        match config {
            ExpandConfig::Static(shape) => {
                // TODO: Validate shape values are positive or -1 per ONNX spec - Negative values other than -1 are invalid - Missing constraint validation
                // TODO: Validate broadcasting rules - Per spec, input shape and target shape must be compatible for broadcasting - Missing broadcast validation
                if shape.is_empty() {
                    // Empty shape means scalar output
                    node.outputs[0].ty = ArgType::ScalarTensor(input_elem_type);
                } else {
                    node.outputs[0].ty = ArgType::Tensor(TensorType {
                        dtype: input_elem_type,
                        rank: shape.len(),
                        static_shape: Some(shape.iter().map(|&dim| Some(dim as usize)).collect()),
                    });
                }
            }
            ExpandConfig::Runtime(_) => {
                // When the shape cannot be determined statically, infer the rank from the shape input
                let output_rank = match &node.inputs[1].ty {
                    // Scalar shape input means expanding to a scalar (rank 0)
                    ArgType::ScalarTensor(_) | ArgType::ScalarNative(_) => 0,
                    ArgType::Shape(rank) => *rank,
                    ArgType::Tensor(tensor) => {
                        if let Some(static_shape) = &tensor.static_shape
                            && let Some(Some(rank)) = static_shape.first()
                        {
                            *rank
                        } else {
                            // Check if output already has a rank set from ONNX
                            match &node.outputs[0].ty {
                                ArgType::Tensor(TensorType { rank, .. }) if *rank > 0 => *rank,
                                _ => {
                                    // Fallback: use the input tensor's rank as the output rank.
                                    // Per ONNX spec, output rank = max(input_rank, len(shape)).
                                    // When len(shape) is unknown, this assumes same-rank
                                    // broadcasting (len(shape) <= input_rank), which is correct
                                    // for the known real-world case (SDXL UNet: 1D timestep
                                    // expanded to 1D [batch_size]).
                                    match &node.inputs[0].ty {
                                        ArgType::Tensor(t) => t.rank,
                                        // Shape is always 1D
                                        ArgType::Shape(_) => 1,
                                        other => {
                                            return Err(ProcessError::Custom(format!(
                                                "Cannot determine output rank for Expand node {} \
                                                 with fully dynamic shape tensor and input type {:?}",
                                                node.name, other
                                            )));
                                        }
                                    }
                                }
                            }
                        }
                    }
                };

                if output_rank == 0 {
                    node.outputs[0].ty = ArgType::ScalarTensor(input_elem_type);
                } else {
                    node.outputs[0].ty = ArgType::Tensor(TensorType {
                        dtype: input_elem_type,
                        rank: output_rank,
                        static_shape: None,
                    });
                }
            }
        }

        Ok(())
    }

    fn is_noop(&self, node: &RawNode) -> bool {
        // Expand is a no-op when output shape == input shape (no actual broadcasting)
        if let (ArgType::Tensor(in_t), ArgType::Tensor(out_t)) =
            (&node.inputs[0].ty, &node.outputs[0].ty)
            && let (Some(in_shape), Some(out_shape)) = (&in_t.static_shape, &out_t.static_shape)
        {
            return in_shape == out_shape;
        }
        false
    }

    fn extract_config(&self, node: &RawNode, _opset: usize) -> Result<Self::Config, ProcessError> {
        // Extract config
        let config = match node.inputs[1].value() {
            Some(tensor_data) => match tensor_data.to_i64_vec() {
                Ok(shape) => ExpandConfig::Static(shape),
                Err(_) => {
                    return Err(ProcessError::Custom(
                        "Expand: shape data type must be int32 or int64".to_string(),
                    ));
                }
            },
            None => {
                // Runtime shape - store reference instead of cloning the argument
                ExpandConfig::Runtime(RuntimeInputRef::new(node.inputs[1].name.clone(), 1))
            }
        };
        Ok(config)
    }

    fn build_node(&self, builder: RawNode, opset: usize) -> Node {
        let config = self
            .extract_config(&builder, opset)
            .expect("Config extraction failed");

        Node::Expand(ExpandNode {
            name: builder.name,
            inputs: builder.inputs,
            outputs: builder.outputs,
            config,
        })
    }
}

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

    fn create_test_node(
        input_rank: usize,
        shape_value: Option<Vec<i64>>,
        shape_type: Option<ArgType>,
    ) -> TestNodeBuilder {
        let mut builder = TestNodeBuilder::new(NodeType::Expand, "test_expand")
            .input_tensor_f32("input", input_rank, None)
            .output_tensor_f32("output", 0, None); // Rank 0 will be updated

        if let Some(shape) = shape_value {
            builder = builder.input_tensor_i64_data("shape", shape.clone(), vec![shape.len()]);
        } else if let Some(st) = shape_type {
            // Use the provided custom shape type
            builder = builder.add_input("shape", st);
        } else {
            // Default case with dynamic shape
            builder = builder.input_tensor_i64("shape", 1, Some(vec![3]));
        }

        builder
    }

    #[test]
    fn test_expand_with_constant_shape() {
        let mut node = create_test_node(2, Some(vec![2, 3, 4]), None).build_with_graph_data(16);

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

        match &node.outputs[0].ty {
            ArgType::Tensor(tensor) => {
                assert_eq!(tensor.dtype, DType::F32);
                assert_eq!(tensor.rank, 3);
                assert_eq!(tensor.static_shape, Some(vec![Some(2), Some(3), Some(4)]));
            }
            _ => panic!("Expected tensor output"),
        }
    }

    #[test]
    fn test_expand_with_dynamic_shape() {
        let mut node = create_test_node(2, None, None).build();

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

        match &node.outputs[0].ty {
            ArgType::Tensor(tensor) => {
                assert_eq!(tensor.dtype, DType::F32);
                assert_eq!(tensor.rank, 3);
                assert_eq!(tensor.static_shape, None);
            }
            _ => panic!("Expected tensor output"),
        }
    }

    #[test]
    fn test_expand_with_incorrect_inputs() {
        let mut node = create_test_node(2, Some(vec![2, 3, 4]), None).build_with_graph_data(16);
        // Remove one input to make it invalid
        node.inputs.pop();

        let processor = ExpandProcessor;
        let spec = processor.spec();
        let result = crate::processor::validate_node_spec(&node, 16, &spec);
        assert!(matches!(
            result,
            Err(ProcessError::InvalidInputCount {
                expected: 2,
                actual: 1
            })
        ));
    }

    // Tests for expand_config function

    #[test]
    fn test_expand_config_with_static_shape() {
        let node = create_test_node(2, Some(vec![2, 3, 4]), None).build_with_graph_data(16);
        let mut node = node;
        let processor = ExpandProcessor;
        let prefs = OutputPreferences::new();
        let config = processor.extract_config(&node, 16).unwrap();
        processor.infer_types(&mut node, 16, &prefs).unwrap();

        match config {
            ExpandConfig::Static(shape) => {
                assert_eq!(*shape, vec![2, 3, 4]);
            }
            ExpandConfig::Runtime(_) => panic!("Expected Static config, got Runtime"),
        }
    }

    #[test]
    fn test_expand_config_with_runtime_shape() {
        let node = create_test_node(2, None, None).build();
        let mut node = node;
        let processor = ExpandProcessor;
        let prefs = OutputPreferences::new();
        let config = processor.extract_config(&node, 16).unwrap();
        processor.infer_types(&mut node, 16, &prefs).unwrap();

        match config {
            ExpandConfig::Static(_) => panic!("Expected Runtime config, got Static"),
            ExpandConfig::Runtime(name) => {
                assert_eq!(name.name, "shape");
            }
        }
    }

    #[test]
    fn test_expand_config_with_shape_type() {
        let shape_type = ArgType::Shape(3);
        let node = create_test_node(2, None, Some(shape_type)).build();
        let mut node = node;
        let processor = ExpandProcessor;
        let prefs = OutputPreferences::new();
        let config = processor.extract_config(&node, 16).unwrap();
        processor.infer_types(&mut node, 16, &prefs).unwrap();

        match config {
            ExpandConfig::Static(_) => panic!("Expected Runtime config, got Static"),
            ExpandConfig::Runtime(name) => {
                assert_eq!(name.name, "shape");
            }
        }
    }

    #[test]
    fn test_expand_config_with_invalid_shape_rank() {
        let invalid_shape_type = ArgType::Tensor(TensorType {
            dtype: DType::I64,
            rank: 2, // Invalid rank, should be 1
            static_shape: None,
        });
        let node = create_test_node(2, None, Some(invalid_shape_type)).build();
        let mut node = node;
        let processor = ExpandProcessor;
        let prefs = OutputPreferences::new();
        let _config = processor.extract_config(&node, 16).unwrap();
        let result = processor.infer_types(&mut node, 16, &prefs);
        assert!(matches!(result, Err(ProcessError::Custom(_))));
    }

    #[test]
    fn test_expand_config_with_invalid_shape_type() {
        let invalid_shape_type = ArgType::Tensor(TensorType {
            dtype: DType::F32, // Invalid element type, should be Int64
            rank: 1,
            static_shape: None,
        });
        let node = create_test_node(2, None, Some(invalid_shape_type)).build();
        let mut node = node;
        let processor = ExpandProcessor;
        let prefs = OutputPreferences::new();
        let _config = processor.extract_config(&node, 16).unwrap();
        let result = processor.infer_types(&mut node, 16, &prefs);
        assert!(matches!(result, Err(ProcessError::Custom(_))));
    }

    #[test]
    fn test_expand_scalar_native_shape_outputs_scalar() {
        // ScalarNative shape input means expanding to rank 0
        let shape_type = ArgType::ScalarNative(DType::I64);
        let node = create_test_node(2, None, Some(shape_type)).build();
        let mut node = node;
        let processor = ExpandProcessor;
        let prefs = OutputPreferences::new();
        processor.infer_types(&mut node, 16, &prefs).unwrap();

        assert_eq!(node.outputs[0].ty, ArgType::ScalarTensor(DType::F32));
    }

    #[test]
    fn test_expand_config_with_invalid_value_type() {
        // Create a node with shape input that has Float32 type instead of Int64
        let node = TestNodeBuilder::new(NodeType::Expand, "test_expand")
            .input_tensor_f32("input", 2, None)
            .input_tensor_f32_data("shape", vec![2.0, 3.0, 4.0], vec![3]) // Wrong type - Float32 instead of Int64
            .output_tensor_f32("output", 0, None)
            .build_with_graph_data(16);

        let node = node;
        let processor = ExpandProcessor;
        let result = processor.extract_config(&node, 16);
        match result {
            Err(ProcessError::Custom(_)) => {}
            _ => panic!("Expected ProcessError::Custom for invalid shape data type"),
        }
    }

    #[test]
    fn test_expand_update_outputs_with_shape_input() {
        // Test Expand with Shape type as shape input
        let mut node = create_test_node(2, None, Some(ArgType::Shape(4))).build();

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

        match &node.outputs[0].ty {
            ArgType::Tensor(tensor) => {
                assert_eq!(tensor.dtype, DType::F32);
                assert_eq!(tensor.rank, 4); // Shape(4) means output will be rank 4
                assert_eq!(tensor.static_shape, None); // Dynamic shape
            }
            _ => panic!("Expected tensor output"),
        }
    }

    #[test]
    fn test_expand_update_outputs_with_shape_input_static_value() {
        // Test Expand with shape input that has static values
        let mut node = TestNodeBuilder::new(NodeType::Expand, "test_expand")
            .input_tensor_f32("input", 2, None)
            .input_tensor_i64_data("shape", vec![5, 10, 15], vec![3]) // Static shape values
            .output_tensor_f32("output", 0, None)
            .build_with_graph_data(16);

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

        match &node.outputs[0].ty {
            ArgType::Tensor(tensor) => {
                assert_eq!(tensor.dtype, DType::F32);
                assert_eq!(tensor.rank, 3);
                assert_eq!(tensor.static_shape, Some(vec![Some(5), Some(10), Some(15)]));
            }
            _ => panic!("Expected tensor output"),
        }
    }

    #[test]
    fn test_expand_preserves_input_element_type() {
        // Test that Expand preserves the input element type for different types

        // Test Float32 -> Float32
        {
            let mut node = TestNodeBuilder::new(NodeType::Expand, "test_expand")
                .input_tensor_f32("input", 2, None)
                .input_tensor_i64_data("shape", vec![2, 3, 4], vec![3])
                .output_tensor_f32("output", 0, None)
                .build_with_graph_data(16);

            // Initially set output to wrong type
            node.outputs[0].ty = ArgType::Tensor(TensorType {
                dtype: DType::I64, // Wrong type
                rank: 0,
                static_shape: None,
            });

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

            match &node.outputs[0].ty {
                ArgType::Tensor(tensor) => {
                    assert_eq!(
                        tensor.dtype,
                        DType::F32,
                        "Expand should preserve Float32 input type"
                    );
                    assert_eq!(tensor.rank, 3);
                }
                _ => panic!("Expected tensor output"),
            }
        }

        // Test Int64 -> Int64
        {
            let mut node = TestNodeBuilder::new(NodeType::Expand, "test_expand")
                .input_tensor_i64("input", 2, None)
                .input_tensor_i64_data("shape", vec![2, 3, 4], vec![3])
                .output_tensor_i64("output", 0, None)
                .build_with_graph_data(16);

            // Initially set output to wrong type
            node.outputs[0].ty = ArgType::Tensor(TensorType {
                dtype: DType::F32, // Wrong type
                rank: 0,
                static_shape: None,
            });

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

            match &node.outputs[0].ty {
                ArgType::Tensor(tensor) => {
                    assert_eq!(
                        tensor.dtype,
                        DType::I64,
                        "Expand should preserve Int64 input type"
                    );
                    assert_eq!(tensor.rank, 3);
                }
                _ => panic!("Expected tensor output"),
            }
        }

        // Test Bool -> Bool
        {
            let mut node = TestNodeBuilder::new(NodeType::Expand, "test_expand")
                .input_tensor_bool("input", 2, None)
                .input_tensor_i64_data("shape", vec![2, 3, 4], vec![3])
                .output_tensor_bool("output", 0, None)
                .build_with_graph_data(16);

            // Initially set output to wrong type
            node.outputs[0].ty = ArgType::Tensor(TensorType {
                dtype: DType::F32, // Wrong type
                rank: 0,
                static_shape: None,
            });

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

            match &node.outputs[0].ty {
                ArgType::Tensor(tensor) => {
                    assert_eq!(
                        tensor.dtype,
                        DType::Bool(BoolStore::Native),
                        "Expand should preserve Bool input type"
                    );
                    assert_eq!(tensor.rank, 3);
                }
                _ => panic!("Expected tensor output"),
            }
        }
    }

    #[test]
    fn test_expand_with_mismatched_output_type() {
        // Test that Expand corrects output type even when initially set incorrectly
        // This simulates the case where ONNX might have wrong type info
        let mut node = TestNodeBuilder::new(NodeType::Expand, "test_expand")
            .input_tensor_i64("input", 2, None) // Input is Int64
            .input_tensor_i64_data("shape", vec![2, 3], vec![2])
            .output_tensor_f32("output", 0, None) // Output incorrectly set to Float32
            .build_with_graph_data(16);

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

        match &node.outputs[0].ty {
            ArgType::Tensor(tensor) => {
                assert_eq!(
                    tensor.dtype,
                    DType::I64,
                    "Expand should use input type (Int64) not initial output type (Float32)"
                );
                assert_eq!(tensor.rank, 2);
                assert_eq!(tensor.static_shape, Some(vec![Some(2), Some(3)]));
            }
            _ => panic!("Expected tensor output"),
        }
    }

    #[test]
    fn test_expand_shape_input_static_target() {
        // When input[0] is ArgType::Shape (e.g., output of a Shape node),
        // Expand should treat it as a 1D int64 tensor.
        // This pattern occurs in models like piper-tts/VITS (issue #266).
        let mut node = TestNodeBuilder::new(NodeType::Expand, "test_expand")
            .add_input("input", ArgType::Shape(2))
            .input_tensor_i64_data("shape", vec![2, 3], vec![2])
            .output_tensor_i64("output", 0, None)
            .build_with_graph_data(16);

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

        match &node.outputs[0].ty {
            ArgType::Tensor(tensor) => {
                assert_eq!(tensor.dtype, DType::I64);
                assert_eq!(tensor.rank, 2);
                assert_eq!(tensor.static_shape, Some(vec![Some(2), Some(3)]));
            }
            _ => panic!("Expected tensor output"),
        }
    }

    #[test]
    fn test_expand_scalar_input_static_shape() {
        let mut node = TestNodeBuilder::new(NodeType::Expand, "test_expand")
            .input_scalar_f32("input")
            .input_tensor_i64_data("shape", vec![2, 3], vec![2])
            .output_tensor_f32("output", 0, None)
            .build_with_graph_data(16);

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

        match &node.outputs[0].ty {
            ArgType::Tensor(tensor) => {
                assert_eq!(tensor.dtype, DType::F32);
                assert_eq!(tensor.rank, 2);
                assert_eq!(tensor.static_shape, Some(vec![Some(2), Some(3)]));
            }
            _ => panic!("Expected tensor output"),
        }
    }

    #[test]
    fn test_expand_scalar_input_runtime_shape() {
        let mut node = TestNodeBuilder::new(NodeType::Expand, "test_expand")
            .input_scalar_i64("input")
            .input_tensor_i64("shape", 1, Some(vec![4]))
            .output_tensor_i64("output", 0, None)
            .build();

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

        match &node.outputs[0].ty {
            ArgType::Tensor(tensor) => {
                assert_eq!(tensor.dtype, DType::I64);
                assert_eq!(tensor.rank, 4);
                assert_eq!(tensor.static_shape, None);
            }
            _ => panic!("Expected tensor output"),
        }
    }

    #[test]
    fn test_expand_same_static_shape_is_noop() {
        let node = TestNodeBuilder::new(NodeType::Expand, "test")
            .input_tensor_f32("input", 3, Some(vec![2, 3, 4]))
            .input_tensor_i64("shape", 1, Some(vec![3]))
            .output_tensor_f32("output", 3, Some(vec![2, 3, 4]))
            .build();
        assert!(ExpandProcessor.is_noop(&node));
    }

    #[test]
    fn test_expand_different_static_shape_is_not_noop() {
        let node = TestNodeBuilder::new(NodeType::Expand, "test")
            .input_tensor_f32("input", 3, Some(vec![1, 3, 4]))
            .input_tensor_i64("shape", 1, Some(vec![3]))
            .output_tensor_f32("output", 3, Some(vec![2, 3, 4]))
            .build();
        assert!(!ExpandProcessor.is_noop(&node));
    }

    #[test]
    fn test_expand_no_static_shape_is_not_noop() {
        let node = TestNodeBuilder::new(NodeType::Expand, "test")
            .input_tensor_f32("input", 3, None)
            .input_tensor_i64("shape", 1, Some(vec![3]))
            .output_tensor_f32("output", 3, None)
            .build();
        assert!(!ExpandProcessor.is_noop(&node));
    }

    #[test]
    fn test_expand_with_fully_dynamic_shape_tensor() {
        // Shape input is a 1D Int64 tensor with static_shape = None (fully dynamic).
        // This happens in SDXL UNet where the shape comes from a Where/ConstantOfShape chain
        // with no ONNX value_info. The fallback uses the input tensor's rank.
        let shape_type = ArgType::Tensor(TensorType {
            dtype: DType::I64,
            rank: 1,
            static_shape: None,
        });
        let mut node = create_test_node(2, None, Some(shape_type)).build();

        // Clear the output rank so the ONNX-value_info fallback also fails
        node.outputs[0].ty = ArgType::Tensor(TensorType {
            dtype: DType::F32,
            rank: 0,
            static_shape: None,
        });

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

        match &node.outputs[0].ty {
            ArgType::Tensor(tensor) => {
                assert_eq!(tensor.dtype, DType::F32);
                // Falls back to input rank (2) when shape length is unknown
                assert_eq!(tensor.rank, 2);
                assert_eq!(tensor.static_shape, None);
            }
            _ => panic!("Expected tensor output"),
        }
    }

    #[test]
    fn test_expand_scalar_input_fully_dynamic_shape_errors() {
        // Scalar input with fully dynamic shape (static_shape = None) should error
        // because we cannot determine len(shape) and input rank is meaningless for scalars.
        let shape_type = ArgType::Tensor(TensorType {
            dtype: DType::I64,
            rank: 1,
            static_shape: None,
        });
        let mut node = TestNodeBuilder::new(NodeType::Expand, "test_expand")
            .input_scalar_f32("input")
            .add_input("shape", shape_type)
            .output_tensor_f32("output", 0, None)
            .build();

        node.outputs[0].ty = ArgType::Tensor(TensorType {
            dtype: DType::F32,
            rank: 0,
            static_shape: None,
        });

        let processor = ExpandProcessor;
        let prefs = OutputPreferences::new();
        let result = processor.infer_types(&mut node, 16, &prefs);
        assert!(matches!(result, Err(ProcessError::Custom(_))));
    }

    #[test]
    fn test_expand_static_empty_shape_outputs_scalar() {
        // Expanding with an empty shape [] produces a scalar output
        let mut node = TestNodeBuilder::new(NodeType::Expand, "test_expand")
            .input_scalar_f32("input")
            .input_tensor_i64_data("shape", vec![], vec![0])
            .output_tensor_f32("output", 0, None)
            .build_with_graph_data(16);

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

        assert_eq!(node.outputs[0].ty, ArgType::ScalarTensor(DType::F32));
    }

    // TODO: Add test for invalid shape values - Test negative values other than -1 (e.g., -2, -3) should return error - Missing constraint validation test
    // TODO: Add test for shape with value -1 - Per spec, -1 means copy from input dimension - Missing edge case test
    // TODO: Add test for incompatible broadcasting - Test case where input shape cannot be broadcast to target shape - Missing broadcast validation test
    // TODO: Add test for zero in target shape - Test behavior when target shape contains 0 - Missing edge case test
    // TODO: Add test for expanding scalar to tensor - Test input with rank 0 expanded to higher rank - Missing edge case test
    // TODO: Add test for different data types - Spec supports many types (all numeric types, bool, strings) - Only testing f32, i64, bool
    // TODO: Add test for opset < 8 - Should fail per spec, Expand introduced in opset 8 - Missing opset validation test
    // TODO: Add test for unexpected attributes - Should validate and reject unknown attributes - Missing attribute validation test
}