torsh_tensor/

conv.rs

//! Convolution and signal processing operations for tensors

use crate::{FloatElement, Tensor};
use torsh_core::error::{Result, TorshError};
use torsh_core::TensorElement;

impl<T: FloatElement> Tensor<T> {
    /// im2col (unfold): gather the sliding windows of an `[N, C, H, W]` tensor
    /// into per-group patch matrices.
    ///
    /// The result is shaped `[groups, N * out_h * out_w, (C / groups) * kh * kw]`:
    /// one row per (batch item, output position), one column per kernel tap, so a
    /// convolution is a single batched matrix product against the reshaped
    /// weight. Positions that fall in the zero padding contribute zeros.
    ///
    /// # Autograd
    ///
    /// The gather is recorded, and its backward pass is col2im: every patch
    /// element scatter-adds its gradient back onto the input element it was read
    /// from, so overlapping windows accumulate. This is what lets a convolution
    /// built as `im2col_2d(...).matmul(weight)` produce exact gradients for both
    /// the input and the weight.
    ///
    /// # Arguments
    /// * `kernel` - `(kh, kw)`
    /// * `stride` - `(sh, sw)`
    /// * `padding` - `(ph, pw)` zero padding applied on both sides of each axis
    /// * `dilation` - `(dh, dw)` spacing between kernel taps
    /// * `groups` - number of channel groups; `C` must be divisible by it
    pub fn im2col_2d(
        &self,
        kernel: (usize, usize),
        stride: (usize, usize),
        padding: (usize, usize),
        dilation: (usize, usize),
        groups: usize,
    ) -> Result<Self> {
        let shape_binding = self.shape();
        let dims = shape_binding.dims();
        if dims.len() != 4 {
            return Err(TorshError::InvalidShape(format!(
                "im2col_2d expects a 4-D [batch, channels, height, width] tensor, got {}-D",
                dims.len()
            )));
        }
        let (batch, channels, height, width) = (dims[0], dims[1], dims[2], dims[3]);

        if groups == 0 || channels % groups != 0 {
            return Err(TorshError::InvalidArgument(format!(
                "im2col_2d: {channels} channels cannot be split into {groups} groups"
            )));
        }
        if kernel.0 == 0 || kernel.1 == 0 || stride.0 == 0 || stride.1 == 0 {
            return Err(TorshError::InvalidArgument(
                "im2col_2d: kernel and stride must be positive".to_string(),
            ));
        }
        if dilation.0 == 0 || dilation.1 == 0 {
            return Err(TorshError::InvalidArgument(
                "im2col_2d: dilation must be positive".to_string(),
            ));
        }

        let span_h = dilation.0 * (kernel.0 - 1) + 1;
        let span_w = dilation.1 * (kernel.1 - 1) + 1;
        let padded_h = height + 2 * padding.0;
        let padded_w = width + 2 * padding.1;
        if padded_h < span_h || padded_w < span_w {
            return Err(TorshError::InvalidShape(format!(
                "im2col_2d: a {span_h}x{span_w} kernel span does not fit a \
                 {padded_h}x{padded_w} padded input"
            )));
        }
        let out_h = (padded_h - span_h) / stride.0 + 1;
        let out_w = (padded_w - span_w) / stride.1 + 1;

        let per_group = channels / groups;
        let patch_len = per_group * kernel.0 * kernel.1;
        let rows = batch * out_h * out_w;

        let input_data = self.to_vec()?;
        let mut patches = vec![<T as num_traits::Zero>::zero(); groups * rows * patch_len];

        for batch_index in 0..batch {
            for group in 0..groups {
                let group_base = group * rows * patch_len;
                for out_y in 0..out_h {
                    for out_x in 0..out_w {
                        let row = group_base
                            + (batch_index * out_h * out_w + out_y * out_w + out_x) * patch_len;
                        for channel in 0..per_group {
                            let global_channel = group * per_group + channel;
                            let channel_base =
                                (batch_index * channels + global_channel) * height * width;
                            for ky in 0..kernel.0 {
                                let in_y = out_y * stride.0 + ky * dilation.0;
                                if in_y < padding.0 {
                                    continue;
                                }
                                let in_y = in_y - padding.0;
                                if in_y >= height {
                                    continue;
                                }
                                for kx in 0..kernel.1 {
                                    let in_x = out_x * stride.1 + kx * dilation.1;
                                    if in_x < padding.1 {
                                        continue;
                                    }
                                    let in_x = in_x - padding.1;
                                    if in_x >= width {
                                        continue;
                                    }
                                    let column = (channel * kernel.0 + ky) * kernel.1 + kx;
                                    patches[row + column] =
                                        input_data[channel_base + in_y * width + in_x];
                                }
                            }
                        }
                    }
                }
            }
        }

        let mut result = Self::from_data(patches, vec![groups, rows, patch_len], self.device())?;

        if crate::should_record_grad(self.requires_grad) {
            result.requires_grad = true;
            result.operation = crate::core_ops::Operation::Im2Col {
                input: std::sync::Arc::new(self.clone()),
                config: crate::core_ops::Im2ColConfig {
                    input_shape: [batch, channels, height, width],
                    kernel,
                    stride,
                    padding,
                    dilation,
                    groups,
                    output: (out_h, out_w),
                },
            };
        }

        Ok(result)
    }

    /// Broadcast a per-output-channel bias across a contiguous convolution output.
    ///
    /// `output_data` is laid out row-major as `[batch, out_channels, spatial...]`,
    /// so each `(batch, channel)` pair owns a contiguous run of `spatial_size`
    /// elements. Every element of channel `c`'s run receives `bias_data[c]` added
    /// to it, identically for every batch element. This is the standard channel-wise
    /// broadcast of a `[C_out]` bias vector across an `[N, C_out, *spatial]` tensor.
    ///
    /// Iterating contiguous channel-sized chunks (instead of recomputing a flat
    /// index per element) keeps the access pattern cache-friendly and lets the
    /// compiler auto-vectorize the inner scalar add.
    fn add_channel_bias(
        output_data: &mut [T],
        bias_data: &[T],
        out_channels: usize,
        spatial_size: usize,
    ) {
        // `chunks_mut(0)` panics, and an empty buffer has nothing to broadcast.
        if spatial_size == 0 || out_channels == 0 {
            return;
        }

        // Each chunk is one (batch, channel) block; chunks advance in row-major
        // order, so the channel is the chunk index modulo the channel count.
        for (block_index, block) in output_data.chunks_mut(spatial_size).enumerate() {
            let channel = block_index % out_channels;
            let bias = bias_data[channel];
            for value in block.iter_mut() {
                *value = *value + bias;
            }
        }
    }

