onnx-ir 0.21.0

//! # ConvTranspose (2D)
//!
//! 2D transposed convolution (deconvolution) operation.
//!
//! **ONNX Spec**: <https://onnx.ai/onnx/operators/onnx__ConvTranspose.html>
//!
//! ## Opset Versions
//! - **Opset 1**: Initial version with basic transposed convolution support
//! - **Opset 11**: No changes to ConvTranspose operator itself (broader ONNX updates)
//!
//! ## Implementation Notes
//! - Weight tensor layout: Implementation expects [out_channels, in_channels, kernel_h, kernel_w]
//! - Padding order: ONNX `pads` format is [H_begin, W_begin, H_end, W_end]

use derive_new::new;
use onnx_ir_derive::NodeBuilder;

use crate::ir::{Argument, Node, RawNode};

use crate::processor::{
    InputSpec, NodeProcessor, NodeSpec, OutputPreferences, OutputSpec, ProcessError,
};

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

/// Configuration for ConvTranspose2d operations.
#[derive(Debug, Clone, PartialEq, Eq, new)]
#[allow(clippy::too_many_arguments)]
pub struct ConvTranspose2dConfig {
    /// Size of the kernel.
    pub kernel_size: [usize; 2],
    /// Stride of the convolutional kernel.
    pub stride: [usize; 2],
    /// Dilation of the convolutional kernel.
    pub dilation: [usize; 2],
    /// Padding.
    pub padding: [usize; 2],
    /// Output padding.
    pub padding_out: [usize; 2],
    /// Groups.
    pub groups: usize,
}

pub(crate) struct Convtranspose2dProcessor;

impl NodeProcessor for Convtranspose2dProcessor {
    type Config = ConvTranspose2dConfig;

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

    fn lift_constants(&self, node: &mut RawNode, _opset: usize) -> Result<(), ProcessError> {
        // Lift weight (input[1]) and optional bias (input[2])
        if node.inputs.len() > 1 && node.inputs[1].is_constant() {
            node.inputs[1].to_static()?;
        }
        if node.inputs.len() > 2 && node.inputs[2].is_constant() {
            node.inputs[2].to_static()?;
        }

        Ok(())
    }

    fn infer_types(
        &self,
        node: &mut RawNode,
        _opset: usize,
        _output_preferences: &OutputPreferences,
    ) -> Result<(), ProcessError> {
        // Validate attributes before extracting config
        for (key, value) in node.attrs.iter() {
            match key.as_str() {
                "kernel_shape" | "strides" | "pads" | "dilations" | "group" | "output_padding" => {}
                "auto_pad" => {
                    let auto_pad = value.clone().into_string();
                    if auto_pad != "NOTSET" {
                        return Err(ProcessError::InvalidAttribute {
                            name: "auto_pad".to_string(),
                            reason: format!("Unsupported 'auto_pad' value: {auto_pad}"),
                        });
                    }
                }
                _ => {
                    return Err(ProcessError::InvalidAttribute {
                        name: key.clone(),
                        reason: format!("Unexpected attribute for ConvTranspose2d: {key}"),
                    });
                }
            }
        }

        // Output type inference
        crate::processor::same_as_input(node);

        Ok(())
    }

