burn-backend 0.22.0

#![allow(clippy::single_range_in_vec_init)]
use super::{ConvOptions, ConvTransposeOptions, PadMode};
use crate::{Backend, TensorMetadata, tensor::FloatTensor};
use burn_std::{MetadataError, Shape, Slice};

use alloc::{vec, vec::Vec};
#[cfg(not(feature = "std"))]
#[allow(unused_imports)]
use num_traits::Float as _;

/// Explicitly pads an input when convolution padding is asymmetric.
///
/// Backends whose native convolution operation only supports symmetric padding can use this
/// helper to apply the per-side padding to the input and reset the returned convolution options
/// to zero padding. Symmetric options and their input are returned unchanged.
pub fn pad_asymmetric_conv_input<B: Backend, const N: usize>(
    x: FloatTensor<B>,
    options: ConvOptions<N>,
) -> (FloatTensor<B>, ConvOptions<N>) {
    if !options.is_asymmetric() {
        return (x, options);
    }

    let mut padding = vec![(0, 0); N + 2];
    padding[2..].copy_from_slice(&options.padding);
    let x = B::float_pad(x, &padding, PadMode::Constant(0.0));
    let options = ConvOptions::new(options.stride, [0; N], options.dilation, options.groups);

    (x, options)
}

fn crop_asymmetric_input_grad<B: Backend, const N: usize>(
    grad: FloatTensor<B>,
    input_shape: &Shape,
    padding: &[(usize, usize); N],
) -> FloatTensor<B> {
    let mut ranges = input_shape
        .iter()
        .map(|&size| Slice::from(0..size))
        .collect::<Vec<_>>();
    for (dim, (begin, _)) in padding.iter().enumerate() {
        ranges[dim + 2] = Slice::from(*begin..*begin + input_shape[dim + 2]);
    }
    B::float_slice(grad, &ranges)
}

/// Calculate the expected output shape `[batch_size, channels_out, spatial_dims, ..]` for a pooling operation.
pub fn calculate_pool_output_shape<const N: usize>(
    in_shape: &Shape,
    kernel_size: &[usize; N],
    stride: &[usize; N],
    padding: &[usize; N],
    dilation: &[usize; N],
    ceil_mode: bool,
) -> Result<Shape, MetadataError> {
    if in_shape.rank() != N + 2 {
        return Err(MetadataError::RankMismatch {
            left: in_shape.rank(),
            right: N + 2,
        });
    }

    let mut out_shape = in_shape.clone();
    // Spatial dims
    for (i, size_i) in out_shape[2..].iter_mut().enumerate() {
        *size_i = calculate_pool_output_size(
            kernel_size[i],
            stride[i],
            padding[i],
            dilation[i],
            *size_i,
            ceil_mode,
        );
    }

    Ok(out_shape)
}

/// Calculate the expected output shape `[batch_size, channels_out, spatial_dims, ..]` for a convolution.
pub fn calculate_conv_output_shape<const N: usize>(
    in_shape: &Shape,
    weight_shape: &Shape,
    stride: &[usize; N],
    padding: &[(usize, usize); N],
    dilation: &[usize; N],
) -> Result<Shape, MetadataError> {
    if weight_shape.rank() != N + 2 {
        return Err(MetadataError::RankMismatch {
            left: weight_shape.rank(),
            right: N + 2,
        });
    }

    if in_shape.rank() != N + 2 {
        return Err(MetadataError::RankMismatch {
            left: in_shape.rank(),
            right: N + 2,
        });
    }

    let kernel_size = &weight_shape[2..];

    let mut out_shape = in_shape.clone();
    // Spatial dims
    for (i, size_i) in out_shape[2..].iter_mut().enumerate() {
        *size_i = calculate_conv_output_size_asymmetric(
            kernel_size[i],
            stride[i],
            padding[i],
            dilation[i],
            *size_i,
        );
    }
    // Output channels
    out_shape[1] = weight_shape[0];

    Ok(out_shape)
}

/// Calculate the expected output shape `[batch_size, channels_out, spatial_dims, ..]` for a transposed convolution.
pub fn calculate_conv_transpose_output_shape<const N: usize>(
    in_shape: &Shape,
    weight_shape: &Shape,
    stride: &[usize; N],
    padding: &[usize; N],
    padding_out: &[usize; N],
    dilation: &[usize; N],
    groups: usize,
) -> Result<Shape, MetadataError> {
    if weight_shape.rank() != N + 2 {
        return Err(MetadataError::RankMismatch {
            left: weight_shape.rank(),
            right: N + 2,
        });
    }

    if in_shape.rank() != N + 2 {
        return Err(MetadataError::RankMismatch {
            left: in_shape.rank(),
            right: N + 2,
        });
    }

    let kernel_size = &weight_shape[2..];

    let mut out_shape = in_shape.clone();
    // Spatial dims
    for (i, size_i) in out_shape[2..].iter_mut().enumerate() {
        *size_i = calculate_conv_transpose_output_size(
            kernel_size[i],
            stride[i],
            padding[i],
            padding_out[i],
            dilation[i],
            *size_i,
        );
    }
    // Output channels
    out_shape[1] = weight_shape[1] * groups;

    Ok(out_shape)
}

/// Calculate the expected padding size required when applying a convolution.
pub fn calculate_conv_padding(
    kernel_size: usize,
    stride: usize,
    size_in: usize,
    size_out: usize,
) -> usize {
    let kernel_size = kernel_size as f32;
    let stride = stride as f32;
    let size_in = size_in as f32;
    let size_out = size_out as f32;

    let padding = stride * (size_out - 1.) - size_in + kernel_size;
    let padding = (padding / 2.).ceil();

    padding as usize
}

/// Calculate the expected output size when doing a convolution operation.
pub fn calculate_conv_output_size(
    kernel_size: usize,
    stride: usize,
    padding: usize,
    dilation: usize,
    size_in: usize,
) -> usize {
    calculate_conv_output_size_asymmetric(
        kernel_size,
        stride,
        (padding, padding),
        dilation,
        size_in,
    )
}

/// Calculate the expected output size with explicit `(begin, end)` padding.
pub fn calculate_conv_output_size_asymmetric(
    kernel_size: usize,
    stride: usize,
    padding: (usize, usize),
    dilation: usize,
    size_in: usize,
) -> usize {
    (size_in + padding.0 + padding.1 - dilation * (kernel_size - 1) - 1) / stride + 1
}

/// Calculate the expected output sizes when doing a convolution operation.
pub fn calculate_conv_output_sizes(
    kernel_size: &[usize],
    stride: &[usize],
    padding: &[(usize, usize)],
    dilation: &[usize],
    size_in: &[usize],
) -> Vec<usize> {
    size_in
        .iter()
        .enumerate()
        .map(|(i, size_in)| {
            calculate_conv_output_size_asymmetric(
                kernel_size[i],
                stride[i],
                padding[i],
                dilation[i],
                *size_in,
            )
        })
        .collect()
}