    /// 1D convolution operation
    pub fn conv1d(
        &self,
        weight: &Self,
        bias: Option<&Self>,
        stride: usize,
        padding: usize,
        dilation: usize,
        groups: usize,
    ) -> Result<Self> {
        // Input shape: (N, C_in, L)
        // Weight shape: (C_out, C_in/groups, kernel_size)
        // Output shape: (N, C_out, L_out)

        let input_shape_obj = self.shape();
        let input_shape = input_shape_obj.dims();
        let weight_shape_obj = weight.shape();
        let weight_shape = weight_shape_obj.dims();

        if input_shape.len() != 3 {
            return Err(TorshError::InvalidArgument(format!(
                "Expected 3D input tensor for conv1d, got {}D",
                input_shape.len()
            )));
        }

        if weight_shape.len() != 3 {
            return Err(TorshError::InvalidArgument(format!(
                "Expected 3D weight tensor for conv1d, got {}D",
                weight_shape.len()
            )));
        }

        let batch_size = input_shape[0];
        let in_channels = input_shape[1];
        let input_length = input_shape[2];

        let out_channels = weight_shape[0];
        let kernel_size = weight_shape[2];

        // Check groups
        if in_channels % groups != 0 || out_channels % groups != 0 {
            return Err(TorshError::InvalidArgument(
                "in_channels and out_channels must be divisible by groups".to_string(),
            ));
        }

        if weight_shape[1] != in_channels / groups {
            return Err(TorshError::InvalidArgument(format!(
                "Weight tensor has wrong number of input channels: expected {}, got {}",
                in_channels / groups,
                weight_shape[1]
            )));
        }

        // Calculate output length
        let effective_kernel = (kernel_size - 1) * dilation + 1;
        let padded_length = input_length + 2 * padding;
        let output_length = (padded_length - effective_kernel) / stride + 1;

        // Initialize output
        let mut output_data =
            vec![<T as TensorElement>::zero(); batch_size * out_channels * output_length];

        // Perform convolution. Both operands are borrowed out of storage once
        // for the whole op instead of taking a read guard per multiply-add.
        self.with_operand_slices(weight, |input_data, weight_data| {
            for n in 0..batch_size {
                for g in 0..groups {
                    let out_ch_start = g * (out_channels / groups);
                    let out_ch_end = (g + 1) * (out_channels / groups);
                    let in_ch_start = g * (in_channels / groups);
                    let in_ch_end = (g + 1) * (in_channels / groups);

                    for oc in out_ch_start..out_ch_end {
                        for ol in 0..output_length {
                            let mut sum = <T as TensorElement>::zero();

                            for ic in in_ch_start..in_ch_end {
                                let ic_rel = ic - in_ch_start;
                                for k in 0..kernel_size {
                                    let il = (ol * stride + k * dilation) as i32 - padding as i32;

                                    if il >= 0 && (il as usize) < input_length {
                                        let input_idx = n * in_channels * input_length
                                            + ic * input_length
                                            + il as usize;
                                        let weight_idx = oc * (in_channels / groups) * kernel_size
                                            + ic_rel * kernel_size
                                            + k;

                                        let input_val =
                                            *input_data.get(input_idx).ok_or_else(|| {
                                                TorshError::IndexOutOfBounds {
                                                    index: input_idx,
                                                    size: input_data.len(),
                                                }
                                            })?;
                                        let weight_val =
                                            *weight_data.get(weight_idx).ok_or_else(|| {
                                                TorshError::IndexOutOfBounds {
                                                    index: weight_idx,
                                                    size: weight_data.len(),
                                                }
                                            })?;
                                        sum = sum + input_val * weight_val;
                                    }
                                }
                            }

                            let output_idx =
                                n * out_channels * output_length + oc * output_length + ol;
                            output_data[output_idx] = sum;
                        }
                    }
                }
            }
            Ok(())
        })?;

        // Create output tensor
        let mut output = Tensor::from_data(
            output_data,
            vec![batch_size, out_channels, output_length],
            self.device(),
        )?;

        // Add bias if provided
        if let Some(b) = bias {
            if b.shape().dims() != [out_channels] {
                return Err(TorshError::InvalidArgument(format!(
                    "Bias must have shape [{}], got {:?}",
                    out_channels,
                    b.shape().dims()
                )));
            }

            // Broadcast the per-output-channel bias across every spatial position
            // of each channel, for every batch element.
            let bias_data = b.to_vec()?;
            let mut output_data = output.to_vec()?;
            Self::add_channel_bias(&mut output_data, &bias_data, out_channels, output_length);

            // Recreate tensor with modified data
            output = Tensor::from_data(
                output_data,
                vec![batch_size, out_channels, output_length],
                self.device(),
            )?;
        }

        // Track operation for autograd
        if crate::should_record_grad(
            self.requires_grad
                || weight.requires_grad
                || (bias.is_some() && bias.expect("bias checked with is_some").requires_grad),
        ) {
            use std::sync::Arc;
            output.requires_grad = true;
            output.operation = crate::Operation::Custom(
                "conv1d".to_string(),
                vec![
                    Arc::downgrade(&Arc::new(self.clone())),
                    Arc::downgrade(&Arc::new(weight.clone())),
                ],
            );
        }

        Ok(output)
    }

