```python
import functools

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline


_CUDA_SOURCE = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>
#include <vector>

template <typename scalar_t>
__global__ void pack_input_kernel(
    const scalar_t* __restrict__ src,
    scalar_t* __restrict__ dst,
    int64_t channels,
    int64_t spatial
) {
    __shared__ scalar_t tile[32][33];

    const int tx = threadIdx.x;
    const int ty = threadIdx.y;
    const int64_t n = blockIdx.z;
    const int64_t c0 = int64_t(blockIdx.y) * 32;
    const int64_t s0 = int64_t(blockIdx.x) * 32;

    #pragma unroll
    for (int j = 0; j < 32; j += 8) {
        const int64_t c = c0 + ty + j;
        const int64_t s = s0 + tx;
        if (c < channels && s < spatial) {
            tile[ty + j][tx] = src[(n * channels + c) * spatial + s];
        }
    }

    __syncthreads();

    #pragma unroll
    for (int j = 0; j < 32; j += 8) {
        const int64_t c = c0 + tx;
        const int64_t s = s0 + ty + j;
        if (c < channels && s < spatial) {
            dst[(n * spatial + s) * channels + c] = tile[tx][ty + j];
        }
    }
}

template <typename scalar_t>
__global__ void pack_weight_kernel(
    const scalar_t* __restrict__ src,
    scalar_t* __restrict__ dst,
    int64_t count,
    int64_t input_per_group,
    int64_t output_per_group,
    int64_t kh,
    int64_t kw
) {
    for (
        int64_t index = int64_t(blockIdx.x) * blockDim.x + threadIdx.x;
        index < count;
        index += int64_t(blockDim.x) * gridDim.x
    ) {
        // Destination physical layout: [output_channel, kh, kw, input_channel].
        int64_t rem = index;
        const int64_t ic = rem % input_per_group;
        rem /= input_per_group;
        const int64_t x = rem % kw;
        rem /= kw;
        const int64_t y = rem % kh;
        const int64_t oc = rem / kh;

        const int64_t group = oc / output_per_group;
        const int64_t local_oc = oc % output_per_group;
        const int64_t global_ic = group * input_per_group + ic;

        const int64_t source_index =
            ((global_ic * output_per_group + local_oc) * kh
             + (kh - 1 - y)) * kw + (kw - 1 - x);

        dst[index] = src[source_index];
    }
}

std::vector<torch::Tensor> pack_transpose_conv(
    torch::Tensor input,
    torch::Tensor weight,
    int64_t groups
) {
    const c10::cuda::CUDAGuard guard(input.device());
    const auto stream = at::cuda::getCurrentCUDAStream();

    const int64_t n = input.size(0);
    const int64_t cin = input.size(1);
    const int64_t spatial = input.size(2) * input.size(3);
    const int64_t kh = weight.size(2);
    const int64_t kw = weight.size(3);
    const int64_t input_per_group = cin / groups;
    const int64_t output_per_group = weight.size(1);
    const int64_t cout = output_per_group * groups;

    torch::Tensor packed_input;
    const bool needs_input_pack =
        !input.is_contiguous(at::MemoryFormat::ChannelsLast);

    if (needs_input_pack) {
        packed_input = torch::empty(
            input.sizes(),
            input.options().memory_format(at::MemoryFormat::ChannelsLast)
        );
    } else {
        packed_input = input;
    }

    auto packed_weight = torch::empty(
        {cout, input_per_group, kh, kw},
        weight.options().memory_format(at::MemoryFormat::ChannelsLast)
    );

    const int64_t count = weight.numel();
    const int weight_blocks =
        static_cast<int>(std::min<int64_t>((count + 255) / 256, 4096));

    AT_DISPATCH_FLOATING_TYPES_AND2(
        at::ScalarType::Half,
        at::ScalarType::BFloat16,
        input.scalar_type(),
        "pack_transpose_conv",
        [&] {
            if (needs_input_pack) {
                const dim3 threads(32, 8);
                const dim3 blocks(
                    (spatial + 31) / 32,
                    (cin + 31) / 32,
                    n
                );
                pack_input_kernel<scalar_t><<<blocks, threads, 0, stream>>>(
                    input.data_ptr<scalar_t>(),
                    packed_input.data_ptr<scalar_t>(),
                    cin,
                    spatial
                );
                C10_CUDA_KERNEL_LAUNCH_CHECK();
            }

            pack_weight_kernel<scalar_t><<<weight_blocks, 256, 0, stream>>>(
                weight.data_ptr<scalar_t>(),
                packed_weight.data_ptr<scalar_t>(),
                count,
                input_per_group,
                output_per_group,
                kh,
                kw
            );
            C10_CUDA_KERNEL_LAUNCH_CHECK();
        }
    );

    return {packed_input, packed_weight};
}
"""


@functools.lru_cache(maxsize=1)
def _extension():
    return load_inline(
        name="asymmetric_transpose_conv_pack_v1",
        cpp_sources="""
#include <torch/extension.h>
#include <vector>
std::vector<torch::Tensor> pack_transpose_conv(
    torch::Tensor input, torch::Tensor weight, int64_t groups);
""",
        cuda_sources=_CUDA_SOURCE,
        functions=["pack_transpose_conv"],
        extra_cflags=["-O3"],
        extra_cuda_cflags=["-O3"],
        verbose=False,
    )


class _PackTransposeConv(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, weight, groups):
        ctx.groups = groups
        ctx.weight_shape = tuple(weight.shape)
        packed_x, packed_weight = _extension().pack_transpose_conv(
            x, weight, groups
        )
        return packed_x, packed_weight

    @staticmethod
    def backward(ctx, grad_x, grad_weight):
        original_grad_weight = None
        if grad_weight is not None and ctx.needs_input_grad[1]:
            cin, output_per_group, kh, kw = ctx.weight_shape
            original_grad_weight = (
                grad_weight.reshape(
                    ctx.groups, output_per_group, cin // ctx.groups, kh, kw
                )
                .flip(-2, -1)
                .permute(0, 2, 1, 3, 4)
                .reshape(ctx.weight_shape)
                .contiguous()
            )

        return (
            grad_x if ctx.needs_input_grad[0] else None,
            original_grad_weight,
            None,
        )


class ModelNew(nn.Module):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: tuple,
        stride: tuple = (1, 1),
        padding: tuple = (0, 0),
        output_padding: tuple = (0, 0),
        dilation: tuple = (1, 1),
        groups: int = 1,
        bias: bool = False,
    ):
        super().__init__()
        self.conv_transpose2d = nn.ConvTranspose2d(
            in_channels,
            out_channels,
            kernel_size,
            stride=stride,
            padding=padding,
            output_padding=output_padding,
            dilation=dilation,
            groups=groups,
            bias=bias,
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        conv = self.conv_transpose2d
        weight = conv.weight

        ph = conv.dilation[0] * (conv.kernel_size[0] - 1) - conv.padding[0]
        pw = conv.dilation[1] * (conv.kernel_size[1] - 1) - conv.padding[1]

        fast_path = (
            x.is_cuda
            and x.ndim == 4
            and x.device == weight.device
            and x.dtype == weight.dtype
            and x.dtype in (
                torch.float16,
                torch.bfloat16,
                torch.float32,
                torch.float64,
            )
            and conv.stride == (1, 1)
            and conv.output_padding == (0, 0)
            and ph >= 0
            and pw >= 0
            and 0 < x.shape[0] <= 65535
            and x.shape[1] == conv.in_channels
            and 0 < x.shape[1] <= 65535 * 32
            and x.shape[2] > 0
            and x.shape[3] > 0
            and weight.is_contiguous()
            and (
                x.is_contiguous()
                or x.is_contiguous(memory_format=torch.channels_last)
            )
        )

        if not fast_path:
            return conv(x)

        packed_x, packed_weight = _PackTransposeConv.apply(
            x, weight, conv.groups
        )

        # Unit-stride transposed convolution equals a forward convolution
        # with reversed spatial kernels and exchanged channel dimensions.
        return F.conv2d(
            packed_x,
            packed_weight,
            bias=conv.bias,
            stride=1,
            padding=(ph, pw),
            dilation=conv.dilation,
            groups=conv.groups,
        )
```