/// Calculate the expected output size when doing a pooling operation.
///
/// # Arguments
///
/// * `kernel_size` - Size of the pooling kernel
/// * `stride` - Stride of the pooling operation
/// * `padding` - Padding applied to input
/// * `dilation` - Dilation of the pooling kernel
/// * `size_in` - Input size (height or width)
/// * `ceil_mode` - If true, use ceiling instead of floor for output size calculation.
///   This allows the last pooling window to go out-of-bounds if needed, as long as it
///   starts inside the input or left padding (matches PyTorch and ONNX).
pub fn calculate_pool_output_size(
    kernel_size: usize,
    stride: usize,
    padding: usize,
    dilation: usize,
    size_in: usize,
    ceil_mode: bool,
) -> usize {
    let size_padded = size_in + 2 * padding;
    let kernel_extent = dilation * (kernel_size - 1) + 1;
    if ceil_mode {
        // Ceiling division: (a + b - 1) / b. Adding `stride - 1` before subtracting
        // allows a kernel up to `stride - 1` larger than the padded input, where one
        // window still fits (PyTorch rejects anything larger).
        let size_out = (size_padded + stride - 1 - kernel_extent) / stride + 1;
        // Drop the last window if it would start at or past the end of the input
        // (in the trailing padding, or beyond the input without padding)
        if (size_out - 1) * stride >= size_in + padding {
            size_out - 1
        } else {
            size_out
        }
    } else {
        // Floor division (default)
        (size_padded - kernel_extent) / stride + 1
    }
}

/// Calculate the expected output size when doing a transposed convolution operation.
pub fn calculate_conv_transpose_output_size(
    kernel_size: usize,
    stride: usize,
    padding: usize,
    padding_out: usize,
    dilation: usize,
    size_in: usize,
) -> usize {
    (size_in - 1) * stride + (dilation * (kernel_size - 1) + 1) + padding_out - 2 * padding
}

/// Calculate the original input size that was used for a transposed convolution.
/// This is used during the backward pass to recover the correct gradient shape.
fn calculate_conv_transpose_input_size(
    kernel_size: usize,
    stride: usize,
    padding: usize,
    padding_out: usize,
    dilation: usize,
    size_out: usize,
) -> usize {
    // We solve the forward formula for size_in:
    // size_out = (size_in - 1) * stride + (dilation * (kernel_size - 1) + 1) + padding_out - 2 * padding
    (size_out + 2 * padding - dilation * (kernel_size - 1) - padding_out - 1) / stride + 1
}

/// Calculate the original input sizes that were used for a transposed convolution.
fn calculate_conv_transpose_input_sizes<const D: usize>(
    kernel_size: [usize; D],
    stride: [usize; D],
    padding: [usize; D],
    padding_out: [usize; D],
    dilation: [usize; D],
    size_out: [usize; D],
) -> [usize; D] {
    let mut res = [0; D];
    for i in 0..D {
        res[i] = calculate_conv_transpose_input_size(
            kernel_size[i],
            stride[i],
            padding[i],
            padding_out[i],
            dilation[i],
            size_out[i],
        );
    }
    res
}

/// Calculate the [1D convolution](crate::ops::ModuleOps::conv1d) backward pass, returning the gradient for `x`.
pub(crate) fn conv1d_x_backward<B: Backend>(
    x: FloatTensor<B>,
    weight: FloatTensor<B>,
    output_grad: FloatTensor<B>,
    options: ConvOptions<1>,
) -> FloatTensor<B> {
    if options.is_asymmetric() {
        let input_shape = x.shape();
        let padding = options.padding;
        let (x, options) = pad_asymmetric_conv_input::<B, 1>(x, options);
        let grad = conv1d_x_backward_symmetric::<B>(x, weight, output_grad, options);
        return crop_asymmetric_input_grad::<B, 1>(grad, &input_shape, &padding);
    }

    conv1d_x_backward_symmetric::<B>(x, weight, output_grad, options)
}

fn conv1d_x_backward_symmetric<B: Backend>(
    x: FloatTensor<B>,
    weight: FloatTensor<B>,
    output_grad: FloatTensor<B>,
    options: ConvOptions<1>,
) -> FloatTensor<B> {
    let weight_shape = weight.shape();

    let [_batch_size, _, length_in] = x.shape().dims();
    let [_batch_size, _channels_out, length_out] = output_grad.shape().dims();
    let [_, _, kernel_size] = weight_shape.dims();

    let padding_out = calculate_padding_out(
        kernel_size,
        options.stride[0],
        options.padding_begin()[0],
        options.dilation[0],
        length_in,
        length_out,
    );

    B::conv_transpose1d(
        output_grad,
        weight,
        None,
        ConvTransposeOptions::new(
            options.stride,
            options.padding_begin(),
            [padding_out],
            options.dilation,
            options.groups,
        ),
    )
}

/// Calculate the [1D convolution](crate::ops::ModuleOps::conv1d) backward pass, returning the gradient for `weight`.
pub(crate) fn conv1d_weight_backward<B: Backend>(
    x: FloatTensor<B>,
    weight: FloatTensor<B>,
    output_grad: FloatTensor<B>,
    options: ConvOptions<1>,
) -> FloatTensor<B> {
    if options.is_asymmetric() {
        let (x, options) = pad_asymmetric_conv_input::<B, 1>(x, options);
        return conv1d_weight_backward_symmetric::<B>(x, weight, output_grad, options);
    }

    conv1d_weight_backward_symmetric::<B>(x, weight, output_grad, options)
}

fn conv1d_weight_backward_symmetric<B: Backend>(
    x: FloatTensor<B>,
    weight: FloatTensor<B>,
    output_grad: FloatTensor<B>,
    options: ConvOptions<1>,
) -> FloatTensor<B> {
    let weight_dtype = weight.dtype();
    let weight_shape = weight.shape();
    let weight_device = weight.device();

    match options.groups == 1 {
        true => conv1d_weight_grad_no_groups::<B>(x, output_grad, weight_shape, options),
        false => conv1d_weight_grad_groups::<B>(
            x,
            B::float_zeros(weight_shape, &weight_device, weight_dtype.into()),
            output_grad,
            options,
        ),
    }
}

/// Calculate the [1D convolution](crate::ops::ModuleOps::conv1d) backward pass, returning the gradient for `bias`.
pub(crate) fn conv1d_bias_backward<B: Backend>(
    x: FloatTensor<B>,
    bias: FloatTensor<B>,
    output_grad: FloatTensor<B>,
) -> FloatTensor<B> {
    let [batch_size, _, _length_in] = x.shape().dims();
    let [_batch_size, channels_out, length_out] = output_grad.shape().dims();

    let grad = B::float_swap_dims(output_grad, 0, 1);
    let grad = B::float_reshape(grad, Shape::new([channels_out, batch_size * length_out]));
    let grad = B::float_sum_dim(grad, 1);

    B::float_reshape(grad, bias.shape())
}

/// Calculate the [2D convolution](crate::ops::ModuleOps::conv2d) backward pass, returning the gradient for `x`.
pub(crate) fn conv2d_x_backward<B: Backend>(
    x: FloatTensor<B>,
    weight: FloatTensor<B>,
    output_grad: FloatTensor<B>,
    options: ConvOptions<2>,
) -> FloatTensor<B> {
    if options.is_asymmetric() {
        let input_shape = x.shape();
        let padding = options.padding;
        let (x, options) = pad_asymmetric_conv_input::<B, 2>(x, options);
        let grad = conv2d_x_backward_symmetric::<B>(x, weight, output_grad, options);
        return crop_asymmetric_input_grad::<B, 2>(grad, &input_shape, &padding);
    }

    conv2d_x_backward_symmetric::<B>(x, weight, output_grad, options)
}