    /// 2D convolution operation
    pub fn conv2d(
        &self,
        weight: &Self,
        bias: Option<&Self>,
        stride: (usize, usize),
        padding: (usize, usize),
        dilation: (usize, usize),
        groups: usize,
    ) -> Result<Self> {
        // Input shape: (N, C_in, H, W)
        // Weight shape: (C_out, C_in/groups, kernel_h, kernel_w)
        // Output shape: (N, C_out, H_out, W_out)

        let input_shape_obj = self.shape();
        let input_shape = input_shape_obj.dims();
        let weight_shape_obj = weight.shape();
        let weight_shape = weight_shape_obj.dims();

        if input_shape.len() != 4 {
            return Err(TorshError::InvalidArgument(format!(
                "Expected 4D input tensor for conv2d, got {}D",
                input_shape.len()
            )));
        }

        if weight_shape.len() != 4 {
            return Err(TorshError::InvalidArgument(format!(
                "Expected 4D weight tensor for conv2d, got {}D",
                weight_shape.len()
            )));
        }

        let batch_size = input_shape[0];
        let in_channels = input_shape[1];
        let input_height = input_shape[2];
        let input_width = input_shape[3];

        let out_channels = weight_shape[0];
        let kernel_height = weight_shape[2];
        let kernel_width = weight_shape[3];

        // Check groups
        if in_channels % groups != 0 || out_channels % groups != 0 {
            return Err(TorshError::InvalidArgument(
                "in_channels and out_channels must be divisible by groups".to_string(),
            ));
        }

        if weight_shape[1] != in_channels / groups {
            return Err(TorshError::InvalidArgument(format!(
                "Weight tensor has wrong number of input channels: expected {}, got {}",
                in_channels / groups,
                weight_shape[1]
            )));
        }

        // Calculate output dimensions
        let effective_kernel_h = (kernel_height - 1) * dilation.0 + 1;
        let effective_kernel_w = (kernel_width - 1) * dilation.1 + 1;
        let padded_height = input_height + 2 * padding.0;
        let padded_width = input_width + 2 * padding.1;
        let output_height = (padded_height - effective_kernel_h) / stride.0 + 1;
        let output_width = (padded_width - effective_kernel_w) / stride.1 + 1;

        // Initialize output
        let mut output_data = vec![
            <T as TensorElement>::zero();
            batch_size * out_channels * output_height * output_width
        ];

        // Perform convolution. Both operands are borrowed out of storage once
        // for the whole op — no per-operand copy, no per-MAC lock.
        self.with_operand_slices(weight, |self_data, weight_data| {
            for n in 0..batch_size {
                for g in 0..groups {
                    let out_ch_start = g * (out_channels / groups);
                    let out_ch_end = (g + 1) * (out_channels / groups);
                    let in_ch_start = g * (in_channels / groups);
                    let in_ch_end = (g + 1) * (in_channels / groups);

                    for oc in out_ch_start..out_ch_end {
                        for oh in 0..output_height {
                            for ow in 0..output_width {
                                let mut sum = <T as TensorElement>::zero();

                                for ic in in_ch_start..in_ch_end {
                                    let ic_rel = ic - in_ch_start;
                                    for kh in 0..kernel_height {
                                        for kw in 0..kernel_width {
                                            let ih = (oh * stride.0 + kh * dilation.0) as i32
                                                - padding.0 as i32;
                                            let iw = (ow * stride.1 + kw * dilation.1) as i32
                                                - padding.1 as i32;

                                            if ih >= 0
                                                && (ih as usize) < input_height
                                                && iw >= 0
                                                && (iw as usize) < input_width
                                            {
                                                let input_idx =
                                                    n * in_channels * input_height * input_width
                                                        + ic * input_height * input_width
                                                        + ih as usize * input_width
                                                        + iw as usize;
                                                let weight_idx = oc
                                                    * (in_channels / groups)
                                                    * kernel_height
                                                    * kernel_width
                                                    + ic_rel * kernel_height * kernel_width
                                                    + kh * kernel_width
                                                    + kw;

                                                sum = sum
                                                    + self_data[input_idx]
                                                        * weight_data[weight_idx];
                                            }
                                        }
                                    }
                                }

                                let output_idx = n * out_channels * output_height * output_width
                                    + oc * output_height * output_width
                                    + oh * output_width
                                    + ow;
                                output_data[output_idx] = sum;
                            }
                        }
                    }
                }
            }
            Ok(())
        })?;

        // Create output tensor
        let mut output = Tensor::from_data(
            output_data,
            vec![batch_size, out_channels, output_height, output_width],
            self.device(),
        )?;

        // Add bias if provided
        if let Some(b) = bias {
            if b.shape().dims() != [out_channels] {
                return Err(TorshError::InvalidArgument(format!(
                    "Bias must have shape [{}], got {:?}",
                    out_channels,
                    b.shape().dims()
                )));
            }

            // Broadcast the per-output-channel bias across the (H, W) spatial map
            // of each channel, for every batch element.
            let bias_data = b.to_vec()?;
            let mut output_data = output.to_vec()?;
            Self::add_channel_bias(
                &mut output_data,
                &bias_data,
                out_channels,
                output_height * output_width,
            );

            // Create new output tensor with bias added
            output = Tensor::from_data(
                output_data,
                vec![batch_size, out_channels, output_height, output_width],
                self.device(),
            )?;
        }

        // Track operation for autograd
        if crate::should_record_grad(
            self.requires_grad
                || weight.requires_grad
                || (bias.is_some() && bias.expect("bias checked with is_some").requires_grad),
        ) {
            use std::sync::Arc;
            output.requires_grad = true;
            output.operation = crate::Operation::Custom(
                "conv2d".to_string(),
                vec![
                    Arc::downgrade(&Arc::new(self.clone())),
                    Arc::downgrade(&Arc::new(weight.clone())),
                ],
            );
        }

        Ok(output)
    }

