onnx-ir 0.21.0

//! # ArgMax
//!
//! Computes the indices of the maximum elements along the specified axis.
//!
//! **ONNX Spec**: <https://onnx.ai/onnx/operators/onnx__ArgMax.html>
//!
//! ## Opset Versions
//!
//! - **Opset 13**: Added bfloat16 to type constraints
//! - **Opset 12**: Added `select_last_index` attribute
//! - **Opset 11**: Changed `axis` range to support negative indices [-r, r-1]

use onnx_ir_derive::NodeBuilder;

use crate::ir::Argument;

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

/// Configuration for ArgMax operations
#[derive(Debug, Clone, new)]
pub struct ArgMaxConfig {
    /// Axis along which to find the maximum
    pub axis: usize,
    /// Whether to keep dimensions after reduction
    pub keepdims: bool,
    /// When true, return the index of the LAST occurrence of the maximum
    /// along the axis (ONNX `select_last_index=1`). Default false returns
    /// the first occurrence.
    pub select_last_index: bool,
}

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

pub(crate) struct ArgMaxProcessor;

impl NodeProcessor for ArgMaxProcessor {
    type Config = ArgMaxConfig;

    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> {
        // Extract the input tensor type
        let tensor = match &node.inputs[0].ty {
            ArgType::Tensor(tensor) => tensor,
            _ => {
                return Err(ProcessError::TypeMismatch {
                    expected: "Tensor".to_string(),
                    actual: format!("{:?}", node.inputs[0].ty),
                });
            }
        };

        // Get config values before mutating node. Propagate the error so
        // invalid attributes (e.g., select_last_index out of range) surface
        // as a clean ProcessError rather than an `.expect` panic.
        let config = self.extract_config(node, opset)?;
        let keepdims = config.keepdims;

        // For burn compatibility, argmax always outputs a tensor
        // When keepdims=false, we still output a tensor but with adjusted rank
        if keepdims {
            // keepdims=true: output rank same as input rank (dimension becomes 1)
            node.outputs[0].ty = ArgType::Tensor(TensorType {
                dtype: DType::I64,
                rank: tensor.rank,
                static_shape: None,
            });
        } else if tensor.rank == 1 {
            // keepdims=false on 1D tensor: output is scalar
            node.outputs[0].ty = ArgType::ScalarTensor(DType::I64);
        } else {
            // keepdims=false on nD tensor (n > 1): output rank is input rank - 1
            node.outputs[0].ty = ArgType::Tensor(TensorType {
                dtype: DType::I64,
                rank: tensor.rank - 1,
                static_shape: None,
            });
        }

        Ok(())
    }

    fn extract_config(&self, node: &RawNode, _opset: usize) -> Result<Self::Config, ProcessError> {
        let tensor = match &node.inputs[0].ty {
            ArgType::Tensor(tensor) => tensor,
            _ => {
                return Err(ProcessError::TypeMismatch {
                    expected: "Tensor".to_string(),
                    actual: format!("{:?}", node.inputs[0].ty),
                });
            }
        };

        let mut axis: i64 = 0;
        let mut keepdims = true;
        let mut select_last_index = false;

        // Extract and validate attributes
        for (key, value) in node.attrs.iter() {
            match key.as_str() {
                "axis" => axis = value.clone().into_i64(),
                "keepdims" => {
                    let keepdims_val = value.clone().into_i64();

                    // Validate keepdims value
                    if keepdims_val != 0 && keepdims_val != 1 {
                        return Err(ProcessError::InvalidAttribute {
                            name: "keepdims".to_string(),
                            reason: "Only keepdims=0 or keepdims=1 is supported for argmax in burn"
                                .to_string(),
                        });
                    }

                    keepdims = keepdims_val != 0;
                }
                "select_last_index" => {
                    let v = value.clone().into_i64();
                    if v != 0 && v != 1 {
                        return Err(ProcessError::InvalidAttribute {
                            name: "select_last_index".to_string(),
                            reason: "Only select_last_index=0 or 1 is supported".to_string(),
                        });
                    }
                    select_last_index = v != 0;
                }
                _ => {
                    // Unknown attributes are ignored (could add warning here)
                }
            }
        }

        if axis < 0 {
            axis += tensor.rank as i64;
        }

        let config = ArgMaxConfig::new(axis as usize, keepdims, select_last_index);
        Ok(config)
    }