    fn extract_config(&self, node: &RawNode, _opset: usize) -> Result<Self::Config, ProcessError> {
        let mut kernel_shape = Vec::new();
        let mut stride = vec![1, 1]; // Default stride to 1
        let mut pads = vec![0, 0, 0, 0]; // Default padding to 0
        let mut dilations = vec![1, 1]; // Default dilation to 1
        let mut group: usize = 1; // Default group to 1
        let mut output_padding = vec![0, 0]; // Default output padding to 0

        // Extract attributes
        for (key, value) in node.attrs.iter() {
            match key.as_str() {
                "kernel_shape" => kernel_shape = value.clone().into_i64s(),
                "strides" => stride = value.clone().into_i64s(),
                "pads" => pads = value.clone().into_i64s(),
                "dilations" => dilations = value.clone().into_i64s(),
                "group" => group = value.clone().into_i64() as usize,
                "output_padding" => output_padding = value.clone().into_i64s(),
                "auto_pad" => {}
                _ => {}
            }
        }

        // ONNX pads format: [H_begin, W_begin, H_end, W_end] = [top, left, bottom, right]
        let [top, left, bottom, right] = [pads[0], pads[1], pads[2], pads[3]];
        if left < 0 || top < 0 || right < 0 || bottom < 0 {
            return Err(ProcessError::Custom(
                "Negative pad values are not supported".to_string(),
            ));
        } else if (left != right) || (top != bottom) {
            return Err(ProcessError::Custom(
                "Asymmetric padding is not supported".to_string(),
            ));
        }

        let kernel_size = if kernel_shape.is_empty() {
            let weight_shape = node.inputs[1]
                .value()
                .ok_or_else(|| {
                    ProcessError::Custom(
                        "ConvTranspose2d: weight tensor must be present".to_string(),
                    )
                })?
                .shape
                .to_vec();

            if weight_shape.len() != 4 {
                return Err(ProcessError::Custom(format!(
                    "expected to infer kernel shape from a weight tensor of rank 4 but got shape {weight_shape:?}"
                )));
            }

            [weight_shape[2], weight_shape[3]]
        } else {
            [kernel_shape[0] as _, kernel_shape[1] as _]
        };

        let config = ConvTranspose2dConfig::new(
            kernel_size,
            [stride[0] as usize, stride[1] as usize],
            [dilations[0] as usize, dilations[1] as usize],
            [pads[0] as usize, pads[1] as usize],
            [output_padding[0] as usize, output_padding[1] as usize],
            group,
        );

        Ok(config)
    }

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