fn conv2d_x_backward_symmetric<B: Backend>(
    x: FloatTensor<B>,
    weight: FloatTensor<B>,
    output_grad: FloatTensor<B>,
    options: ConvOptions<2>,
) -> FloatTensor<B> {
    let weight_shape = weight.shape();

    let [_batch_size, _channels_in, height_in, width_in] = x.shape().dims();
    let [_, _, height_out, width_out] = output_grad.shape().dims();
    let [_channels_out, _, kernel_size_1, kernel_size_2] = weight_shape.dims();

    let padding_1_out = calculate_padding_out(
        kernel_size_1,
        options.stride[0],
        options.padding_begin()[0],
        options.dilation[0],
        height_in,
        height_out,
    );
    let padding_2_out = calculate_padding_out(
        kernel_size_2,
        options.stride[1],
        options.padding_begin()[1],
        options.dilation[1],
        width_in,
        width_out,
    );

    B::conv_transpose2d(
        output_grad,
        weight,
        None,
        ConvTransposeOptions::new(
            options.stride,
            options.padding_begin(),
            [padding_1_out, padding_2_out],
            options.dilation,
            options.groups,
        ),
    )
}

/// Calculate the [2D convolution](crate::ops::ModuleOps::conv2d) backward pass, returning the gradient for `weight`.
pub(crate) fn conv2d_weight_backward<B: Backend>(
    x: FloatTensor<B>,
    weight: FloatTensor<B>,
    output_grad: FloatTensor<B>,
    options: ConvOptions<2>,
) -> FloatTensor<B> {
    if options.is_asymmetric() {
        let (x, options) = pad_asymmetric_conv_input::<B, 2>(x, options);
        return conv2d_weight_backward_symmetric::<B>(x, weight, output_grad, options);
    }

    conv2d_weight_backward_symmetric::<B>(x, weight, output_grad, options)
}

fn conv2d_weight_backward_symmetric<B: Backend>(
    x: FloatTensor<B>,
    weight: FloatTensor<B>,
    output_grad: FloatTensor<B>,
    options: ConvOptions<2>,
) -> FloatTensor<B> {
    let weight_dtype = weight.dtype();
    let weight_shape = weight.shape();
    let weight_device = weight.device();

    match options.groups == 1 {
        true => conv2d_weight_grad_no_groups::<B>(x, output_grad, weight_shape, options),
        false => conv2d_weight_grad_groups::<B>(
            x,
            B::float_zeros(weight_shape, &weight_device, weight_dtype.into()),
            output_grad,
            options,
        ),
    }
}

/// Calculate the [2D convolution](crate::ops::ModuleOps::conv2d) backward pass, returning the gradient for `bias`.
pub(crate) fn conv2d_bias_backward<B: Backend>(
    x: FloatTensor<B>,
    bias: FloatTensor<B>,
    output_grad: FloatTensor<B>,
) -> FloatTensor<B> {
    let [batch_size, _, _, _] = x.shape().dims();
    let [_, channels_out, height_out, width_out] = output_grad.shape().dims();

    let grad = B::float_swap_dims(output_grad, 0, 1);
    let grad = B::float_reshape(
        grad,
        Shape::new([channels_out, batch_size * height_out * width_out]),
    );
    let grad = B::float_sum_dim(grad, 1);

    B::float_reshape(grad, bias.shape())
}

/// Calculate the [3D convolution](crate::ops::ModuleOps::conv3d) backward pass, returning the gradient for `x`.
pub(crate) fn conv3d_x_backward<B: Backend>(
    x: FloatTensor<B>,
    weight: FloatTensor<B>,
    output_grad: FloatTensor<B>,
    options: ConvOptions<3>,
) -> FloatTensor<B> {
    let weight_shape = weight.shape();

    let [_batch_size, _channels_in, depth_in, height_in, width_in] = x.shape().dims();
    let [_, _, depth_out, height_out, width_out] = output_grad.shape().dims();
    let [
        _channels_out,
        _,
        kernel_size_1,
        kernel_size_2,
        kernel_size_3,
    ] = weight_shape.dims();

    let padding_1_out = calculate_padding_out(
        kernel_size_1,
        options.stride[0],
        options.padding_begin()[0],
        options.dilation[0],
        depth_in,
        depth_out,
    );
    let padding_2_out = calculate_padding_out(
        kernel_size_2,
        options.stride[1],
        options.padding_begin()[1],
        options.dilation[1],
        height_in,
        height_out,
    );
    let padding_3_out = calculate_padding_out(
        kernel_size_3,
        options.stride[2],
        options.padding_begin()[2],
        options.dilation[2],
        width_in,
        width_out,
    );

    B::conv_transpose3d(
        output_grad,
        weight,
        None,
        ConvTransposeOptions::new(
            options.stride,
            options.padding_begin(),
            [padding_1_out, padding_2_out, padding_3_out],
            options.dilation,
            options.groups,
        ),
    )
}

/// Calculate the [3D convolution](crate::ops::ModuleOps::conv3d) backward pass, returning the gradient for `weight`.
pub(crate) fn conv3d_weight_backward<B: Backend>(
    x: FloatTensor<B>,
    weight: FloatTensor<B>,
    output_grad: FloatTensor<B>,
    options: ConvOptions<3>,
) -> FloatTensor<B> {
    let weight_dtype = weight.dtype();
    let weight_shape = weight.shape();
    let weight_device = weight.device();

    match options.groups == 1 {
        true => conv3d_weight_grad_no_groups::<B>(x, output_grad, weight_shape, options),
        false => conv3d_weight_grad_groups::<B>(
            x,
            B::float_zeros(weight_shape, &weight_device, weight_dtype.into()),
            output_grad,
            options,
        ),
    }
}

/// Calculate the [3D convolution](crate::ops::ModuleOps::conv3d) backward pass, returning the gradient for `bias`.
pub(crate) fn conv3d_bias_backward<B: Backend>(
    x: FloatTensor<B>,
    bias: FloatTensor<B>,
    output_grad: FloatTensor<B>,
) -> FloatTensor<B> {
    let [batch_size, _channels_in, _depth_in, _height_in, _width_in] = x.shape().dims();
    let [_, channels_out, depth_out, height_out, width_out] = output_grad.shape().dims();

    let grad = B::float_swap_dims(output_grad, 0, 1);
    let grad = B::float_reshape(
        grad,
        Shape::new([
            channels_out,
            batch_size * depth_out * height_out * width_out,
        ]),
    );
    let grad = B::float_sum_dim(grad, 1);

    B::float_reshape(grad, bias.shape())
}