    /// 3D convolution operation
    pub fn conv3d(
        &self,
        weight: &Self,
        bias: Option<&Self>,
        stride: (usize, usize, usize),
        padding: (usize, usize, usize),
        dilation: (usize, usize, usize),
        groups: usize,
    ) -> Result<Self> {
        // Input shape: (N, C_in, D, H, W)
        // Weight shape: (C_out, C_in/groups, kernel_d, kernel_h, kernel_w)
        // Output shape: (N, C_out, D_out, H_out, W_out)

        let input_shape_obj = self.shape();
        let input_shape = input_shape_obj.dims();
        let weight_shape_obj = weight.shape();
        let weight_shape = weight_shape_obj.dims();

        if input_shape.len() != 5 {
            return Err(TorshError::InvalidArgument(format!(
                "Expected 5D input tensor for conv3d, got {}D",
                input_shape.len()
            )));
        }

        if weight_shape.len() != 5 {
            return Err(TorshError::InvalidArgument(format!(
                "Expected 5D weight tensor for conv3d, got {}D",
                weight_shape.len()
            )));
        }

        let batch_size = input_shape[0];
        let in_channels = input_shape[1];
        let input_depth = input_shape[2];
        let input_height = input_shape[3];
        let input_width = input_shape[4];

        let out_channels = weight_shape[0];
        let kernel_depth = weight_shape[2];
        let kernel_height = weight_shape[3];
        let kernel_width = weight_shape[4];

        // Check groups
        if in_channels % groups != 0 || out_channels % groups != 0 {
            return Err(TorshError::InvalidArgument(
                "in_channels and out_channels must be divisible by groups".to_string(),
            ));
        }

        if weight_shape[1] != in_channels / groups {
            return Err(TorshError::InvalidArgument(format!(
                "Weight tensor has wrong number of input channels: expected {}, got {}",
                in_channels / groups,
                weight_shape[1]
            )));
        }

        // Calculate output dimensions
        let effective_kernel_d = (kernel_depth - 1) * dilation.0 + 1;
        let effective_kernel_h = (kernel_height - 1) * dilation.1 + 1;
        let effective_kernel_w = (kernel_width - 1) * dilation.2 + 1;
        let padded_depth = input_depth + 2 * padding.0;
        let padded_height = input_height + 2 * padding.1;
        let padded_width = input_width + 2 * padding.2;
        let output_depth = (padded_depth - effective_kernel_d) / stride.0 + 1;
        let output_height = (padded_height - effective_kernel_h) / stride.1 + 1;
        let output_width = (padded_width - effective_kernel_w) / stride.2 + 1;

        // Initialize output
        let output_size = batch_size * out_channels * output_depth * output_height * output_width;
        let mut output_data = vec![<T as TensorElement>::zero(); output_size];

        let self_data = self.to_vec()?;
        let weight_data = weight.to_vec()?;

        // Perform convolution
        for n in 0..batch_size {
            for g in 0..groups {
                let out_ch_start = g * (out_channels / groups);
                let out_ch_end = (g + 1) * (out_channels / groups);
                let in_ch_start = g * (in_channels / groups);
                let in_ch_end = (g + 1) * (in_channels / groups);

                for oc in out_ch_start..out_ch_end {
                    for od in 0..output_depth {
                        for oh in 0..output_height {
                            for ow in 0..output_width {
                                let mut sum = <T as TensorElement>::zero();

                                for ic in in_ch_start..in_ch_end {
                                    let ic_rel = ic - in_ch_start;
                                    for kd in 0..kernel_depth {
                                        for kh in 0..kernel_height {
                                            for kw in 0..kernel_width {
                                                let id = (od * stride.0 + kd * dilation.0) as i32
                                                    - padding.0 as i32;
                                                let ih = (oh * stride.1 + kh * dilation.1) as i32
                                                    - padding.1 as i32;
                                                let iw = (ow * stride.2 + kw * dilation.2) as i32
                                                    - padding.2 as i32;

                                                if id >= 0
                                                    && (id as usize) < input_depth
                                                    && ih >= 0
                                                    && (ih as usize) < input_height
                                                    && iw >= 0
                                                    && (iw as usize) < input_width
                                                {
                                                    let input_idx = n
                                                        * in_channels
                                                        * input_depth
                                                        * input_height
                                                        * input_width
                                                        + ic * input_depth
                                                            * input_height
                                                            * input_width
                                                        + id as usize * input_height * input_width
                                                        + ih as usize * input_width
                                                        + iw as usize;
                                                    let weight_idx = oc
                                                        * (in_channels / groups)
                                                        * kernel_depth
                                                        * kernel_height
                                                        * kernel_width
                                                        + ic_rel
                                                            * kernel_depth
                                                            * kernel_height
                                                            * kernel_width
                                                        + kd * kernel_height * kernel_width
                                                        + kh * kernel_width
                                                        + kw;

                                                    sum = sum
                                                        + self_data[input_idx]
                                                            * weight_data[weight_idx];
                                                }
                                            }
                                        }
                                    }
                                }

                                let output_idx =
                                    n * out_channels * output_depth * output_height * output_width
                                        + oc * output_depth * output_height * output_width
                                        + od * output_height * output_width
                                        + oh * output_width
                                        + ow;
                                output_data[output_idx] = sum;
                            }
                        }
                    }
                }
            }
        }

        // Create output tensor
        let mut output = Tensor::from_data(
            output_data,
            vec![
                batch_size,
                out_channels,
                output_depth,
                output_height,
                output_width,
            ],
            self.device(),
        )?;

        // Add bias if provided
        if let Some(b) = bias {
            if b.shape().dims() != [out_channels] {
                return Err(TorshError::InvalidArgument(format!(
                    "Bias must have shape [{}], got {:?}",
                    out_channels,
                    b.shape().dims()
                )));
            }

            // Broadcast the per-output-channel bias across the (D, H, W) spatial
            // volume of each channel, for every batch element.
            let bias_data = b.to_vec()?;
            let mut output_data = output.to_vec()?;
            Self::add_channel_bias(
                &mut output_data,
                &bias_data,
                out_channels,
                output_depth * output_height * output_width,
            );

            // Create new output tensor with bias added
            output = Tensor::from_data(
                output_data,
                vec![
                    batch_size,
                    out_channels,
                    output_depth,
                    output_height,
                    output_width,
                ],
                self.device(),
            )?;
        }

        // Track operation for autograd
        if crate::should_record_grad(
            self.requires_grad
                || weight.requires_grad
                || (bias.is_some() && bias.expect("bias checked with is_some").requires_grad),
        ) {
            use std::sync::Arc;
            output.requires_grad = true;
            output.operation = crate::Operation::Custom(
                "conv3d".to_string(),
                vec![
                    Arc::downgrade(&Arc::new(self.clone())),
                    Arc::downgrade(&Arc::new(weight.clone())),
                ],
            );
        }

        Ok(output)
    }