        Node::ConvTranspose2d(ConvTranspose2dNode {
            name: builder.name,
            inputs: builder.inputs,
            outputs: builder.outputs,
            config,
        })
    }
}

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

    #[allow(clippy::too_many_arguments)]
    fn create_test_node(
        kernel_shape: Vec<i64>,
        strides: Vec<i64>,
        pads: Vec<i64>,
        dilations: Vec<i64>,
        output_padding: Vec<i64>,
        group: i64,
        has_bias: bool,
        auto_pad: Option<&str>,
    ) -> TestNodeBuilder {
        // Create weight tensor data
        // ONNX ConvTranspose weight: [in_channels, out_channels/groups, k_h, k_w]
        let weight_shape = vec![2, 4, 2, 2]; // [C=in_channels, M/groups=out_channels/groups, k_h, k_w]
        let weight_data = vec![0.0; 32]; // 2*4*2*2 = 32

        let has_kernel_shape = !kernel_shape.is_empty();

        // Start building the node with input and weight
        let mut builder = TestNodeBuilder::new(NodeType::ConvTranspose2d, "test_convtranspose2d")
            .input_tensor_f32("data", 4, None)
            .input_tensor_f32_data("weight", weight_data, weight_shape)
            .output_tensor_f32("output", 4, None);

        // Add bias if needed
        if has_bias {
            builder = builder.input_tensor_f32("bias", 1, None);
        }

        if has_kernel_shape {
            builder = builder.attr_ints("kernel_shape", kernel_shape);
        }

        // Add attributes
        builder = builder
            .attr_ints("strides", strides)
            .attr_ints("pads", pads)
            .attr_ints("dilations", dilations)
            .attr_ints("output_padding", output_padding)
            .attr_int("group", group);

        if let Some(auto_pad) = auto_pad {
            builder = builder.attr_string("auto_pad", auto_pad);
        }

        builder
    }

    #[test]
    fn test_conv_transpose2d_config_basic() {
        let node = create_test_node(
            vec![2, 2],
            vec![1, 1],
            vec![0, 0, 0, 0],
            vec![1, 1],
            vec![0, 0],
            1,
            false,
            None,
        )
        .build_with_graph_data(16);
        let mut node = node;
        let processor = Convtranspose2dProcessor;
        let prefs = OutputPreferences::new();
        let config = processor.extract_config(&node, 16).unwrap();
        processor.infer_types(&mut node, 16, &prefs).unwrap();

        assert_eq!(config.kernel_size, [2, 2]);
        assert_eq!(config.stride, [1, 1]);
        assert_eq!(config.dilation, [1, 1]);
        assert_eq!(config.padding, [0, 0]);
        assert_eq!(config.padding_out, [0, 0]);
        assert_eq!(config.groups, 1);
    }

    #[test]
    fn test_conv_transpose2d_config_with_padding() {
        let node = create_test_node(
            vec![3, 3],
            vec![2, 2],
            vec![1, 1, 1, 1],
            vec![1, 1],
            vec![0, 0],
            1,
            false,
            None,
        )
        .build_with_graph_data(16);
        let mut node = node;
        let processor = Convtranspose2dProcessor;
        let prefs = OutputPreferences::new();
        let config = processor.extract_config(&node, 16).unwrap();
        processor.infer_types(&mut node, 16, &prefs).unwrap();

        assert_eq!(config.padding, [1, 1]);
        assert_eq!(config.stride, [2, 2]);
    }

    #[test]
    fn test_conv_transpose2d_config_with_output_padding() {
        let node = create_test_node(
            vec![2, 2],
            vec![2, 2],
            vec![0, 0, 0, 0],
            vec![1, 1],
            vec![1, 1],
            1,
            false,
            None,
        )
        .build_with_graph_data(16);
        let mut node = node;
        let processor = Convtranspose2dProcessor;
        let prefs = OutputPreferences::new();
        let config = processor.extract_config(&node, 16).unwrap();
        processor.infer_types(&mut node, 16, &prefs).unwrap();

        assert_eq!(config.padding_out, [1, 1]);
    }

    #[test]
    fn test_conv_transpose2d_config_with_groups() {
        let node = create_test_node(
            vec![2, 2],
            vec![1, 1],
            vec![0, 0, 0, 0],
            vec![1, 1],
            vec![0, 0],
            2,
            false,
            None,
        )
        .build_with_graph_data(16);
        let mut node = node;
        let processor = Convtranspose2dProcessor;
        let prefs = OutputPreferences::new();
        let config = processor.extract_config(&node, 16).unwrap();
        processor.infer_types(&mut node, 16, &prefs).unwrap();

        assert_eq!(config.groups, 2);
    }

    #[test]
    fn test_conv_transpose2d_config_with_asymmetric_padding() {
        let node = create_test_node(
            vec![2, 2],
            vec![1, 1],
            vec![1, 1, 2, 2],
            vec![1, 1],
            vec![0, 0],
            1,
            false,
            None,
        )
        .build_with_graph_data(16);
        let processor = Convtranspose2dProcessor;
        let result = processor.extract_config(&node, 16);
        assert!(
            matches!(result, Err(ProcessError::Custom(ref msg)) if msg.contains("Asymmetric padding is not supported"))
        );
    }

    #[test]
    fn test_conv_transpose2d_config_autopad_not_set() {
        let node = create_test_node(
            vec![2, 2],
            vec![1, 1],
            vec![0, 0, 0, 0],
            vec![1, 1],
            vec![0, 0],
            1,
            false,
            Some("NOTSET"),
        )
        .build_with_graph_data(16);
        let mut node = node;
        let processor = Convtranspose2dProcessor;
        let prefs = OutputPreferences::new();
        let config = processor.extract_config(&node, 16).unwrap();
        processor.infer_types(&mut node, 16, &prefs).unwrap();

        assert_eq!(config.kernel_size, [2, 2]);
        assert_eq!(config.stride, [1, 1]);
        assert_eq!(config.dilation, [1, 1]);
        assert_eq!(config.padding, [0, 0]);
        assert_eq!(config.padding_out, [0, 0]);
        assert_eq!(config.groups, 1);
    }

    #[test]
    fn test_conv_transpose2d_config_autopad_not_supported() {
        let node = create_test_node(
            vec![2, 2],
            vec![1, 1],
            vec![0, 0, 0, 0],
            vec![1, 1],
            vec![0, 0],
            1,
            false,
            Some("SAME_UPPER"),
        )
        .build_with_graph_data(16);
        let mut node = node;
        let processor = Convtranspose2dProcessor;
        let prefs = OutputPreferences::new();
        let result = processor.infer_types(&mut node, 16, &prefs);
        assert!(matches!(result, Err(ProcessError::InvalidAttribute { .. })));
    }

    #[test]
    fn test_conv_transpose2d_config_kernel_shape_not_set() {
        let node = create_test_node(
            vec![],
            vec![1, 1],
            vec![0, 0, 0, 0],
            vec![1, 1],
            vec![0, 0],
            1,
            false,
            None,
        )
        .build_with_graph_data(16);
        let mut node = node;
        let processor = Convtranspose2dProcessor;
        let prefs = OutputPreferences::new();
        let config = processor.extract_config(&node, 16).unwrap();
        processor.infer_types(&mut node, 16, &prefs).unwrap();

        assert_eq!(config.kernel_size, [2, 2]); // Inferred via weight tensor shape
    }
}