```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

# Custom CUDA kernel that fuses the entire element-wise chain after the
# transposed convolution:
#     Mish -> add (add_value) -> Hardtanh(-1, 1) -> scale
# into a single memory-bound pass over the tensor (vectorized with float4).
#
# Mish identity used for a cheap, numerically stable evaluation:
#     mish(z) = z * tanh(softplus(z))
#     tanh(ln(1+e^z)) = ((1+e^z)^2 - 1) / ((1+e^z)^2 + 1)
#                     = (e^z * (e^z + 2)) / (e^z * (e^z + 2) + 2)
fused_epilogue_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>

namespace {

__device__ __forceinline__ float epilogue(float z, float add_value, float scale) {
    float m;
    if (z > 20.0f) {
        // tanh(softplus(z)) == 1.0f exactly in fp32 for z > 20, so mish(z) == z
        m = z;
    } else {
        float e = expf(z);
        float num = e * (e + 2.0f);   // e^{2z} + 2 e^z
        m = z * (num / (num + 2.0f)); // mish
    }
    float v = m + add_value;               // add
    v = fminf(fmaxf(v, -1.0f), 1.0f);      // hardtanh(min=-1, max=1)
    return v * scale;                      // scale
}

__global__ void fused_epilogue_vec4_kernel(const float4* __restrict__ in,
                                           float4* __restrict__ out,
                                           int n4,
                                           float add_value,
                                           float scale) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n4) {
        float4 v = in[i];
        v.x = epilogue(v.x, add_value, scale);
        v.y = epilogue(v.y, add_value, scale);
        v.z = epilogue(v.z, add_value, scale);
        v.w = epilogue(v.w, add_value, scale);
        out[i] = v;
    }
}

__global__ void fused_epilogue_scalar_kernel(const float* __restrict__ in,
                                             float* __restrict__ out,
                                             long n,
                                             float add_value,
                                             float scale) {
    long i = (long)blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) {
        out[i] = epilogue(in[i], add_value, scale);
    }
}

}  // namespace

torch::Tensor fused_mish_add_hardtanh_scale(torch::Tensor x, double add_value, double scale) {
    TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");

    auto xc = x.contiguous();
    auto out = torch::empty_like(xc);
    long n = xc.numel();
    if (n == 0) {
        return out;
    }

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    const int block = 256;
    const float av = (float)add_value;
    const float sc = (float)scale;

    if ((n & 3) == 0) {
        long n4 = n >> 2;
        long grid = (n4 + block - 1) / block;
        fused_epilogue_vec4_kernel<<<(unsigned int)grid, block, 0, stream>>>(
            reinterpret_cast<const float4*>(xc.data_ptr<float>()),
            reinterpret_cast<float4*>(out.data_ptr<float>()),
            (int)n4, av, sc);
    } else {
        long grid = (n + block - 1) / block;
        fused_epilogue_scalar_kernel<<<(unsigned int)grid, block, 0, stream>>>(
            xc.data_ptr<float>(), out.data_ptr<float>(), n, av, sc);
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
"""

fused_epilogue_cpp_source = (
    "torch::Tensor fused_mish_add_hardtanh_scale(torch::Tensor x, double add_value, double scale);"
)

# Compile the fused Mish + Add + Hardtanh + Scale CUDA kernel
fused_epilogue = load_inline(
    name="fused_mish_add_hardtanh_scale_v1",
    cpp_sources=fused_epilogue_cpp_source,
    cuda_sources=fused_epilogue_cuda_source,
    functions=["fused_mish_add_hardtanh_scale"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized version of Model.

    The transposed convolution is kept as the cuDNN-backed nn.ConvTranspose2d
    (it is a highly tuned implicit-GEMM / tensor-core implementation), while
    the four following element-wise operators (Mish, add, Hardtanh, scale) are
    replaced by a single fused custom CUDA kernel. This eliminates 3 full
    read+write passes over the large (128 x 64 x 256 x 256) intermediate
    tensors and 3 kernel launches.
    """
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, output_padding, add_value, scale):
        super(ModelNew, self).__init__()
        torch.backends.cudnn.benchmark = True
        self.conv_transpose = nn.ConvTranspose2d(in_channels, out_channels, kernel_size, stride, padding, output_padding)
        self.add_value = float(add_value)
        self.scale = float(scale)

    def forward(self, x):
        x = self.conv_transpose(x)
        if x.is_cuda and x.dtype == torch.float32:
            return fused_epilogue.fused_mish_add_hardtanh_scale(x, self.add_value, self.scale)
        # Fallback (CPU / non-float32) path: original PyTorch ops
        x = torch.nn.functional.mish(x)
        x = x + self.add_value
        x = torch.nn.functional.hardtanh(x, min_val=-1, max_val=1)
        x = x * self.scale
        return x
```