    /// Depthwise 2D convolution operation
    /// Each input channel is convolved with its own kernel independently
    pub fn depthwise_conv2d(
        &self,
        weight: &Self,
        bias: Option<&Self>,
        stride: (usize, usize),
        padding: (usize, usize),
        dilation: (usize, usize),
    ) -> Result<Self> {
        // Input shape: (N, C_in, H, W)
        // Weight shape: (C_in, 1, kernel_h, kernel_w) - each channel has its own kernel
        // Output shape: (N, C_in, H_out, W_out)

        let input_shape_obj = self.shape();
        let input_shape = input_shape_obj.dims();
        let weight_shape_obj = weight.shape();
        let weight_shape = weight_shape_obj.dims();

        if input_shape.len() != 4 {
            return Err(TorshError::InvalidArgument(format!(
                "Expected 4D input tensor for depthwise_conv2d, got {}D",
                input_shape.len()
            )));
        }

        if weight_shape.len() != 4 {
            return Err(TorshError::InvalidArgument(format!(
                "Expected 4D weight tensor for depthwise_conv2d, got {}D",
                weight_shape.len()
            )));
        }

        let batch_size = input_shape[0];
        let in_channels = input_shape[1];
        let input_height = input_shape[2];
        let input_width = input_shape[3];

        let kernel_height = weight_shape[2];
        let kernel_width = weight_shape[3];

        // For depthwise conv, weight should have shape (C_in, 1, kernel_h, kernel_w)
        if weight_shape[0] != in_channels || weight_shape[1] != 1 {
            return Err(TorshError::InvalidArgument(format!(
                "Weight tensor must have shape ({}, 1, kernel_h, kernel_w), got ({}, {}, {}, {})",
                in_channels, weight_shape[0], weight_shape[1], weight_shape[2], weight_shape[3]
            )));
        }

        // Calculate output dimensions
        let effective_kernel_h = (kernel_height - 1) * dilation.0 + 1;
        let effective_kernel_w = (kernel_width - 1) * dilation.1 + 1;
        let padded_height = input_height + 2 * padding.0;
        let padded_width = input_width + 2 * padding.1;
        let output_height = (padded_height - effective_kernel_h) / stride.0 + 1;
        let output_width = (padded_width - effective_kernel_w) / stride.1 + 1;

        // Initialize output
        let mut output_data = vec![
            <T as TensorElement>::zero();
            batch_size * in_channels * output_height * output_width
        ];

        // Perform depthwise convolution. Both operands are borrowed out of
        // storage once for the whole op instead of per multiply-add.
        self.with_operand_slices(weight, |input_data, weight_data| {
            for n in 0..batch_size {
                for c in 0..in_channels {
                    for oh in 0..output_height {
                        for ow in 0..output_width {
                            let mut sum = <T as TensorElement>::zero();

                            for kh in 0..kernel_height {
                                for kw in 0..kernel_width {
                                    let ih =
                                        (oh * stride.0 + kh * dilation.0) as i32 - padding.0 as i32;
                                    let iw =
                                        (ow * stride.1 + kw * dilation.1) as i32 - padding.1 as i32;

                                    if ih >= 0
                                        && (ih as usize) < input_height
                                        && iw >= 0
                                        && (iw as usize) < input_width
                                    {
                                        let input_idx =
                                            n * in_channels * input_height * input_width
                                                + c * input_height * input_width
                                                + ih as usize * input_width
                                                + iw as usize;
                                        let weight_idx = c * kernel_height * kernel_width
                                            + kh * kernel_width
                                            + kw;

                                        let input_val =
                                            *input_data.get(input_idx).ok_or_else(|| {
                                                TorshError::IndexOutOfBounds {
                                                    index: input_idx,
                                                    size: input_data.len(),
                                                }
                                            })?;
                                        let weight_val =
                                            *weight_data.get(weight_idx).ok_or_else(|| {
                                                TorshError::IndexOutOfBounds {
                                                    index: weight_idx,
                                                    size: weight_data.len(),
                                                }
                                            })?;
                                        sum = sum + input_val * weight_val;
                                    }
                                }
                            }

                            let output_idx = n * in_channels * output_height * output_width
                                + c * output_height * output_width
                                + oh * output_width
                                + ow;
                            output_data[output_idx] = sum;
                        }
                    }
                }
            }
            Ok(())
        })?;

        // Create output tensor
        let mut output = Tensor::from_data(
            output_data,
            vec![batch_size, in_channels, output_height, output_width],
            self.device(),
        )?;

        // Add bias if provided
        if let Some(b) = bias {
            if b.shape().dims() != [in_channels] {
                return Err(TorshError::InvalidArgument(format!(
                    "Bias must have shape [{}], got {:?}",
                    in_channels,
                    b.shape().dims()
                )));
            }

            // Broadcast the per-channel bias across the (H, W) spatial map of each
            // channel, for every batch element. Depthwise conv keeps one output
            // channel per input channel, so the channel count is `in_channels`.
            let bias_data = b.to_vec()?;
            let mut output_data = output.to_vec()?;
            Self::add_channel_bias(
                &mut output_data,
                &bias_data,
                in_channels,
                output_height * output_width,
            );

            // Create new output tensor with bias added
            output = Tensor::from_data(
                output_data,
                vec![batch_size, in_channels, output_height, output_width],
                self.device(),
            )?;
        }

        // Track operation for autograd
        if crate::should_record_grad(
            self.requires_grad
                || weight.requires_grad
                || (bias.is_some() && bias.expect("bias checked with is_some").requires_grad),
        ) {
            use std::sync::Arc;
            output.requires_grad = true;
            output.operation = crate::Operation::Custom(
                "depthwise_conv2d".to_string(),
                vec![
                    Arc::downgrade(&Arc::new(self.clone())),
                    Arc::downgrade(&Arc::new(weight.clone())),
                ],
            );
        }

        Ok(output)
    }

    /// Separable 2D convolution operation
    /// Factorized into depthwise convolution followed by pointwise (1x1) convolution
    pub fn separable_conv2d(
        &self,
        depthwise_weight: &Self,
        pointwise_weight: &Self,
        bias: Option<&Self>,
        stride: (usize, usize),
        padding: (usize, usize),
        dilation: (usize, usize),
    ) -> Result<Self> {
        // Step 1: Depthwise convolution
        let depthwise_output = self.depthwise_conv2d(
            depthwise_weight,
            None, // No bias in depthwise step
            stride,
            padding,
            dilation,
        )?;

        // Step 2: Pointwise (1x1) convolution
        let output = depthwise_output.conv2d(
            pointwise_weight,
            bias,
            (1, 1), // stride = 1 for pointwise
            (0, 0), // padding = 0 for pointwise
            (1, 1), // dilation = 1 for pointwise
            1,      // groups = 1 for pointwise
        )?;

        // Track operation for autograd
        if crate::should_record_grad(
            self.requires_grad
                || depthwise_weight.requires_grad
                || pointwise_weight.requires_grad
                || (bias.is_some() && bias.expect("bias checked with is_some").requires_grad),
        ) {
            use std::sync::Arc;
            let mut tracked_output = output;
            tracked_output.requires_grad = true;
            tracked_output.operation = crate::Operation::Custom(
                "separable_conv2d".to_string(),
                vec![
                    Arc::downgrade(&Arc::new(self.clone())),
                    Arc::downgrade(&Arc::new(depthwise_weight.clone())),
                    Arc::downgrade(&Arc::new(pointwise_weight.clone())),
                ],
            );
            Ok(tracked_output)
        } else {
            Ok(output)
        }
    }