    fn build_node(&self, builder: RawNode, opset: usize) -> Node {
        // `build_node` runs after `infer_types`, which already calls
        // `extract_config` and propagates errors via `?`. Any config that
        // reaches here has therefore already been validated, so the
        // `.expect` is only reachable if a processor is constructed
        // out-of-pipeline or if a future validation is added to
        // `extract_config` without being mirrored in `infer_types`. The
        // trait signature `fn build_node(...) -> Node` can't return
        // `Result`, so this is the least bad option; lifting the trait
        // to `Result<Node, ProcessError>` would require a codebase-wide
        // refactor across every processor.
        let config = self
            .extract_config(&builder, opset)
            .expect("Config extraction failed");

        Node::ArgMax(ArgMaxNode {
            name: builder.name,
            inputs: builder.inputs,
            outputs: builder.outputs,
            config,
        })
    }
}

#[cfg(test)]
mod tests {
    #![allow(clippy::bool_assert_comparison)]

    use super::*;
    use crate::ir::{Argument, DType, NodeType};
    use crate::node::test_utils::TestNodeBuilder;

    fn create_test_node(axis: i64, select_last_index: i64, keepdims: i64) -> RawNode {
        TestNodeBuilder::new(NodeType::ArgMax, "test_argmax")
            .input_tensor_f32("data", 3, None)
            .output_tensor_i64("output", 3, None)
            .attr_int("axis", axis)
            .attr_int("select_last_index", select_last_index)
            .attr_int("keepdims", keepdims)
            .build()
    }

    #[test]
    fn test_argmax_config_basic() {
        let mut node = create_test_node(0, 0, 1);

        let processor = ArgMaxProcessor;

        // Extract config first, then infer types
        let config = processor.extract_config(&node, 16).unwrap();

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

        assert_eq!(config.axis, 0);
        assert_eq!(config.keepdims, true);
    }

    #[test]
    fn test_argmax_config_negative_axis() {
        let mut node = create_test_node(-2, 0, 1);

        let processor = ArgMaxProcessor;

        // Extract config first, then infer types
        let config = processor.extract_config(&node, 16).unwrap();

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

        assert_eq!(config.axis, 1); // -2 + 3 = 1
        assert_eq!(config.keepdims, true);
    }

