```python
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

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

torch.backends.cudnn.benchmark = True
torch.backends.cudnn.allow_tf32 = True

# Fused CUDA kernel: phase-interleave (transposed-conv output reconstruction)
# + bias subtraction + tanh in a single pass over memory.
fused_source = r"""
#include <torch/extension.h>
#include <cuda_runtime.h>

__global__ void phase_interleave_bias_tanh_kernel(
    const float* __restrict__ y,     // (B, 4*CO, YS, YS) contiguous
    const float* __restrict__ bias,  // (CO,)
    float* __restrict__ out,         // (B, CO, OS, OS)
    const int CO,
    const int YS,
    const int OS)
{
    const int total = OS * OS;
    const long idx = (long)blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;

    const int oh = (int)(idx / OS);
    const int ow = (int)(idx - (long)oh * OS);

    const int bc = blockIdx.y;          // b * CO + co
    const int b  = bc / CO;
    const int co = bc - b * CO;

    const int p = oh & 1;
    const int q = ow & 1;

    const long yc = ((long)b * 4 + (p * 2 + q)) * CO + co;
    const long yidx = (yc * YS + ((oh >> 1) + p)) * YS + ((ow >> 1) + q);
    const long oidx = (long)bc * total + idx;

    out[oidx] = tanhf(y[yidx] - bias[co]);
}

torch::Tensor fused_interleave_bias_tanh(torch::Tensor y, torch::Tensor bias, int64_t co) {
    const int B  = y.size(0);
    const int CO = (int)co;
    const int YS = y.size(2);
    const int OS = 2 * YS - 1;

    auto out = torch::empty({B, CO, OS, OS}, y.options());

    const long total = (long)OS * OS;
    const int block = 256;
    const int gx = (int)((total + block - 1) / block);
    dim3 grid(gx, B * CO);

    phase_interleave_bias_tanh_kernel<<<grid, block>>>(
        y.data_ptr<float>(),
        bias.data_ptr<float>(),
        out.data_ptr<float>(),
        CO, YS, OS);
    return out;
}
"""

fused_cpp_source = "torch::Tensor fused_interleave_bias_tanh(torch::Tensor y, torch::Tensor bias, int64_t co);"

fused_op = load_inline(
    name="fused_interleave_bias_tanh",
    cpp_sources=fused_cpp_source,
    cuda_sources=fused_source,
    functions=["fused_interleave_bias_tanh"],
    verbose=False,
    extra_cuda_cflags=["-O3", "--use_fast_math"],
)


class ModelNew(nn.Module):
    """
    Optimized version:
      ConvTranspose2d(k=4,s=2,p=1,op=1) is computed exactly as a single
      Conv2d(k=2,p=1) with 4*out_channels phase kernels, followed by one
      custom CUDA kernel that interleaves the 4 phases, subtracts the bias
      and applies tanh in a single fused memory pass.
    """
    def __init__(self, in_channels, out_channels, kernel_size, bias_shape, stride=2, padding=1, output_padding=1):
        super(ModelNew, self).__init__()
        self.conv_transpose = nn.ConvTranspose2d(in_channels, out_channels, kernel_size,
                                                 stride=stride, padding=padding,
                                                 output_padding=output_padding)
        self.bias = nn.Parameter(torch.randn(bias_shape))
        self._fast = (isinstance(kernel_size, int) and kernel_size == 4 and
                      stride == 2 and padding == 1 and output_padding == 1)
        self._pw = None
        self._pb = None
        self._key = None

    def _build_phase_weight(self):
        w = self.conv_transpose.weight  # (CI, CO, 4, 4)
        # phase (p,q) kernel taps: w[ci, co, rows[p]], cols[q]
        kee = w[:, :, 1::2, 1::2].flip([2, 3]).permute(1, 0, 2, 3)  # taps (3-2a, 3-2b)
        keo = w[:, :, 1::2, 0::2].flip([2, 3]).permute(1, 0, 2, 3)  # taps (3-2a, 2-2b)
        koe = w[:, :, 0::2, 1::2].flip([2, 3]).permute(1, 0, 2, 3)  # taps (2-2a, 3-2b)
        koo = w[:, :, 0::2, 0::2].flip([2, 3]).permute(1, 0, 2, 3)  # taps (2-2a, 2-2b)
        self._pw = torch.cat([kee, keo, koe, koo], dim=0).contiguous()
        cb = self.conv_transpose.bias
        self._pb = cb.repeat(4).contiguous() if cb is not None else None
        self._key = (w.data_ptr(), w._version, w.device, w.dtype,
                     None if cb is None else (cb.data_ptr(), cb._version))

    def forward(self, x):
        if (self._fast and x.is_cuda and x.dtype == torch.float32 and x.dim() == 4
                and self.bias.numel() == self.conv_transpose.out_channels):
            w = self.conv_transpose.weight
            cb = self.conv_transpose.bias
            key = (w.data_ptr(), w._version, w.device, w.dtype,
                   None if cb is None else (cb.data_ptr(), cb._version))
            if key != self._key:
                self._build_phase_weight()
            y = F.conv2d(x, self._pw, self._pb, stride=1, padding=1)
            return fused_op.fused_interleave_bias_tanh(
                y, self.bias.view(-1), self.conv_transpose.out_channels)
        x = self.conv_transpose(x)
        x = x - self.bias
        return torch.tanh(x)
```