    /// Transposed (deconvolution) 2D convolution operation
    #[allow(clippy::too_many_arguments)]
    pub fn conv_transpose2d(
        &self,
        weight: &Self,
        bias: Option<&Self>,
        stride: (usize, usize),
        padding: (usize, usize),
        output_padding: (usize, usize),
        dilation: (usize, usize),
        groups: usize,
    ) -> Result<Self> {
        // Input shape: (N, C_in, H, W)
        // Weight shape: (C_in, C_out/groups, kernel_h, kernel_w)
        // Output shape: (N, C_out, H_out, W_out)

        let input_shape_obj = self.shape();
        let input_shape = input_shape_obj.dims();
        let weight_shape_obj = weight.shape();
        let weight_shape = weight_shape_obj.dims();

        if input_shape.len() != 4 {
            return Err(TorshError::InvalidArgument(format!(
                "Expected 4D input tensor for conv_transpose2d, got {}D",
                input_shape.len()
            )));
        }

        if weight_shape.len() != 4 {
            return Err(TorshError::InvalidArgument(format!(
                "Expected 4D weight tensor for conv_transpose2d, got {}D",
                weight_shape.len()
            )));
        }

        let batch_size = input_shape[0];
        let in_channels = input_shape[1];
        let input_height = input_shape[2];
        let input_width = input_shape[3];

        let out_channels = weight_shape[1] * groups;
        let kernel_height = weight_shape[2];
        let kernel_width = weight_shape[3];

        // Check groups
        if in_channels % groups != 0 || out_channels % groups != 0 {
            return Err(TorshError::InvalidArgument(
                "in_channels and out_channels must be divisible by groups".to_string(),
            ));
        }

        if weight_shape[0] != in_channels {
            return Err(TorshError::InvalidArgument(format!(
                "Weight tensor has wrong number of input channels: expected {}, got {}",
                in_channels, weight_shape[0]
            )));
        }

        // Calculate output dimensions
        let effective_kernel_h = (kernel_height - 1) * dilation.0 + 1;
        let effective_kernel_w = (kernel_width - 1) * dilation.1 + 1;
        let output_height =
            (input_height - 1) * stride.0 - 2 * padding.0 + effective_kernel_h + output_padding.0;
        let output_width =
            (input_width - 1) * stride.1 - 2 * padding.1 + effective_kernel_w + output_padding.1;

        // Initialize output
        let mut output_data = vec![
            <T as TensorElement>::zero();
            batch_size * out_channels * output_height * output_width
        ];

        let self_data = self.to_vec()?;
        let weight_data = weight.to_vec()?;

        // Perform transposed convolution
        for n in 0..batch_size {
            for g in 0..groups {
                let in_ch_start = g * (in_channels / groups);
                let in_ch_end = (g + 1) * (in_channels / groups);
                let out_ch_start = g * (out_channels / groups);
                let out_ch_end = (g + 1) * (out_channels / groups);

                for ic in in_ch_start..in_ch_end {
                    for ih in 0..input_height {
                        for iw in 0..input_width {
                            let input_val = self_data[n * in_channels * input_height * input_width
                                + ic * input_height * input_width
                                + ih * input_width
                                + iw];

                            for oc in out_ch_start..out_ch_end {
                                let oc_rel = oc - out_ch_start;
                                for kh in 0..kernel_height {
                                    for kw in 0..kernel_width {
                                        let oh = ih * stride.0 + kh * dilation.0;
                                        let ow = iw * stride.1 + kw * dilation.1;

                                        if oh >= padding.0 && ow >= padding.1 {
                                            let oh_final = oh - padding.0;
                                            let ow_final = ow - padding.1;

                                            if oh_final < output_height && ow_final < output_width {
                                                let weight_idx = ic
                                                    * (out_channels / groups)
                                                    * kernel_height
                                                    * kernel_width
                                                    + oc_rel * kernel_height * kernel_width
                                                    + kh * kernel_width
                                                    + kw;
                                                let output_idx =
                                                    n * out_channels * output_height * output_width
                                                        + oc * output_height * output_width
                                                        + oh_final * output_width
                                                        + ow_final;

                                                output_data[output_idx] = output_data[output_idx]
                                                    + input_val * weight_data[weight_idx];
                                            }
                                        }
                                    }
                                }
                            }
                        }
                    }
                }
            }
        }

        // Create output tensor
        let mut output = Tensor::from_data(
            output_data,
            vec![batch_size, out_channels, output_height, output_width],
            self.device(),
        )?;

        // Add bias if provided
        if let Some(b) = bias {
            if b.shape().dims() != [out_channels] {
                return Err(TorshError::InvalidArgument(format!(
                    "Bias must have shape [{}], got {:?}",
                    out_channels,
                    b.shape().dims()
                )));
            }

            // Broadcast the per-output-channel bias across the (H, W) spatial map
            // of each channel, for every batch element.
            let bias_data = b.to_vec()?;
            let mut output_data = output.to_vec()?;
            Self::add_channel_bias(
                &mut output_data,
                &bias_data,
                out_channels,
                output_height * output_width,
            );

            // Create new output tensor with bias added
            output = Tensor::from_data(
                output_data,
                vec![batch_size, out_channels, output_height, output_width],
                self.device(),
            )?;
        }

        // Track operation for autograd
        if crate::should_record_grad(
            self.requires_grad
                || weight.requires_grad
                || (bias.is_some() && bias.expect("bias checked with is_some").requires_grad),
        ) {
            use std::sync::Arc;
            output.requires_grad = true;
            output.operation = crate::Operation::Custom(
                "conv_transpose2d".to_string(),
                vec![
                    Arc::downgrade(&Arc::new(self.clone())),
                    Arc::downgrade(&Arc::new(weight.clone())),
                ],
            );
        }

        Ok(output)
    }