    #[test]
    fn test_argmax_config_multiple_inputs() {
        let mut node = create_test_node(0, 0, 1);
        node.inputs.push(Argument {
            name: "extra".to_string(),
            ty: ArgType::Tensor(TensorType {
                dtype: DType::F32,
                rank: 1,
                static_shape: None,
            }),
            value_source: crate::ir::ValueSource::Dynamic,
            value_store: None,
        });

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

    #[test]
    fn test_argmax_config_keepdims_supported() {
        let mut node_keepdims_0 = create_test_node(0, 0, 0);

        let processor = ArgMaxProcessor;

        // Extract config first, then infer types
        let extracted_config_0 = processor.extract_config(&node_keepdims_0, 16).unwrap();

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

        assert_eq!(extracted_config_0.axis, 0);
        assert_eq!(extracted_config_0.keepdims, false);

        let mut node_keepdims_1 = create_test_node(0, 0, 1);

        let processor = ArgMaxProcessor;

        // Extract config first, then infer types
        let extracted_config_1 = processor.extract_config(&node_keepdims_1, 16).unwrap();

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

        assert_eq!(extracted_config_1.axis, 0);
        assert_eq!(extracted_config_1.keepdims, true);
    }

    #[test]
    fn test_argmax_config_keepdims_invalid() {
        let node = create_test_node(0, 0, 2); // Invalid keepdims value

        let processor = ArgMaxProcessor;

        // Validation should fail during config extraction
        let result = processor.extract_config(&node, 16);
        assert!(matches!(result, Err(ProcessError::InvalidAttribute { .. })));
    }

    #[test]
    fn test_argmax_config_select_last_index_supported() {
        // select_last_index=1 is supported via a flip-based rewrite; the
        // config just propagates the flag.
        let node = create_test_node(0, 1, 1);
        let processor = ArgMaxProcessor;
        let config = processor.extract_config(&node, 16).unwrap();
        assert_eq!(config.select_last_index, true);

        let node0 = create_test_node(0, 0, 1);
        let config0 = processor.extract_config(&node0, 16).unwrap();
        assert_eq!(config0.select_last_index, false);
    }

    #[test]
    fn test_argmax_config_select_last_index_invalid() {
        // Only 0 and 1 are allowed values.
        let mut node = create_test_node(0, 1, 1);
        node.attrs.insert(
            "select_last_index".to_string(),
            crate::ir::AttributeValue::Int64(2),
        );
        let processor = ArgMaxProcessor;
        let result = processor.extract_config(&node, 16);
        assert!(matches!(result, Err(ProcessError::InvalidAttribute { .. })));
    }

    #[test]
    fn test_argmax_update_outputs_keepdims_0() {
        // Test argmax with keepdims=0 - output rank should be reduced but minimum 1 for burn
        let mut node = TestNodeBuilder::new(NodeType::ArgMax, "test_argmax_keepdims_0")
            .attr_int("axis", 1)
            .attr_int("keepdims", 0)
            .input_tensor_f32("data", 2, None) // 2D input
            .output_tensor_i64("output", 2, None) // Will be updated by processor
            .build();

        let processor = ArgMaxProcessor;

        // Extract config first, then infer types
        let _config = processor.extract_config(&node, 16).unwrap();

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

        // Should output tensor with rank 1 (2 - 1 = 1, max(1, 1) = 1)
        match &node.outputs[0].ty {
            ArgType::Tensor(tensor) => {
                assert_eq!(tensor.rank, 1);
                assert_eq!(tensor.dtype, crate::ir::DType::I64);
            }
            other => panic!("Expected tensor output, got {:?}", other),
        }
    }

    #[test]
    fn test_argmax_update_outputs_keepdims_1() {
        // Test argmax with keepdims=1 - output rank should be same as input
        let mut node = TestNodeBuilder::new(NodeType::ArgMax, "test_argmax_keepdims_1")
            .attr_int("axis", 0)
            .attr_int("keepdims", 1)
            .input_tensor_f32("data", 3, None) // 3D input
            .output_tensor_i64("output", 3, None) // Will be updated by processor
            .build();

        let processor = ArgMaxProcessor;

        // Extract config first, then infer types
        let _config = processor.extract_config(&node, 16).unwrap();

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

        // Should output tensor with same rank as input (3)
        match &node.outputs[0].ty {
            ArgType::Tensor(tensor) => {
                assert_eq!(tensor.rank, 3);
                assert_eq!(tensor.dtype, crate::ir::DType::I64);
            }
            other => panic!("Expected tensor output, got {:?}", other),
        }
    }

    #[test]
    fn test_argmax_update_outputs_keepdims_0_scalar() {
        // Test argmax with keepdims=0 on 1D tensor - should output scalar
        let mut node = TestNodeBuilder::new(NodeType::ArgMax, "test_argmax_1d_keepdims_0")
            .attr_int("axis", 0)
            .attr_int("keepdims", 0)
            .input_tensor_f32("data", 1, None) // 1D input
            .output_tensor_i64("output", 1, None) // Will be updated by processor
            .build();

        let processor = ArgMaxProcessor;

        // Extract config first, then infer types
        let _config = processor.extract_config(&node, 16).unwrap();

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

        // Should output ScalarTensor (stays on device)
        match &node.outputs[0].ty {
            ArgType::ScalarTensor(elem_type) => {
                assert_eq!(*elem_type, crate::ir::DType::I64);
            }
            other => panic!("Expected ScalarTensor output, got {:?}", other),
        }
    }
}