/// Calculate the [1D convolution transpose](crate::ops::ModuleOps::conv_transpose1d) backward pass, returning the gradient for `x`.
pub(crate) fn conv_transpose1d_x_backward<B: Backend>(
    weight: FloatTensor<B>,
    output_grad: FloatTensor<B>,
    options: ConvTransposeOptions<1>,
) -> FloatTensor<B> {
    let [batch_size, _c_out, out_length] = output_grad.shape().dims();
    let [c_in, _c_out_groups, kernel_size] = weight.shape().dims();

    let grad = B::conv1d(
        output_grad,
        weight,
        None,
        ConvOptions::new(
            options.stride,
            options.padding,
            options.dilation,
            options.groups,
        ),
    );

    if options.padding_out[0] == 0 {
        return grad;
    }

    let exp_length = calculate_conv_transpose_input_size(
        kernel_size,
        options.stride[0],
        options.padding[0],
        options.padding_out[0],
        options.dilation[0],
        out_length,
    );

    B::float_slice(
        grad,
        &[
            Slice::from(0..batch_size),
            Slice::from(0..c_in),
            Slice::from(0..exp_length),
        ],
    )
}

/// Calculate the [1D convolution transpose](crate::ops::ModuleOps::conv_transpose1d) backward pass, returning the gradient for `weight`.
pub(crate) fn conv_transpose1d_weight_backward<B: Backend>(
    x: FloatTensor<B>,
    weight: FloatTensor<B>,
    output_grad: FloatTensor<B>,
    options: ConvTransposeOptions<1>,
) -> FloatTensor<B> {
    let weight_dtype = weight.dtype();
    let weight_shape = weight.shape();
    let weight_device = weight.device();

    match options.groups == 1 {
        true => conv_transpose1d_weight_grad_no_groups::<B>(x, output_grad, weight_shape, options),
        false => conv_transpose1d_weight_grad_groups::<B>(
            x,
            B::float_zeros(weight_shape, &weight_device, weight_dtype.into()),
            output_grad,
            options,
        ),
    }
}

/// Calculate the [1D convolution transpose](crate::ops::ModuleOps::conv_transpose1d) backward pass, returning the gradient for `bias`.
pub(crate) fn conv_transpose1d_bias_backward<B: Backend>(
    x: FloatTensor<B>,
    bias: FloatTensor<B>,
    output_grad: FloatTensor<B>,
) -> FloatTensor<B> {
    let [batch_size, _channels_in, _] = x.shape().dims();
    let [_, channels_out, length_out] = output_grad.shape().dims();

    let grad = B::float_swap_dims(output_grad, 0, 1);
    let grad = B::float_reshape(grad, Shape::new([channels_out, batch_size * length_out]));
    let grad = B::float_sum_dim(grad, 1);

    B::float_reshape(grad, bias.shape())
}

/// Calculate the [2D convolution transpose](crate::ops::ModuleOps::conv_transpose2d) backward pass, returning the gradient for `x`.
pub(crate) fn conv_transpose2d_x_backward<B: Backend>(
    weight: FloatTensor<B>,
    output_grad: FloatTensor<B>,
    options: ConvTransposeOptions<2>,
) -> FloatTensor<B> {
    let [batch_size, _c_out, out_h, out_w] = output_grad.shape().dims();
    let [c_in, _c_out_groups, k_h, k_w] = weight.shape().dims();

    let grad = B::conv2d(
        output_grad,
        weight,
        None,
        ConvOptions::new(
            options.stride,
            options.padding,
            options.dilation,
            options.groups,
        ),
    );

    if options.padding_out[0] == 0 && options.padding_out[1] == 0 {
        return grad;
    }

    let [exp_h, exp_w] = calculate_conv_transpose_input_sizes(
        [k_h, k_w],
        options.stride,
        options.padding,
        options.padding_out,
        options.dilation,
        [out_h, out_w],
    );

    B::float_slice(
        grad,
        &[
            Slice::from(0..batch_size),
            Slice::from(0..c_in),
            Slice::from(0..exp_h),
            Slice::from(0..exp_w),
        ],
    )
}

/// Calculate the [2D convolution transpose](crate::ops::ModuleOps::conv_transpose2d) backward pass, returning the gradient for `weight`.
pub(crate) fn conv_transpose2d_weight_backward<B: Backend>(
    x: FloatTensor<B>,
    weight: FloatTensor<B>,
    output_grad: FloatTensor<B>,
    options: ConvTransposeOptions<2>,
) -> FloatTensor<B> {
    let weight_dtype = weight.dtype();
    let weight_shape = weight.shape();
    let weight_device = weight.device();

    match options.groups == 1 {
        true => conv_transpose2d_weight_grad_no_groups::<B>(x, output_grad, weight_shape, options),
        false => conv_transpose2d_weight_grad_groups::<B>(
            x,
            B::float_zeros(weight_shape, &weight_device, weight_dtype.into()),
            output_grad,
            options,
        ),
    }
}

/// Calculate the [2D convolution transpose](crate::ops::ModuleOps::conv_transpose2d) backward pass, returning the gradient for `bias`.
pub(crate) fn conv_transpose2d_bias_backward<B: Backend>(
    x: FloatTensor<B>,
    bias: FloatTensor<B>,
    output_grad: FloatTensor<B>,
) -> FloatTensor<B> {
    let [batch_size, _channels_in, _, _] = x.shape().dims();
    let [_, channels_out, height_out, width_out] = output_grad.shape().dims();

    let grad = B::float_swap_dims(output_grad, 0, 1);
    let grad = B::float_reshape(
        grad,
        Shape::new([channels_out, batch_size * height_out * width_out]),
    );
    let grad = B::float_sum_dim(grad, 1);

    B::float_reshape(grad, bias.shape())
}

/// Calculate the [3D convolution transpose](crate::ops::ModuleOps::conv_transpose3d) backward pass, returning the gradient for `x`.
pub(crate) fn conv_transpose3d_x_backward<B: Backend>(
    weight: FloatTensor<B>,
    output_grad: FloatTensor<B>,
    options: ConvTransposeOptions<3>,
) -> FloatTensor<B> {
    let [batch_size, _c_out, out_d, out_h, out_w] = output_grad.shape().dims();
    let [c_in, _c_out_groups, k_d, k_h, k_w] = weight.shape().dims();

    let grad = B::conv3d(
        output_grad,
        weight,
        None,
        ConvOptions::new(
            options.stride,
            options.padding,
            options.dilation,
            options.groups,
        ),
    );

    if options.padding_out[0] == 0 && options.padding_out[1] == 0 && options.padding_out[2] == 0 {
        return grad;
    }

    let [exp_d, exp_h, exp_w] = calculate_conv_transpose_input_sizes(
        [k_d, k_h, k_w],
        options.stride,
        options.padding,
        options.padding_out,
        options.dilation,
        [out_d, out_h, out_w],
    );

    B::float_slice(
        grad,
        &[
            Slice::from(0..batch_size),
            Slice::from(0..c_in),
            Slice::from(0..exp_d),
            Slice::from(0..exp_h),
            Slice::from(0..exp_w),
        ],
    )
}

/// Calculate the [3D convolution transpose](crate::ops::ModuleOps::conv_transpose3d) backward pass, returning the gradient for `weight`.
pub(crate) fn conv_transpose3d_weight_backward<B: Backend>(
    x: FloatTensor<B>,
    weight: FloatTensor<B>,
    output_grad: FloatTensor<B>,
    options: ConvTransposeOptions<3>,
) -> FloatTensor<B> {
    let weight_dtype = weight.dtype();
    let weight_shape = weight.shape();
    let weight_device = weight.device();

    match options.groups == 1 {
        true => conv_transpose3d_weight_grad_no_groups::<B>(x, output_grad, weight_shape, options),
        false => conv_transpose3d_weight_grad_groups::<B>(
            x,
            B::float_zeros(weight_shape, &weight_device, weight_dtype.into()),
            output_grad,
            options,
        ),
    }
}