    /// 1D cross-correlation operation
    /// Computes the cross-correlation between two 1D signals
    #[allow(clippy::needless_range_loop)]
    pub fn xcorr1d(&self, other: &Self, mode: CorrelationMode) -> Result<Self> {
        let self_shape_ref = self.shape();
        let other_shape_ref = other.shape();
        let self_shape = self_shape_ref.dims();
        let other_shape = other_shape_ref.dims();

        if self_shape.len() != 1 || other_shape.len() != 1 {
            return Err(TorshError::InvalidArgument(
                "xcorr1d requires 1D tensors".to_string(),
            ));
        }

        let n = self_shape[0];
        let m = other_shape[0];

        let (output_size, lag_start) = match mode {
            CorrelationMode::Full => (n + m - 1, -(m as i32 - 1)),
            CorrelationMode::Valid => {
                if n < m || m < n {
                    return Err(TorshError::InvalidArgument(
                        "Valid mode requires both tensors to have the same size or one to be smaller".to_string(),
                    ));
                }
                (std::cmp::max(n, m) - std::cmp::min(n, m) + 1, 0)
            }
            CorrelationMode::Same => (n, -((m as i32 - 1) / 2)),
        };

        let mut output_data = vec![<T as TensorElement>::zero(); output_size];
        let self_data = self.to_vec()?;
        let other_data = other.to_vec()?;

        // Compute cross-correlation: (f ★ g)[lag] = Σ_i f[i] * g[i - lag]
        for i in 0..output_size {
            let mut sum = <T as TensorElement>::zero();
            let lag = lag_start + i as i32;

            for j in 0..n {
                let other_idx = j as i32 - lag;
                if other_idx >= 0 && (other_idx as usize) < m {
                    sum = sum + self_data[j] * other_data[other_idx as usize];
                }
            }
            output_data[i] = sum;
        }

        let output = Tensor::from_data(output_data, vec![output_size], self.device())?;

        Ok(output)
    }

    /// 1D auto-correlation operation
    /// Computes the auto-correlation of a 1D signal
    pub fn autocorr1d(&self, max_lag: Option<usize>) -> Result<Self> {
        let shape_ref = self.shape();
        let shape = shape_ref.dims();
        if shape.len() != 1 {
            return Err(TorshError::InvalidArgument(
                "autocorr1d requires 1D tensor".to_string(),
            ));
        }

        let n = shape[0];
        let max_lag = max_lag.unwrap_or(n - 1).min(n - 1);

        let self_data = self.to_vec()?;
        let mut output_data = Vec::with_capacity(max_lag + 1);

        // Directly compute auto-correlation: R[k] = Σ_n x[n] * x[n-k]
        for lag in 0..=max_lag {
            let mut sum = <T as TensorElement>::zero();

            for i in lag..n {
                sum = sum + self_data[i] * self_data[i - lag];
            }

            output_data.push(sum);
        }

        let output = Tensor::from_data(output_data, vec![max_lag + 1], self.device())?;
        Ok(output)
    }

    /// 2D cross-correlation operation
    /// Computes the 2D cross-correlation between two signals
    pub fn xcorr2d(&self, other: &Self, mode: CorrelationMode) -> Result<Self> {
        let self_shape_ref = self.shape();
        let other_shape_ref = other.shape();
        let self_shape = self_shape_ref.dims();
        let other_shape = other_shape_ref.dims();

        if self_shape.len() != 2 || other_shape.len() != 2 {
            return Err(TorshError::InvalidArgument(
                "xcorr2d requires 2D tensors".to_string(),
            ));
        }

        let (h1, w1) = (self_shape[0], self_shape[1]);
        let (h2, w2) = (other_shape[0], other_shape[1]);

        let (out_h, out_w, start_h, start_w) = match mode {
            CorrelationMode::Full => (h1 + h2 - 1, w1 + w2 - 1, 0, 0),
            CorrelationMode::Valid => {
                if h1 < h2 || w1 < w2 {
                    return Err(TorshError::InvalidArgument(
                        "Valid mode requires first tensor to be larger than or equal to second"
                            .to_string(),
                    ));
                }
                (h1 - h2 + 1, w1 - w2 + 1, h2 - 1, w2 - 1)
            }
            CorrelationMode::Same => (h1, w1, (h2 - 1) / 2, (w2 - 1) / 2),
        };

        let mut output_data = vec![<T as TensorElement>::zero(); out_h * out_w];
        let self_data = self.to_vec()?;
        let other_data = other.to_vec()?;

        // Compute 2D cross-correlation
        for i in 0..out_h {
            for j in 0..out_w {
                let mut sum = <T as TensorElement>::zero();
                let actual_i = i + start_h;
                let actual_j = j + start_w;

                for ki in 0..h2 {
                    for kj in 0..w2 {
                        let src_i = actual_i as i32 - ki as i32;
                        let src_j = actual_j as i32 - kj as i32;

                        if src_i >= 0
                            && (src_i as usize) < h1
                            && src_j >= 0
                            && (src_j as usize) < w1
                        {
                            let self_idx = src_i as usize * w1 + src_j as usize;
                            let other_idx = ki * w2 + kj;
                            sum = sum + self_data[self_idx] * other_data[other_idx];
                        }
                    }
                }
                output_data[i * out_w + j] = sum;
            }
        }

        let output = Tensor::from_data(output_data, vec![out_h, out_w], self.device())?;
        Ok(output)
    }

    /// 1D median filter
    /// Applies a median filter with the specified window size
    pub fn median_filter1d(&self, window_size: usize) -> Result<Self> {
        let shape_ref = self.shape();
        let shape = shape_ref.dims();
        if shape.len() != 1 {
            return Err(TorshError::InvalidArgument(
                "median_filter1d requires 1D tensor".to_string(),
            ));
        }

        if window_size == 0 || window_size % 2 == 0 {
            return Err(TorshError::InvalidArgument(
                "Window size must be odd and greater than 0".to_string(),
            ));
        }

        let n = shape[0];
        let half_window = window_size / 2;
        let mut output_data = Vec::with_capacity(n);
        let self_data = self.to_vec()?;

        for i in 0..n {
            let mut window_values = Vec::new();

            // Collect values in the window (with padding by repeating edge values)
            for j in 0..window_size {
                let idx = i as i32 + j as i32 - half_window as i32;
                let actual_idx = if idx < 0 {
                    0
                } else if idx >= n as i32 {
                    n - 1
                } else {
                    idx as usize
                };
                window_values.push(self_data[actual_idx]);
            }

            // Sort to find median
            window_values.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
            output_data.push(window_values[half_window]);
        }

        let output = Tensor::from_data(output_data, vec![n], self.device())?;
        Ok(output)
    }

