```python
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

# Fully-fused ConvTranspose2d(k=4, s=2, p=1, op=1) + bias-subtract + tanh,
# specialized for C_in=C_out=64, H=W=256.
#
# Turn 4 failed correctness because raw TF32 tensor-core rounding gave
# ~1e-2-level errors. This version uses the "3xTF32" trick: every operand is
# split into hi = tf32(x) and lo = x - hi (exact), and the product is
# accumulated as hi*hi + lo*hi + hi*lo. That restores ~fp32 accuracy
# (error ~1e-5) while still running on TF32 tensor cores at ~1/3 the TF32
# peak, which is still several times faster than the fp32 SIMT dgrad kernels.
#
# Pipeline (2 kernels, no cuDNN, output produced directly in NHWC):
#   1) nchw_to_nhwc_pad_kernel: NCHW -> zero-padded NHWC scratch
#      [N][258][290][64] via a bank-conflict-free shared-memory transpose.
#   2) tconv_gemm_kernel: implicit-GEMM transposed convolution using the
#      output-parity decomposition (each output pixel touches a 2x2x64
#      input window and one of 4 pre-swizzled 256x64 weight matrices).
#      wmma 16x16x8 TF32 tensor cores, 3 mma per operand pair (3xTF32).
#      Bias-subtract + tanh fused in the epilogue with float2 stores.
#
# Falls back to the reference implementation for other shapes.

custom_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <mma.h>
#include <cstdint>

using namespace nvcuda;

// round float -> tf32 (returned as float with the low 13 mantissa bits zeroed)
__device__ __forceinline__ float tf32_rn(float x) {
    uint32_t u;
    asm("cvt.rna.tf32.f32 %0, %1;" : "=r"(u) : "f"(x));
    return __uint_as_float(u);
}

// ---------------------------------------------------------------------------
// NCHW -> padded NHWC with zero border.
// xp: [N][RP][CP][C].  Only rows 0 / RP-1 and cols 0 / W+1 must be zero;
// columns W+2..CP-1 are scratch (read only by masked-out fragments).
// Shared tile stride W+9 (odd) keeps the transposed store conflict-free.
// ---------------------------------------------------------------------------
template <int C, int H, int W, int RP, int CP>
__global__ void __launch_bounds__(256)
nchw_to_nhwc_pad_kernel(const float* __restrict__ x,
                        float* __restrict__ xp) {
    constexpr int SW = W + 9;
    extern __shared__ float sh_raw[];
    float (*sh)[SW] = reinterpret_cast<float (*)[SW]>(sh_raw);

    const int n   = blockIdx.x;
    const int rp  = blockIdx.y;
    const int tid = threadIdx.x;

    float* __restrict__ orow = xp + (((long long)n * RP + rp) * CP) * C;

    if (rp == 0 || rp == RP - 1) {
        float4* o4 = reinterpret_cast<float4*>(orow);
        const float4 z = make_float4(0.f, 0.f, 0.f, 0.f);
        const int total4 = CP * (C / 4);
        for (int i = tid; i < total4; i += 256) o4[i] = z;
        return;
    }

    const int r = rp - 1;
    const float* __restrict__ xr = x + (((long long)n * C) * H + r) * W;

    #pragma unroll
    for (int c0 = 0; c0 < C; c0 += 8) {
        const float* src = xr + (long long)c0 * (H * W);
        #pragma unroll
        for (int i = 0; i < 8; ++i)
            sh[c0 + i][tid] = src[(long long)i * (H * W) + tid];
    }
    __syncthreads();

    if (tid < C) {
        orow[tid] = 0.f;                 // left border column
        orow[(W + 1) * C + tid] = 0.f;   // right border column
    }

    #pragma unroll
    for (int c0 = 0; c0 < W; c0 += 4) {
        const int cc = c0 + (tid >> 6);
        const int ci = tid & 63;
        orow[(cc + 1) * C + ci] = sh[ci][cc];
    }
}

// ---------------------------------------------------------------------------
// Fused transposed-conv implicit GEMM (3xTF32 tensor cores) + bias-sub + tanh.
// out[n, y, x, co] = tanh( sum_{t,ci} xp[n, rp, cp, ci] * W[p, t*C+ci, co]
//                          - bias[co] )
// Block: 8 same-parity output rows x 32 same-parity output cols, 256 threads.
// Warp: 1 output row x 32 cols = 2x4 fragments of 16x16.
// Per-parity weights (256x64 floats = 64KB) staged once per block in smem.
// ---------------------------------------------------------------------------
template <int C, int RP, int CP, int HOUT>
__global__ void __launch_bounds__(256, 2)
tconv_gemm_kernel(const float* __restrict__ xp,
                  const float* __restrict__ wq,   // [4][4C][C]
                  const float* __restrict__ bias, // [C]
                  float* __restrict__ out) {
    constexpr int K = 4 * C;
    extern __shared__ float sB[];       // [K][C]

    const int p  = blockIdx.z & 3;
    const int n  = blockIdx.z >> 2;
    const int py = p >> 1;
    const int px = p & 1;
    const int Oy = (HOUT + 1 - py) >> 1;
    const int Ox = (HOUT + 1 - px) >> 1;

    // uniform per-block early exits (before any __syncthreads)
    if (blockIdx.y * 8 >= Oy) return;
    if (blockIdx.x * 32 >= Ox) return;

    {
        const float4* src = reinterpret_cast<const float4*>(wq + (size_t)p * (K * C));
        float4* dst = reinterpret_cast<float4*>(sB);
        #pragma unroll
        for (int i = 0; i < 16; ++i)
            dst[threadIdx.x + i * 256] = src[threadIdx.x + i * 256];
    }
    __syncthreads();

    const int warp = threadIdx.x >> 5;
    const int lane = threadIdx.x & 31;
    const int orow = blockIdx.y * 8 + warp;
    if (orow >= Oy) return;
    const int y   = 2 * orow + py;
    const int obx = blockIdx.x * 32;

    wmma::fragment<wmma::accumulator, 16, 16, 8, float> acc[2][4];
    #pragma unroll
    for (int i = 0; i < 2; ++i)
        #pragma unroll
        for (int j = 0; j < 4; ++j)
            wmma::fill_fragment(acc[i][j], 0.0f);

    const long long nrow = ((long long)n * RP + orow + py) * CP + obx + px;

    #pragma unroll
    for (int t = 0; t < 4; ++t) {
        const long long abase = (nrow + (t >> 1) * CP + (t & 1)) * C;
        #pragma unroll
        for (int c0 = 0; c0 < C; c0 += 8) {
            wmma::fragment<wmma::matrix_a, 16, 16, 8,
                           wmma::precision::tf32, wmma::row_major> af0, af1, ah0, al0, ah1, al1;
            wmma::load_matrix_sync(af0, xp + abase + c0, C);
            wmma::load_matrix_sync(af1, xp + abase + c0 + 16 * C, C);
            #pragma unroll
            for (int e = 0; e < 4; ++e) {
                const float v0 = af0.x[e], v1 = af1.x[e];
                const float h0 = tf32_rn(v0), h1 = tf32_rn(v1);
                ah0.x[e] = h0; al0.x[e] = v0 - h0;
                ah1.x[e] = h1; al1.x[e] = v1 - h1;
            }
            const float* wptr = sB + (t * C + c0) * C;
            #pragma unroll
            for (int j = 0; j < 4; ++j) {
                wmma::fragment<wmma::matrix_b, 16, 16, 8,
                               wmma::precision::tf32, wmma::row_major> bf, bh, bl;
                wmma::load_matrix_sync(bf, wptr + j * 16, C);
                #pragma unroll
                for (int e = 0; e < 4; ++e) {
                    const float v = bf.x[e];
                    const float h = tf32_rn(v);
                    bh.x[e] = h; bl.x[e] = v - h;
                }
                // 3xTF32: hi*hi + lo*hi + hi*lo  (dropped lo*lo ~ 2^-44)
                wmma::mma_sync(acc[0][j], ah0, bh, acc[0][j]);
                wmma::mma_sync(acc[1][j], ah1, bh, acc[1][j]);
                wmma::mma_sync(acc[0][j], al0, bh, acc[0][j]);
                wmma::mma_sync(acc[1][j], al1, bh, acc[1][j]);
                wmma::mma_sync(acc[0][j], ah0, bl, acc[0][j]);
                wmma::mma_sync(acc[1][j], ah1, bl, acc[1][j]);
            }
        }
    }

    // epilogue: fused bias-subtraction + tanh, vectorized NHWC stores
    const int g  = lane >> 2;
    const int tg = lane & 3;
    float bw[4][4];
    #pragma unroll
    for (int j = 0; j < 4; ++j) {
        const int co = 16 * j + 2 * tg;
        bw[j][0] = bias[co];     bw[j][1] = bias[co + 1];
        bw[j][2] = bias[co + 8]; bw[j][3] = bias[co + 9];
    }
    float* __restrict__ obase = out + (((long long)n * HOUT + y) * HOUT) * C;
    #pragma unroll
    for (int i = 0; i < 2; ++i) {
        const int xl0 = obx + 16 * i + g;
        const int xl1 = xl0 + 8;
        const bool v0 = xl0 < Ox;
        const bool v1 = xl1 < Ox;
        if (!(v0 || v1)) continue;
        const long long xr0 = 2 * xl0 + px;
        const long long xr1 = 2 * xl1 + px;
        #pragma unroll
        for (int j = 0; j < 4; ++j) {
            const int co = 16 * j + 2 * tg;
            float* e = acc[i][j].x;
            if (v0) {
                float2* q0 = reinterpret_cast<float2*>(obase + xr0 * C + co);
                q0[0] = make_float2(tanhf(e[0] - bw[j][0]), tanhf(e[1] - bw[j][1]));
                q0[4] = make_float2(tanhf(e[4] - bw[j][2]), tanhf(e[5] - bw[j][3]));
            }
            if (v1) {
                float2* q1 = reinterpret_cast<float2*>(obase + xr1 * C + co);
                q1[0] = make_float2(tanhf(e[2] - bw[j][0]), tanhf(e[3] - bw[j][1]));
                q1[4] = make_float2(tanhf(e[6] - bw[j][2]), tanhf(e[7] - bw[j][3]));
            }
        }
    }
}

// ---------------------------------------------------------------------------
// host wrappers
// ---------------------------------------------------------------------------
torch::Tensor nchw_to_nhwc_pad_cuda(torch::Tensor x) {
    TORCH_CHECK(x.is_cuda() && x.scalar_type() == at::kFloat, "need float32 CUDA tensor");
    TORCH_CHECK(x.dim() == 4, "x must be NCHW");
    auto xc = x.contiguous();
    const int n = xc.size(0), C = xc.size(1), H = xc.size(2), W = xc.size(3);
    TORCH_CHECK(C == 64 && H == 256 && W == 256, "specialized kernel: C=64, H=W=256");

    constexpr int RP = 258, CP = 290, SW = 256 + 9;
    constexpr int SMEM = 64 * SW * 4;
    static bool configured = false;
    if (!configured) {
        cudaFuncSetAttribute((const void*)nchw_to_nhwc_pad_kernel<64, 256, 256, RP, CP>,
                             cudaFuncAttributeMaxDynamicSharedMemorySize, SMEM);
        configured = true;
    }

    auto xp = torch::empty({(long long)n, (long long)RP, (long long)CP, 64LL},
                           xc.options());
    dim3 grid(n, RP);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    nchw_to_nhwc_pad_kernel<64, 256, 256, RP, CP><<<grid, 256, SMEM, stream>>>(
        xc.data_ptr<float>(), xp.data_ptr<float>());
    return xp;
}

torch::Tensor tconv_gemm_cuda(torch::Tensor xp, torch::Tensor wq,
                              torch::Tensor bias, int64_t n_batch) {
    TORCH_CHECK(xp.is_cuda() && wq.is_cuda() && bias.is_cuda(), "CUDA tensors required");
    const int n = (int)n_batch;
    constexpr int HOUT = 513;
    constexpr int SMEM = 256 * 64 * 4;
    static bool configured = false;
    if (!configured) {
        cudaFuncSetAttribute((const void*)tconv_gemm_kernel<64, 258, 290, HOUT>,
                             cudaFuncAttributeMaxDynamicSharedMemorySize, SMEM);
        configured = true;
    }

    auto bias_c = bias.contiguous();
    auto out = torch::empty({(long long)n, (long long)HOUT, (long long)HOUT, 64LL},
                            xp.options());
    dim3 grid(9, 33, n * 4);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    tconv_gemm_kernel<64, 258, 290, HOUT><<<grid, 256, SMEM, stream>>>(
        xp.data_ptr<float>(), wq.data_ptr<float>(), bias_c.data_ptr<float>(),
        out.data_ptr<float>());
    return out;
}
"""

custom_cpp_source = """
torch::Tensor nchw_to_nhwc_pad_cuda(torch::Tensor x);
torch::Tensor tconv_gemm_cuda(torch::Tensor xp, torch::Tensor wq, torch::Tensor bias, int64_t n_batch);
"""

custom_ops = load_inline(
    name="fused_tconv_bias_tanh_3xtf32",
    cpp_sources=custom_cpp_source,
    cuda_sources=custom_source,
    functions=["nchw_to_nhwc_pad_cuda", "tconv_gemm_cuda"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "-use_fast_math"],
)


class ModelNew(nn.Module):
    """
    Optimized model: replaces the cuDNN transposed convolution + bias + tanh
    chain with two custom kernels: an NCHW->padded-NHWC conversion kernel and
    a single implicit-GEMM tensor-core kernel using 3xTF32 precision
    (hi/lo operand splitting, 3 mma per product) that matches fp32 accuracy
    while running on TF32 tensor cores. Bias-subtraction and tanh are fused
    into the GEMM epilogue and the result is emitted directly in
    channels_last layout (no NHWC->NCHW conversion kernel).
    Specialized for C_in=C_out=64, H=W=256, k=4, s=2, p=1, op=1; falls back
    to the reference otherwise.
    """
    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._specialized = (in_channels == 64 and out_channels == 64
                             and kernel_size == 4 and stride == 2
                             and padding == 1 and output_padding == 1)
        # flat kernel-tap indices: index = p*4 + t, p = (py,px) output parity,
        # t = (dr,dc) tap slot; ky[py][dr] = ((3,1),(2,0)), same for kx.
        ky = [3, 3, 1, 1,   3, 3, 1, 1,   2, 2, 0, 0,   2, 2, 0, 0]
        kx = [3, 1, 3, 1,   2, 0, 2, 0,   3, 1, 3, 1,   2, 0, 2, 0]
        self.register_buffer("_ky", torch.tensor(ky, dtype=torch.long), persistent=False)
        self.register_buffer("_kx", torch.tensor(kx, dtype=torch.long), persistent=False)
        self._wq = None
        self._wq_key = None

    def _get_swizzle(self):
        w = self.conv_transpose.weight
        key = (w.data_ptr(), w._version, w.device)
        if self._wq is None or self._wq_key != key:
            Ci, Co = w.shape[0], w.shape[1]
            with torch.no_grad():
                sel = w.detach()[:, :, self._ky, self._kx]      # [Ci, Co, 16]
                wq = sel.permute(2, 0, 1).reshape(4, 4 * Ci, Co)  # [4, 4Ci, Co]
                self._wq = wq.float().contiguous()
            self._wq_key = key
        return self._wq

    def forward(self, x):
        if (not self._specialized) or (not x.is_cuda) or x.shape[2] != 256 or x.shape[3] != 256 \
                or x.dtype != torch.float32:
            return torch.tanh(self.conv_transpose(x) - self.bias)
        wq = self._get_swizzle()
        n = x.shape[0]
        xp = custom_ops.nchw_to_nhwc_pad_cuda(x)
        out = custom_ops.tconv_gemm_cuda(xp, wq, self.bias, n)
        # channels_last NCHW view (values identical to reference)
        return out.permute(0, 3, 1, 2)
```