/// Calculate the [3D convolution transpose](crate::ops::ModuleOps::conv_transpose3d) backward pass, returning the gradient for `bias`.
pub(crate) fn conv_transpose3d_bias_backward<B: Backend>(
    x: FloatTensor<B>,
    bias: FloatTensor<B>,
    output_grad: FloatTensor<B>,
) -> FloatTensor<B> {
    let [batch_size, _channels_in, _, _, _] = x.shape().dims();
    let [_, channels_out, depth_out, height_out, width_out] = output_grad.shape().dims();

    let grad = B::float_swap_dims(output_grad, 0, 1);
    let grad = B::float_reshape(
        grad,
        Shape::new([
            channels_out,
            batch_size * depth_out * height_out * width_out,
        ]),
    );
    let grad = B::float_sum_dim(grad, 1);

    B::float_reshape(grad, bias.shape())
}

/// Execute a 1D convolution using a 2D convolution.
pub(crate) fn conv1d_from_conv2d<B: Backend>(
    x: FloatTensor<B>,
    weight: FloatTensor<B>,
    bias: Option<FloatTensor<B>>,
    options: ConvOptions<1>,
) -> FloatTensor<B> {
    let (x, options) = pad_asymmetric_conv_input::<B, 1>(x, options);

    let [channels_out, _channels_in, kernel_size] = weight.shape().dims();
    let [batch_size, channels_in, length_in] = x.shape().dims();

    let weight = B::float_reshape(
        weight,
        Shape::new([channels_out, channels_in / options.groups, kernel_size, 1]),
    );
    let x = B::float_reshape(x, Shape::new([batch_size, channels_in, length_in, 1]));

    let tensor = B::conv2d(
        x,
        weight,
        bias,
        ConvOptions::new(
            [options.stride[0], 1],
            [options.padding_begin()[0], 0],
            [options.dilation[0], 1],
            options.groups,
        ),
    );
    let [batch_size, channels_out, height_out, _weight_out] = tensor.shape().dims();
    B::float_reshape(tensor, Shape::from([batch_size, channels_out, height_out]))
}

/// Execute a 1D transposed convolution using a 2D transposed convolution.
pub(crate) fn conv_transpose1d_from_conv_transpose2d<B: Backend>(
    x: FloatTensor<B>,
    weight: FloatTensor<B>,
    bias: Option<FloatTensor<B>>,
    options: ConvTransposeOptions<1>,
) -> FloatTensor<B> {
    let [channels_in, channels_out, kernel_size] = weight.shape().dims();
    let [batch_size, _channels_in, length_in] = x.shape().dims();

    let weight = B::float_reshape(
        weight,
        Shape::new([channels_in, channels_out, kernel_size, 1]),
    );
    let x = B::float_reshape(x, Shape::new([batch_size, channels_in, length_in, 1]));

    let tensor = B::conv_transpose2d(
        x,
        weight,
        bias,
        ConvTransposeOptions::new(
            [options.stride[0], 1],
            [options.padding[0], 0],
            [options.padding_out[0], 0],
            [options.dilation[0], 1],
            options.groups,
        ),
    );
    let [batch_size, channels_out, height_out, _weight_out] = tensor.shape().dims();
    B::float_reshape(tensor, Shape::from([batch_size, channels_out, height_out]))
}

fn conv1d_weight_grad_no_groups<B: Backend>(
    x: FloatTensor<B>,
    output_grad: FloatTensor<B>,
    weight_shape: Shape,
    options: ConvOptions<1>,
) -> FloatTensor<B> {
    let x_swapped = B::float_swap_dims(x, 0, 1);
    let output_grad_swapped = B::float_swap_dims(output_grad, 0, 1);
    let weight_grad_swapped = B::conv1d(
        x_swapped,
        output_grad_swapped,
        None,
        ConvOptions::new(options.dilation, options.padding_begin(), options.stride, 1),
    );
    let mut weight_grad = B::float_swap_dims(weight_grad_swapped, 0, 1);

    if weight_grad.shape() != weight_shape {
        let slices = vec![
            Slice::from(0..weight_shape[0]),
            Slice::from(0..weight_shape[1]),
            Slice::from(0..weight_shape[2]),
        ];
        weight_grad = B::float_slice(weight_grad, &slices);
    }
    weight_grad
}

fn conv2d_weight_grad_no_groups<B: Backend>(
    x: FloatTensor<B>,
    output_grad: FloatTensor<B>,
    weight_shape: Shape,
    options: ConvOptions<2>,
) -> FloatTensor<B> {
    let x_swapped = B::float_swap_dims(x, 0, 1);
    let output_grad_swapped = B::float_swap_dims(output_grad, 0, 1);
    let weight_grad_swapped = B::conv2d(
        x_swapped,
        output_grad_swapped,
        None,
        ConvOptions::new(options.dilation, options.padding_begin(), options.stride, 1),
    );
    let mut weight_grad = B::float_swap_dims(weight_grad_swapped, 0, 1);

    if weight_grad.shape() != weight_shape {
        let slices = vec![
            Slice::from(0..weight_shape[0]),
            Slice::from(0..weight_shape[1]),
            Slice::from(0..weight_shape[2]),
            Slice::from(0..weight_shape[3]),
        ];
        weight_grad = B::float_slice(weight_grad, &slices);
    }
    weight_grad
}

fn conv3d_weight_grad_no_groups<B: Backend>(
    x: FloatTensor<B>,
    output_grad: FloatTensor<B>,
    weight_shape: Shape,
    options: ConvOptions<3>,
) -> FloatTensor<B> {
    let x_swapped = B::float_swap_dims(x, 0, 1);
    let output_grad_swapped = B::float_swap_dims(output_grad, 0, 1);
    let weight_grad_swapped = B::conv3d(
        x_swapped,
        output_grad_swapped,
        None,
        ConvOptions::new(options.dilation, options.padding_begin(), options.stride, 1),
    );
    let mut weight_grad = B::float_swap_dims(weight_grad_swapped, 0, 1);

    if weight_grad.shape() != weight_shape {
        let slices = vec![
            Slice::from(0..weight_shape[0]),
            Slice::from(0..weight_shape[1]),
            Slice::from(0..weight_shape[2]),
            Slice::from(0..weight_shape[3]),
            Slice::from(0..weight_shape[4]),
        ];
        weight_grad = B::float_slice(weight_grad, &slices);
    }
    weight_grad
}