    /// 2D median filter
    /// Applies a 2D median filter with the specified window size
    pub fn median_filter2d(&self, window_size: (usize, usize)) -> Result<Self> {
        let shape_ref = self.shape();
        let shape = shape_ref.dims();
        if shape.len() != 2 {
            return Err(TorshError::InvalidArgument(
                "median_filter2d requires 2D tensor".to_string(),
            ));
        }

        let (window_h, window_w) = window_size;
        if window_h == 0 || window_w == 0 || window_h % 2 == 0 || window_w % 2 == 0 {
            return Err(TorshError::InvalidArgument(
                "Window dimensions must be odd and greater than 0".to_string(),
            ));
        }

        let (h, w) = (shape[0], shape[1]);
        let half_h = window_h / 2;
        let half_w = window_w / 2;
        let mut output_data = Vec::with_capacity(h * w);
        let self_data = self.to_vec()?;

        for i in 0..h {
            for j in 0..w {
                let mut window_values = Vec::new();

                // Collect values in the 2D window
                for di in 0..window_h {
                    for dj in 0..window_w {
                        let row = i as i32 + di as i32 - half_h as i32;
                        let col = j as i32 + dj as i32 - half_w as i32;

                        // Handle boundaries by clamping
                        let actual_row = row.max(0).min(h as i32 - 1) as usize;
                        let actual_col = col.max(0).min(w as i32 - 1) as usize;

                        window_values.push(self_data[actual_row * w + actual_col]);
                    }
                }

                // Sort to find median
                window_values.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
                output_data.push(window_values[window_values.len() / 2]);
            }
        }

        let output = Tensor::from_data(output_data, vec![h, w], self.device())?;
        Ok(output)
    }

    /// 1D Gaussian filter
    /// Applies a Gaussian filter with specified sigma (standard deviation)
    pub fn gaussian_filter1d(&self, sigma: f32, kernel_size: Option<usize>) -> Result<Self> {
        let tensor_shape = self.shape();
        let shape = tensor_shape.dims();
        if shape.len() != 1 {
            return Err(TorshError::InvalidArgument(
                "gaussian_filter1d requires 1D tensor".to_string(),
            ));
        }

        if sigma <= 0.0 {
            return Err(TorshError::InvalidArgument(
                "Sigma must be positive".to_string(),
            ));
        }

        // Calculate kernel size if not provided (6 sigma rule)
        let kernel_size = kernel_size.unwrap_or(((6.0 * sigma) as usize).max(3));
        let kernel_size = if kernel_size % 2 == 0 {
            kernel_size + 1
        } else {
            kernel_size
        };

        // Generate Gaussian kernel
        let half_size = kernel_size / 2;
        let mut kernel = Vec::with_capacity(kernel_size);
        let mut sum = 0.0f32;

        for i in 0..kernel_size {
            let x = i as f32 - half_size as f32;
            let value = (-0.5 * (x / sigma).powi(2)).exp();
            kernel.push(value);
            sum += value;
        }

        // Normalize kernel
        for value in &mut kernel {
            *value /= sum;
        }

        // Create kernel tensor
        let kernel_data: Vec<T> = kernel
            .into_iter()
            .map(|v| {
                T::from(v as f64)
                    .unwrap_or_else(|| T::from(0.0).expect("numeric conversion should succeed"))
            })
            .collect();
        let kernel_tensor = Tensor::from_data(kernel_data, vec![kernel_size], self.device())?;

        // Apply convolution (which is equivalent to correlation for symmetric kernels)
        self.xcorr1d(&kernel_tensor, CorrelationMode::Same)
    }

    /// 2D Gaussian filter
    /// Applies a 2D Gaussian filter with specified sigma values
    pub fn gaussian_filter2d(
        &self,
        sigma: (f32, f32),
        kernel_size: Option<(usize, usize)>,
    ) -> Result<Self> {
        let tensor_shape = self.shape();
        let shape = tensor_shape.dims();
        if shape.len() != 2 {
            return Err(TorshError::InvalidArgument(
                "gaussian_filter2d requires 2D tensor".to_string(),
            ));
        }

        let (sigma_x, sigma_y) = sigma;
        if sigma_x <= 0.0 || sigma_y <= 0.0 {
            return Err(TorshError::InvalidArgument(
                "Sigma values must be positive".to_string(),
            ));
        }

        // Calculate kernel sizes if not provided
        let (kernel_h, kernel_w) = kernel_size.unwrap_or((
            ((6.0 * sigma_y) as usize).max(3),
            ((6.0 * sigma_x) as usize).max(3),
        ));
        let kernel_h = if kernel_h % 2 == 0 {
            kernel_h + 1
        } else {
            kernel_h
        };
        let kernel_w = if kernel_w % 2 == 0 {
            kernel_w + 1
        } else {
            kernel_w
        };

        // Generate 2D Gaussian kernel
        let half_h = kernel_h / 2;
        let half_w = kernel_w / 2;
        let mut kernel = Vec::with_capacity(kernel_h * kernel_w);
        let mut sum = 0.0f32;

        for i in 0..kernel_h {
            for j in 0..kernel_w {
                let y = i as f32 - half_h as f32;
                let x = j as f32 - half_w as f32;
                let value = (-0.5 * ((x / sigma_x).powi(2) + (y / sigma_y).powi(2))).exp();
                kernel.push(value);
                sum += value;
            }
        }

        // Normalize kernel
        for value in &mut kernel {
            *value /= sum;
        }

        // Create kernel tensor
        let kernel_data: Vec<T> = kernel
            .into_iter()
            .map(|v| {
                T::from(v as f64)
                    .unwrap_or_else(|| T::from(0.0).expect("numeric conversion should succeed"))
            })
            .collect();
        let kernel_tensor =
            Tensor::from_data(kernel_data, vec![kernel_h, kernel_w], self.device())?;

        // Apply 2D correlation
        self.xcorr2d(&kernel_tensor, CorrelationMode::Same)
    }
}

/// Correlation modes for signal processing operations
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum CorrelationMode {
    /// Full correlation output
    Full,
    /// Valid correlation output (no padding)
    Valid,
    /// Same size as input (with padding)
    Same,
}