fn conv1d_weight_grad_groups<B: Backend>(
    x: FloatTensor<B>,
    mut weight_grad: FloatTensor<B>,
    output_grad: FloatTensor<B>,
    options: ConvOptions<1>,
) -> FloatTensor<B> {
    let [channels_out, increment_ci, kernel_size] = weight_grad.shape().dims();
    let increment_co = channels_out / options.groups;

    let x_swapped = B::float_swap_dims(x, 0, 1);
    let output_grad_swapped = B::float_swap_dims(output_grad, 0, 1);

    for g in 0..options.groups {
        let start_idx_ci = g * increment_ci;
        let end_idx_ci = (g + 1) * increment_ci;
        let start_idx_co = g * increment_co;
        let end_idx_co = (g + 1) * increment_co;

        let x_slice = vec![Slice::new(
            start_idx_ci as isize,
            Some(end_idx_ci as isize),
            1,
        )];
        let x = B::float_slice(x_swapped.clone(), &x_slice);
        let grad_slice = vec![Slice::new(
            start_idx_co as isize,
            Some(end_idx_co as isize),
            1,
        )];
        let grad = B::float_slice(output_grad_swapped.clone(), &grad_slice);
        let mut weight_grad_tmp = B::conv1d(
            x,
            grad,
            None,
            ConvOptions::new(options.dilation, options.padding_begin(), options.stride, 1),
        );
        weight_grad_tmp = B::float_swap_dims(weight_grad_tmp, 0, 1);
        weight_grad = B::float_slice_assign(
            weight_grad,
            &[
                Slice::from(start_idx_co..end_idx_co),
                Slice::from(0..increment_ci),
                Slice::from(0..kernel_size),
            ],
            weight_grad_tmp,
        );
    }

    weight_grad
}

fn conv2d_weight_grad_groups<B: Backend>(
    x: FloatTensor<B>,
    mut weight_grad: FloatTensor<B>,
    output_grad: FloatTensor<B>,
    options: ConvOptions<2>,
) -> FloatTensor<B> {
    let [channels_out, increment_ci, kernel_size_1, kernel_size_2] = weight_grad.shape().dims();
    let increment_co = channels_out / options.groups;

    let x_swapped = B::float_swap_dims(x, 0, 1);
    let output_grad_swapped = B::float_swap_dims(output_grad, 0, 1);

    for g in 0..options.groups {
        let start_idx_ci = g * increment_ci;
        let end_idx_ci = (g + 1) * increment_ci;
        let start_idx_co = g * increment_co;
        let end_idx_co = (g + 1) * increment_co;

        let x_slice = vec![Slice::new(
            start_idx_ci as isize,
            Some(end_idx_ci as isize),
            1,
        )];
        let x = B::float_slice(x_swapped.clone(), &x_slice);
        let grad_slice = vec![Slice::new(
            start_idx_co as isize,
            Some(end_idx_co as isize),
            1,
        )];
        let grad = B::float_slice(output_grad_swapped.clone(), &grad_slice);
        let mut weight_grad_tmp = B::conv2d(
            x,
            grad,
            None,
            ConvOptions::new(options.dilation, options.padding_begin(), options.stride, 1),
        );
        weight_grad_tmp = B::float_swap_dims(weight_grad_tmp, 0, 1);
        let [_, _, kernel_size_1_tmp, kernel_size_2_tmp] = weight_grad_tmp.shape().dims();

        if kernel_size_1_tmp != kernel_size_1 || kernel_size_2_tmp != kernel_size_2 {
            let slices = vec![
                Slice::from(0..increment_co),
                Slice::from(0..increment_ci),
                Slice::from(0..kernel_size_1),
                Slice::from(0..kernel_size_2),
            ];
            weight_grad_tmp = B::float_slice(weight_grad_tmp, &slices);
        }

        weight_grad = B::float_slice_assign(
            weight_grad,
            &[
                Slice::from(start_idx_co..end_idx_co),
                Slice::from(0..increment_ci),
                Slice::from(0..kernel_size_1),
                Slice::from(0..kernel_size_2),
            ],
            weight_grad_tmp,
        );
    }

    weight_grad
}

fn conv3d_weight_grad_groups<B: Backend>(
    x: FloatTensor<B>,
    mut weight_grad: FloatTensor<B>,
    output_grad: FloatTensor<B>,
    options: ConvOptions<3>,
) -> FloatTensor<B> {
    let [
        channels_out,
        increment_ci,
        kernel_size_1,
        kernel_size_2,
        kernel_size_3,
    ] = weight_grad.shape().dims();
    let increment_co = channels_out / options.groups;

    let x_swapped = B::float_swap_dims(x, 0, 1);
    let output_grad_swapped = B::float_swap_dims(output_grad, 0, 1);

    for g in 0..options.groups {
        let start_idx_ci = g * increment_ci;
        let end_idx_ci = (g + 1) * increment_ci;
        let start_idx_co = g * increment_co;
        let end_idx_co = (g + 1) * increment_co;

        let x_slice = vec![Slice::new(
            start_idx_ci as isize,
            Some(end_idx_ci as isize),
            1,
        )];
        let x = B::float_slice(x_swapped.clone(), &x_slice);
        let grad_slice = vec![Slice::new(
            start_idx_co as isize,
            Some(end_idx_co as isize),
            1,
        )];
        let grad = B::float_slice(output_grad_swapped.clone(), &grad_slice);
        let mut weight_grad_tmp = B::conv3d(
            x,
            grad,
            None,
            ConvOptions::new(options.dilation, options.padding_begin(), options.stride, 1),
        );
        weight_grad_tmp = B::float_swap_dims(weight_grad_tmp, 0, 1);
        let [
            _,
            _,
            kernel_size_1_tmp,
            kernel_size_2_tmp,
            kernel_size_3_tmp,
        ] = weight_grad_tmp.shape().dims();

        if kernel_size_1_tmp != kernel_size_1
            || kernel_size_2_tmp != kernel_size_2
            || kernel_size_3_tmp != kernel_size_3
        {
            let slices = vec![
                Slice::from(0..increment_co),
                Slice::from(0..increment_ci),
                Slice::from(0..kernel_size_1),
                Slice::from(0..kernel_size_2),
                Slice::from(0..kernel_size_3),
            ];
            weight_grad_tmp = B::float_slice(weight_grad_tmp, &slices);
        }

        weight_grad = B::float_slice_assign(
            weight_grad,
            &[
                Slice::from(start_idx_co..end_idx_co),
                Slice::from(0..increment_ci),
                Slice::from(0..kernel_size_1),
                Slice::from(0..kernel_size_2),
                Slice::from(0..kernel_size_3),
            ],
            weight_grad_tmp,
        );
    }

    weight_grad
}

fn conv_transpose1d_weight_grad_no_groups<B: Backend>(
    x: FloatTensor<B>,
    output_grad: FloatTensor<B>,
    weight_shape: Shape,
    options: ConvTransposeOptions<1>,
) -> FloatTensor<B> {
    let x_swapped = B::float_swap_dims(x, 0, 1);
    let output_grad_swapped = B::float_swap_dims(output_grad, 0, 1);
    let weight_grad_swapped = B::conv1d(
        output_grad_swapped,
        x_swapped,
        None,
        ConvOptions::new(options.dilation, options.padding, options.stride, 1),
    );
    let mut weight_grad = B::float_swap_dims(weight_grad_swapped, 0, 1);

    let grad_shape = weight_grad.shape();
    if grad_shape != weight_shape {
        let slices = vec![
            Slice::from(0..weight_shape[0]),
            Slice::from(0..weight_shape[1]),
            Slice::from(0..weight_shape[2]),
        ];
        weight_grad = B::float_slice(weight_grad, &slices);
    }
    weight_grad
}

fn conv_transpose2d_weight_grad_no_groups<B: Backend>(
    x: FloatTensor<B>,
    output_grad: FloatTensor<B>,
    weight_shape: Shape,
    options: ConvTransposeOptions<2>,
) -> FloatTensor<B> {
    let x_swapped = B::float_swap_dims(x, 0, 1);
    let output_grad_swapped = B::float_swap_dims(output_grad, 0, 1);
    let weight_grad_swapped = B::conv2d(
        output_grad_swapped,
        x_swapped,
        None,
        ConvOptions::new(options.dilation, options.padding, options.stride, 1),
    );
    let mut weight_grad = B::float_swap_dims(weight_grad_swapped, 0, 1);

    let grad_shape = weight_grad.shape();
    if grad_shape != weight_shape {
        let slices = vec![
            Slice::from(0..weight_shape[0]),
            Slice::from(0..weight_shape[1]),
            Slice::from(0..weight_shape[2]),
            Slice::from(0..weight_shape[3]),
        ];
        weight_grad = B::float_slice(weight_grad, &slices);
    }
    weight_grad
}

fn conv_transpose3d_weight_grad_no_groups<B: Backend>(
    x: FloatTensor<B>,
    output_grad: FloatTensor<B>,
    weight_shape: Shape,
    options: ConvTransposeOptions<3>,
) -> FloatTensor<B> {
    let x_swapped = B::float_swap_dims(x, 0, 1);
    let output_grad_swapped = B::float_swap_dims(output_grad, 0, 1);
    let weight_grad_swapped = B::conv3d(
        output_grad_swapped,
        x_swapped,
        None,
        ConvOptions::new(options.dilation, options.padding, options.stride, 1),
    );
    let mut weight_grad = B::float_swap_dims(weight_grad_swapped, 0, 1);

    let grad_shape = weight_grad.shape();
    if grad_shape != weight_shape {
        let slices = vec![
            Slice::from(0..weight_shape[0]),
            Slice::from(0..weight_shape[1]),
            Slice::from(0..weight_shape[2]),
            Slice::from(0..weight_shape[3]),
            Slice::from(0..weight_shape[4]),
        ];
        weight_grad = B::float_slice(weight_grad, &slices);
    }
    weight_grad
}

fn conv_transpose1d_weight_grad_groups<B: Backend>(
    x: FloatTensor<B>,
    mut weight_grad: FloatTensor<B>,
    output_grad: FloatTensor<B>,
    options: ConvTransposeOptions<1>,
) -> FloatTensor<B> {
    let [channels_in, increment_co, kernel_size] = weight_grad.shape().dims();
    let increment_ci = channels_in / options.groups;

    let x_swapped = B::float_swap_dims(x, 0, 1);
    let output_grad_swapped = B::float_swap_dims(output_grad, 0, 1);

    for g in 0..options.groups {
        let start_idx_ci = g * increment_ci;
        let end_idx_ci = (g + 1) * increment_ci;
        let start_idx_co = g * increment_co;
        let end_idx_co = (g + 1) * increment_co;

        let x_slice = vec![Slice::new(
            start_idx_ci as isize,
            Some(end_idx_ci as isize),
            1,
        )];
        let x = B::float_slice(x_swapped.clone(), &x_slice);
        let grad_slice = vec![Slice::new(
            start_idx_co as isize,
            Some(end_idx_co as isize),
            1,
        )];
        let grad = B::float_slice(output_grad_swapped.clone(), &grad_slice);
        let mut weight_grad_tmp = B::conv1d(
            grad,
            x,
            None,
            ConvOptions::new(options.dilation, options.padding, options.stride, 1),
        );
        weight_grad_tmp = B::float_swap_dims(weight_grad_tmp, 0, 1);
        let [_, _, kernel_size_tmp] = weight_grad_tmp.shape().dims();

        if kernel_size_tmp != kernel_size {
            let slices = vec![
                Slice::from(0..increment_ci),
                Slice::from(0..increment_co),
                Slice::from(0..kernel_size),
            ];
            weight_grad_tmp = B::float_slice(weight_grad_tmp, &slices);
        }

        weight_grad = B::float_slice_assign(
            weight_grad,
            &[
                Slice::from(start_idx_ci..end_idx_ci),
                Slice::from(0..increment_co),
                Slice::from(0..kernel_size),
            ],
            weight_grad_tmp,
        );
    }

    weight_grad
}

fn conv_transpose2d_weight_grad_groups<B: Backend>(
    x: FloatTensor<B>,
    mut weight_grad: FloatTensor<B>,
    output_grad: FloatTensor<B>,
    options: ConvTransposeOptions<2>,
) -> FloatTensor<B> {
    let [channels_in, increment_co, kernel_size_1, kernel_size_2] = weight_grad.shape().dims();
    let increment_ci = channels_in / options.groups;

    let x_swapped = B::float_swap_dims(x, 0, 1);
    let output_grad_swapped = B::float_swap_dims(output_grad, 0, 1);

    for g in 0..options.groups {
        let start_idx_ci = g * increment_ci;
        let end_idx_ci = (g + 1) * increment_ci;
        let start_idx_co = g * increment_co;
        let end_idx_co = (g + 1) * increment_co;

        let x_slice = vec![Slice::new(
            start_idx_ci as isize,
            Some(end_idx_ci as isize),
            1,
        )];
        let x = B::float_slice(x_swapped.clone(), &x_slice);
        let grad_slice = vec![Slice::new(
            start_idx_co as isize,
            Some(end_idx_co as isize),
            1,
        )];
        let grad = B::float_slice(output_grad_swapped.clone(), &grad_slice);
        let mut weight_grad_tmp = B::conv2d(
            grad,
            x,
            None,
            ConvOptions::new(options.dilation, options.padding, options.stride, 1),
        );
        weight_grad_tmp = B::float_swap_dims(weight_grad_tmp, 0, 1);
        let [_, _, kernel_size_1_tmp, kernel_size_2_tmp] = weight_grad_tmp.shape().dims();

        if kernel_size_1_tmp != kernel_size_1 || kernel_size_2_tmp != kernel_size_2 {
            let slices = vec![
                Slice::from(0..increment_ci),
                Slice::from(0..increment_co),
                Slice::from(0..kernel_size_1),
                Slice::from(0..kernel_size_2),
            ];
            weight_grad_tmp = B::float_slice(weight_grad_tmp, &slices);
        }

        weight_grad = B::float_slice_assign(
            weight_grad,
            &[
                Slice::from(start_idx_ci..end_idx_ci),
                Slice::from(0..increment_co),
                Slice::from(0..kernel_size_1),
                Slice::from(0..kernel_size_2),
            ],
            weight_grad_tmp,
        );
    }

    weight_grad
}

fn conv_transpose3d_weight_grad_groups<B: Backend>(
    x: FloatTensor<B>,
    mut weight_grad: FloatTensor<B>,
    output_grad: FloatTensor<B>,
    options: ConvTransposeOptions<3>,
) -> FloatTensor<B> {
    let [
        channels_in,
        increment_co,
        kernel_size_1,
        kernel_size_2,
        kernel_size_3,
    ] = weight_grad.shape().dims();
    let increment_ci = channels_in / options.groups;

    let x_swapped = B::float_swap_dims(x, 0, 1);
    let output_grad_swapped = B::float_swap_dims(output_grad, 0, 1);

    for g in 0..options.groups {
        let start_idx_ci = g * increment_ci;
        let end_idx_ci = (g + 1) * increment_ci;
        let start_idx_co = g * increment_co;
        let end_idx_co = (g + 1) * increment_co;

        let x_slice = vec![Slice::new(
            start_idx_ci as isize,
            Some(end_idx_ci as isize),
            1,
        )];
        let x = B::float_slice(x_swapped.clone(), &x_slice);
        let grad_slice = vec![Slice::new(
            start_idx_co as isize,
            Some(end_idx_co as isize),
            1,
        )];
        let grad = B::float_slice(output_grad_swapped.clone(), &grad_slice);
        let mut weight_grad_tmp = B::conv3d(
            grad,
            x,
            None,
            ConvOptions::new(options.dilation, options.padding, options.stride, 1),
        );
        weight_grad_tmp = B::float_swap_dims(weight_grad_tmp, 0, 1);
        let [
            _,
            _,
            kernel_size_1_tmp,
            kernel_size_2_tmp,
            kernel_size_3_tmp,
        ] = weight_grad_tmp.shape().dims();

        if kernel_size_1_tmp != kernel_size_1
            || kernel_size_2_tmp != kernel_size_2
            || kernel_size_3_tmp != kernel_size_3
        {
            let slices = vec![
                Slice::from(0..increment_ci),
                Slice::from(0..increment_co),
                Slice::from(0..kernel_size_1),
                Slice::from(0..kernel_size_2),
                Slice::from(0..kernel_size_3),
            ];
            weight_grad_tmp = B::float_slice(weight_grad_tmp, &slices);
        }
        weight_grad = B::float_slice_assign(
            weight_grad,
            &[
                Slice::from(start_idx_ci..end_idx_ci),
                Slice::from(0..increment_co),
                Slice::from(0..kernel_size_1),
                Slice::from(0..kernel_size_2),
                Slice::from(0..kernel_size_3),
            ],
            weight_grad_tmp,
        );
    }

    weight_grad
}

/// Compute the `padding_out` for a transpose conv that exactly recovers the
/// original `size_in` from `size_out`, accounting for any input elements the
/// forward conv dropped. Shared by `conv{1,2,3}d_x_backward` and the CubeCL
/// dgrad fallback so the two paths can't drift.
pub fn calculate_padding_out(
    kernel_size: usize,
    stride: usize,
    padding: usize,
    dilation: usize,
    size_in: usize,
    size_out: usize,
) -> usize {
    if stride <= 1 {
        return 0;
    }

    // Invert the transpose conv output formula to recover the exact number of
    // input elements that a forward conv would drop for this (size_in, size_out).
    //
    // Forward: size_out = floor((size_in + 2*padding - dilated_kernel) / stride) + 1
    // Transpose: trans_out = (size_out - 1)*stride + dilated_kernel + padding_out - 2*padding
    // Setting trans_out == size_in and solving for padding_out:
    let dilated_kernel = dilation * (kernel_size - 1) + 1;
    let base = (size_out as i64 - 1) * stride as i64 + dilated_kernel as i64 - 2 * padding as i64;
    i64::max(0, size_in as i64 - base) as usize
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_calculate_pool_output_size_ceil_mode_drops_window_in_padding() {
        // PyTorch gives 3: a 4th window would start at index 6 (padded coords),
        // in the trailing padding.
        assert_eq!(calculate_pool_output_size(2, 2, 1, 1, 5, true), 3);

        // With dilation 2 and stride 3: a 3rd window would start at padded index 6.
        assert_eq!(calculate_pool_output_size(2, 3, 1, 2, 5, true), 2);

        // No padding, kernel 1, stride 3: a 3rd window would start at 6, past the input.
        assert_eq!(calculate_pool_output_size(1, 3, 0, 1, 5, true), 2);

        // The last window starts inside the input, so it is kept.
        assert_eq!(calculate_pool_output_size(3, 2, 1, 1, 6, true), 4);
        assert_eq!(calculate_pool_output_size(3, 2, 0, 1, 6, true), 3);
    }

    #[test]
    fn test_calculate_pool_output_size_ceil_mode_kernel_larger_than_input() {
        // PyTorch gives 1: the single window starts at padded index 0, in the input
        // without padding and in the left padding with padding 1.
        assert_eq!(calculate_pool_output_size(3, 2, 0, 1, 2, true), 1);
        assert_eq!(calculate_pool_output_size(4, 2, 1, 1, 2, true), 1);
    }

    #[test]
    fn test_calculate_output_size_1() {
        let kernel_size = 3;
        let stride = 1;
        let padding = 1;
        let size_in = 3;
        let dilation = 1;

        let size_out = calculate_conv_output_size(kernel_size, stride, padding, dilation, size_in);

        assert_eq!(size_out, 3);
    }

    #[test]
    fn test_calculate_output_size_2() {
        let kernel_size = 5;
        let stride = 2;
        let padding = 3;
        let size_in = 27;
        let dilation = 1;

        let size_out = calculate_conv_output_size(kernel_size, stride, padding, dilation, size_in);

        assert_eq!(size_out, 15);
    }

    #[test]
    fn test_calculate_output_size_3() {
        let kernel_size = 5;
        let stride = 2;
        let padding = 3;
        let size_in = 27;
        let dilation = 2;

        let size_out = calculate_conv_output_size(kernel_size, stride, padding, dilation, size_in);

        assert_eq!(size_out, 13);
    }

    #[test]
    fn test_calculate_same_padding_1() {
        let kernel_size = 3;
        let stride = 1;
        let size_in = 3;
        let dilation = 1;

        let padding = calculate_conv_padding(kernel_size, stride, size_in, size_in);
        let size_out = calculate_conv_output_size(kernel_size, stride, padding, dilation, size_in);

        assert_eq!(size_in, size_out, "Expected size");
    }

    #[test]
    fn test_calculate_same_padding_2() {
        let kernel_size = 3;
        let stride = 2;
        let size_in = 7;
        let dilation = 1;

        let padding = calculate_conv_padding(kernel_size, stride, size_in, size_in);
        let size_out = calculate_conv_output_size(kernel_size, stride, padding, dilation, size_in);

        assert_eq!(size_in, size_out, "Expected size");
    }

    #[test]
    fn test_calculate_output_padding_1() {
        let kernel_size = 3;
        let stride = 2;
        let size_in = 7;
        let size_out = 10;
        let dilation = 1;

        let padding = calculate_conv_padding(kernel_size, stride, size_in, size_out);
        let size_out_expected =
            calculate_conv_output_size(kernel_size, stride, padding, dilation, size_in);

        assert_eq!(size_out, size_out_expected, "Expected size");
    }

    #[test]
    fn test_expect_conv2d_output_shape() {
        // in channels: 3
        // out channels: 8
        // size in: [27, 3]
        // kernel size: [5, 3]
        let stride = [2, 1];
        let padding = [(3, 3), (1, 1)];
        let dilation = [2, 1];
        let shape = calculate_conv_output_shape(
            &Shape::new([12, 3, 27, 3]),
            &Shape::new([8, 3, 5, 3]),
            &stride,
            &padding,
            &dilation,
        )
        .unwrap();
        assert_eq!(shape, Shape::new([12, 8, 13, 3]))
    }
}