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

# ---------------------------------------------------------------------------
# Enable TF32 for any fallback FP32 matmul path.
# ---------------------------------------------------------------------------
def _enable_tf32():
    try:
        torch.backends.cuda.matmul.allow_tf32 = True
    except Exception:
        pass
    try:
        torch.backends.cudnn.allow_tf32 = True
    except Exception:
        pass
    try:
        torch.set_float32_matmul_precision("high")
    except Exception:
        pass

_enable_tf32()

# ---------------------------------------------------------------------------
# The (1024x8192) @ (8192x8192) GEMM dominates the runtime. FP16-input /
# FP32-accumulate tensor cores reach ~2x the TF32 throughput with essentially
# the same numerical error (both round inputs to 10 explicit mantissa bits).
#
# This extension provides a single fused entry point:
#   1) cast x to FP16 (weight is pre-cast and cached on the Python side)
#   2) cuBLASLt FP16->FP32 GEMM with FP32 compute, an autotuned algorithm
#      (per-shape one-time selection among the top heuristic results), and a
#      fused BIAS epilogue
#   3) one fused kernel: per-row instance normalization (biased variance,
#      matching InstanceNorm2d on a (B,1,1,F) tensor), + residual add,
#      * elementwise multiply. The row is cached in shared memory and
#      streaming loads/stores minimize L2 pollution.
# ---------------------------------------------------------------------------
fused_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <cublas_v2.h>
#include <cublasLt.h>
#include <unordered_map>

namespace {

struct LtPlan {
    cublasLtMatmulAlgo_t algo;
    size_t workspace;
    bool valid;
};

std::unordered_map<int64_t, LtPlan>& plan_cache() {
    static auto* m = new std::unordered_map<int64_t, LtPlan>();
    return *m;
}

torch::Tensor& workspace_tensor() {
    static auto* t = new torch::Tensor();
    return *t;
}

cublasLtHandle_t lt_handle() {
    static cublasLtHandle_t h = nullptr;
    if (h == nullptr) {
        if (cublasLtCreate(&h) != CUBLAS_STATUS_SUCCESS) h = nullptr;
    }
    return h;
}

void* get_workspace(size_t bytes, const torch::TensorOptions& opt) {
    if (bytes == 0) return nullptr;
    auto& wt = workspace_tensor();
    if (!wt.defined() || (size_t)wt.numel() < bytes ||
        wt.device() != opt.device()) {
        wt = torch::empty({(int64_t)bytes}, opt.dtype(torch::kUInt8));
    }
    return wt.data_ptr();
}

}  // namespace

// ---------------------------------------------------------------------------
// Autotuned cuBLASLt GEMM:
//   out (M,N) fp32 = x (M,K) fp16 @ w (N,K)^T fp16  [+ bias epilogue]
// Row-major trick: C(MxN) row-major == C^T(NxM) col-major = W @ X^T.
// Returns true on success (bias applied in the epilogue when provided).
// ---------------------------------------------------------------------------
static bool gemm_lt_fp16(const void* x, const void* w, const float* bias,
                         float* out, int64_t M, int64_t N, int64_t K,
                         const torch::TensorOptions& opt,
                         cudaStream_t stream) {
    cublasLtHandle_t lt = lt_handle();
    if (lt == nullptr) return false;
    if (M <= 0 || N <= 0 || K <= 0) return false;
    if (M >= (1LL << 20) || N >= (1LL << 20) || K >= (1LL << 20)) return false;

    cublasLtMatmulDesc_t opDesc = nullptr;
    cublasLtMatrixLayout_t Adesc = nullptr, Bdesc = nullptr, Cdesc = nullptr;
    bool ok = false;

    if (cublasLtMatmulDescCreate(&opDesc, CUBLAS_COMPUTE_32F, CUDA_R_32F)
            != CUBLAS_STATUS_SUCCESS) return false;
    cublasOperation_t opT = CUBLAS_OP_T, opN = CUBLAS_OP_N;
    cublasLtMatmulDescSetAttribute(opDesc, CUBLASLT_MATMUL_DESC_TRANSA,
                                   &opT, sizeof(opT));
    cublasLtMatmulDescSetAttribute(opDesc, CUBLASLT_MATMUL_DESC_TRANSB,
                                   &opN, sizeof(opN));
    if (bias != nullptr) {
        cublasLtEpilogue_t epi = CUBLASLT_EPILOGUE_BIAS;
        cublasLtMatmulDescSetAttribute(opDesc, CUBLASLT_MATMUL_DESC_EPILOGUE,
                                       &epi, sizeof(epi));
        cublasLtMatmulDescSetAttribute(opDesc, CUBLASLT_MATMUL_DESC_BIAS_POINTER,
                                       &bias, sizeof(bias));
    }
    if (cublasLtMatrixLayoutCreate(&Adesc, CUDA_R_16F, K, N, K) != CUBLAS_STATUS_SUCCESS) goto done;
    if (cublasLtMatrixLayoutCreate(&Bdesc, CUDA_R_16F, K, M, K) != CUBLAS_STATUS_SUCCESS) goto done;
    if (cublasLtMatrixLayoutCreate(&Cdesc, CUDA_R_32F, N, M, N) != CUBLAS_STATUS_SUCCESS) goto done;

    {
        float alpha = 1.0f, beta = 0.0f;
        int64_t key = ((int64_t)M << 42) ^ ((int64_t)N << 21) ^ (int64_t)K
                      ^ (bias ? (1LL << 62) : 0);
        auto& cache = plan_cache();
        auto it = cache.find(key);
        if (it == cache.end()) {
            LtPlan plan{};
            plan.valid = false;
            plan.workspace = 0;
            cublasLtMatmulPreference_t pref = nullptr;
            if (cublasLtMatmulPreferenceCreate(&pref) == CUBLAS_STATUS_SUCCESS) {
                size_t maxWs = 64ull * 1024 * 1024;
                cublasLtMatmulPreferenceSetAttribute(
                    pref, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES,
                    &maxWs, sizeof(maxWs));
                cublasLtMatmulHeuristicResult_t hres[8];
                int found = 0;
                cublasLtMatmulAlgoGetHeuristic(lt, opDesc, Adesc, Bdesc, Cdesc,
                                               Cdesc, pref, 8, hres, &found);
                float bestMs = 1e30f;
                for (int i = 0; i < found; ++i) {
                    if (hres[i].state != CUBLAS_STATUS_SUCCESS) continue;
                    size_t ws = hres[i].workspaceSize;
                    if (ws > maxWs) continue;
                    void* wsp = get_workspace(ws, opt);
                    cublasStatus_t st = cublasLtMatmul(
                        lt, opDesc, &alpha, w, Adesc, x, Bdesc, &beta,
                        out, Cdesc, out, Cdesc, &hres[i].algo, wsp, ws, stream);
                    if (st != CUBLAS_STATUS_SUCCESS) continue;
                    cudaEvent_t e0, e1;
                    cudaEventCreate(&e0);
                    cudaEventCreate(&e1);
                    cudaEventRecord(e0, stream);
                    const int reps = 5;
                    for (int r = 0; r < reps; ++r) {
                        cublasLtMatmul(lt, opDesc, &alpha, w, Adesc, x, Bdesc,
                                       &beta, out, Cdesc, out, Cdesc,
                                       &hres[i].algo, wsp, ws, stream);
                    }
                    cudaEventRecord(e1, stream);
                    cudaEventSynchronize(e1);
                    float ms = 0.f;
                    cudaEventElapsedTime(&ms, e0, e1);
                    cudaEventDestroy(e0);
                    cudaEventDestroy(e1);
                    if (ms < bestMs) {
                        bestMs = ms;
                        plan.algo = hres[i].algo;
                        plan.workspace = ws;
                        plan.valid = true;
                    }
                }
                cublasLtMatmulPreferenceDestroy(pref);
            }
            it = cache.emplace(key, plan).first;
        }
        if (it->second.valid) {
            void* wsp = get_workspace(it->second.workspace, opt);
            cublasStatus_t st = cublasLtMatmul(
                lt, opDesc, &alpha, w, Adesc, x, Bdesc, &beta, out, Cdesc,
                out, Cdesc, &it->second.algo, wsp, it->second.workspace, stream);
            ok = (st == CUBLAS_STATUS_SUCCESS);
        }
    }

done:
    if (Adesc) cublasLtMatrixLayoutDestroy(Adesc);
    if (Bdesc) cublasLtMatrixLayoutDestroy(Bdesc);
    if (Cdesc) cublasLtMatrixLayoutDestroy(Cdesc);
    if (opDesc) cublasLtMatmulDescDestroy(opDesc);
    return ok;
}

// ---------------------------------------------------------------------------
// Fallback: cublasGemmEx FP16 x FP16 -> FP32 (no bias epilogue).
// ---------------------------------------------------------------------------
static void gemm_ex_fp16(const void* x, const void* w, float* out,
                         int64_t M, int64_t N, int64_t K) {
    cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
    float alpha = 1.0f, beta = 0.0f;
    cublasStatus_t st = cublasGemmEx(
        handle, CUBLAS_OP_T, CUBLAS_OP_N,
        (int)N, (int)M, (int)K,
        &alpha,
        w, CUDA_R_16F, (int)K,
        x, CUDA_R_16F, (int)K,
        &beta,
        out, CUDA_R_32F, (int)N,
        CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT);
    TORCH_CHECK(st == CUBLAS_STATUS_SUCCESS,
                "cublasGemmEx failed with status ", (int)st);
}

// ---------------------------------------------------------------------------
// Fused: out = ((x [+ bias] - mean) * rstd + y) * y   (row-wise biased stats)
// Row cached in shared memory; x may be aliased with out (in-place safe).
// ---------------------------------------------------------------------------
template <bool VEC, bool HAS_BIAS>
__global__ void fused_bias_inorm_add_mul_smem_kernel(
    const float* __restrict__ x,
    const float* __restrict__ bias,
    const float* __restrict__ y,
    float* __restrict__ out,
    const int Fdim,
    const float eps)
{
    extern __shared__ float s_row[];
    __shared__ float s_sum[32];
    __shared__ float s_sq[32];

    const int row = blockIdx.x;
    const size_t base = (size_t)row * (size_t)Fdim;
    const float* xr = x + base;
    const float* yr = y + base;
    float* orow = out + base;

    float sum = 0.f;
    float sumsq = 0.f;

    if (VEC) {
        const int F4 = Fdim >> 2;
        const float4* xr4 = reinterpret_cast<const float4*>(xr);
        const float4* b4  = reinterpret_cast<const float4*>(bias);
        float4* s4 = reinterpret_cast<float4*>(s_row);
        #pragma unroll 4
        for (int i = threadIdx.x; i < F4; i += blockDim.x) {
            float4 v = __ldcs(xr4 + i);
            if (HAS_BIAS) {
                const float4 bv = __ldg(b4 + i);
                v.x += bv.x; v.y += bv.y; v.z += bv.z; v.w += bv.w;
            }
            s4[i] = v;
            sum   += v.x + v.y + v.z + v.w;
            sumsq += v.x * v.x + v.y * v.y + v.z * v.z + v.w * v.w;
        }
    } else {
        for (int i = threadIdx.x; i < Fdim; i += blockDim.x) {
            float v = __ldcs(xr + i) + (HAS_BIAS ? __ldg(bias + i) : 0.f);
            s_row[i] = v;
            sum   += v;
            sumsq += v * v;
        }
    }

    const unsigned int mask = 0xffffffffu;
    for (int off = 16; off > 0; off >>= 1) {
        sum   += __shfl_down_sync(mask, sum, off);
        sumsq += __shfl_down_sync(mask, sumsq, off);
    }
    const int lane = threadIdx.x & 31;
    const int wid  = threadIdx.x >> 5;
    if (lane == 0) {
        s_sum[wid] = sum;
        s_sq[wid]  = sumsq;
    }
    __syncthreads();

    const int nwarps = (blockDim.x + 31) >> 5;
    if (wid == 0) {
        float a = (lane < nwarps) ? s_sum[lane] : 0.f;
        float b = (lane < nwarps) ? s_sq[lane]  : 0.f;
        for (int off = 16; off > 0; off >>= 1) {
            a += __shfl_down_sync(mask, a, off);
            b += __shfl_down_sync(mask, b, off);
        }
        if (lane == 0) {
            s_sum[0] = a;
            s_sq[0]  = b;
        }
    }
    __syncthreads();

    const float mean = s_sum[0] / (float)Fdim;
    const float var  = s_sq[0] / (float)Fdim - mean * mean;  // biased
    const float inv  = rsqrtf(var + eps);

    if (VEC) {
        const int F4 = Fdim >> 2;
        const float4* s4  = reinterpret_cast<const float4*>(s_row);
        const float4* yr4 = reinterpret_cast<const float4*>(yr);
        float4* o4 = reinterpret_cast<float4*>(orow);
        #pragma unroll 4
        for (int i = threadIdx.x; i < F4; i += blockDim.x) {
            const float4 xv = s4[i];
            const float4 yv = __ldcs(yr4 + i);
            float4 ov;
            ov.x = ((xv.x - mean) * inv + yv.x) * yv.x;
            ov.y = ((xv.y - mean) * inv + yv.y) * yv.y;
            ov.z = ((xv.z - mean) * inv + yv.z) * yv.z;
            ov.w = ((xv.w - mean) * inv + yv.w) * yv.w;
            __stcs(o4 + i, ov);
        }
    } else {
        for (int i = threadIdx.x; i < Fdim; i += blockDim.x) {
            const float yv = __ldcs(yr + i);
            __stcs(orow + i, ((s_row[i] - mean) * inv + yv) * yv);
        }
    }
}

// Fallback for rows too large for shared memory: two passes over global.
template <bool VEC, bool HAS_BIAS>
__global__ void fused_bias_inorm_add_mul_kernel(
    const float* __restrict__ x,
    const float* __restrict__ bias,
    const float* __restrict__ y,
    float* __restrict__ out,
    const int Fdim,
    const float eps)
{
    const int row = blockIdx.x;
    const size_t base = (size_t)row * (size_t)Fdim;
    const float* xr = x + base;
    const float* yr = y + base;
    float* orow = out + base;

    float sum = 0.f;
    float sumsq = 0.f;

    if (VEC) {
        const int F4 = Fdim >> 2;
        const float4* xr4 = reinterpret_cast<const float4*>(xr);
        const float4* b4  = reinterpret_cast<const float4*>(bias);
        for (int i = threadIdx.x; i < F4; i += blockDim.x) {
            float4 v = xr4[i];
            if (HAS_BIAS) {
                const float4 bv = __ldg(b4 + i);
                v.x += bv.x; v.y += bv.y; v.z += bv.z; v.w += bv.w;
            }
            sum   += v.x + v.y + v.z + v.w;
            sumsq += v.x * v.x + v.y * v.y + v.z * v.z + v.w * v.w;
        }
    } else {
        for (int i = threadIdx.x; i < Fdim; i += blockDim.x) {
            float v = xr[i] + (HAS_BIAS ? __ldg(bias + i) : 0.f);
            sum   += v;
            sumsq += v * v;
        }
    }

    __shared__ float s_sum[32];
    __shared__ float s_sq[32];
    const unsigned int mask = 0xffffffffu;
    for (int off = 16; off > 0; off >>= 1) {
        sum   += __shfl_down_sync(mask, sum, off);
        sumsq += __shfl_down_sync(mask, sumsq, off);
    }
    const int lane = threadIdx.x & 31;
    const int wid  = threadIdx.x >> 5;
    if (lane == 0) {
        s_sum[wid] = sum;
        s_sq[wid]  = sumsq;
    }
    __syncthreads();

    const int nwarps = (blockDim.x + 31) >> 5;
    if (wid == 0) {
        float a = (lane < nwarps) ? s_sum[lane] : 0.f;
        float b = (lane < nwarps) ? s_sq[lane]  : 0.f;
        for (int off = 16; off > 0; off >>= 1) {
            a += __shfl_down_sync(mask, a, off);
            b += __shfl_down_sync(mask, b, off);
        }
        if (lane == 0) {
            s_sum[0] = a;
            s_sq[0]  = b;
        }
    }
    __syncthreads();

    const float mean = s_sum[0] / (float)Fdim;
    const float var  = s_sq[0] / (float)Fdim - mean * mean;
    const float inv  = rsqrtf(var + eps);

    if (VEC) {
        const int F4 = Fdim >> 2;
        const float4* xr4 = reinterpret_cast<const float4*>(xr);
        const float4* b4  = reinterpret_cast<const float4*>(bias);
        const float4* yr4 = reinterpret_cast<const float4*>(yr);
        float4* o4 = reinterpret_cast<float4*>(orow);
        for (int i = threadIdx.x; i < F4; i += blockDim.x) {
            float4 xv = xr4[i];
            if (HAS_BIAS) {
                const float4 bv = __ldg(b4 + i);
                xv.x += bv.x; xv.y += bv.y; xv.z += bv.z; xv.w += bv.w;
            }
            const float4 yv = yr4[i];
            float4 ov;
            ov.x = ((xv.x - mean) * inv + yv.x) * yv.x;
            ov.y = ((xv.y - mean) * inv + yv.y) * yv.y;
            ov.z = ((xv.z - mean) * inv + yv.z) * yv.z;
            ov.w = ((xv.w - mean) * inv + yv.w) * yv.w;
            o4[i] = ov;
        }
    } else {
        for (int i = threadIdx.x; i < Fdim; i += blockDim.x) {
            float xv = xr[i] + (HAS_BIAS ? __ldg(bias + i) : 0.f);
            const float yv = yr[i];
            orow[i] = ((xv - mean) * inv + yv) * yv;
        }
    }
}

torch::Tensor forward_fused(torch::Tensor x, torch::Tensor w_half,
                            torch::Tensor bias, torch::Tensor y, double eps) {
    TORCH_CHECK(x.is_cuda() && y.is_cuda() && w_half.is_cuda(),
                "inputs must be CUDA tensors");
    TORCH_CHECK(w_half.scalar_type() == torch::kHalf, "w_half must be float16");
    TORCH_CHECK(x.dim() == 2 && y.dim() == 2, "x and y must be 2-D");

    auto xc = x.contiguous();
    torch::Tensor xh = (xc.scalar_type() == torch::kHalf)
                       ? xc : xc.to(torch::kHalf);
    auto wh = w_half.contiguous();

    const int64_t M = xh.size(0);
    const int64_t K = xh.size(1);
    const int64_t N = wh.size(0);
    TORCH_CHECK(wh.size(1) == K, "weight shape mismatch");

    auto z = torch::empty({M, N}, xc.options().dtype(torch::kFloat32));
    if (M == 0 || N == 0) return z;

    auto yc = y.contiguous();
    TORCH_CHECK(yc.scalar_type() == torch::kFloat32, "y must be float32");
    TORCH_CHECK(yc.sizes() == z.sizes(), "y shape mismatch");

    const bool has_bias = bias.defined() &&
                          bias.scalar_type() == torch::kFloat32 &&
                          bias.numel() == (int64_t)N;
    torch::Tensor bc;
    const float* bptr = nullptr;
    if (has_bias) {
        bc = bias.contiguous();
        bptr = bc.data_ptr<float>();
    }

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    float* zp = z.data_ptr<float>();

    bool bias_done = false;
    if (K > 0) {
        bias_done = gemm_lt_fp16(xh.data_ptr(), wh.data_ptr(), bptr, zp,
                                 M, N, K, xc.options(), stream);
        if (!bias_done) {
            gemm_ex_fp16(xh.data_ptr(), wh.data_ptr(), zp, M, N, K);
        }
    } else {
        z.zero_();
    }
    // If the Lt path handled the bias epilogue, the fused kernel skips it.
    const float* eff_bias = bias_done ? nullptr : bptr;
    const bool kernel_has_bias = (eff_bias != nullptr);

    const int B = (int)z.size(0);
    const int Fdim = (int)z.size(1);

    // in-place: the fused kernel reads z into shared memory before writing.
    auto out = z;
    const bool vec = (Fdim & 3) == 0;
    const size_t smem = (size_t)Fdim * sizeof(float);

    int threads = 256;
    if (Fdim >= 4096) threads = 512;
    else if (Fdim < 512) threads = 128;

    const float* yp = yc.data_ptr<float>();
    float* op = out.data_ptr<float>();
    const float feps = (float)eps;

    if (smem <= 48 * 1024) {
        if (vec) {
            if (kernel_has_bias)
                fused_bias_inorm_add_mul_smem_kernel<true, true><<<B, threads, smem, stream>>>(
                    zp, eff_bias, yp, op, Fdim, feps);
            else
                fused_bias_inorm_add_mul_smem_kernel<true, false><<<B, threads, smem, stream>>>(
                    zp, nullptr, yp, op, Fdim, feps);
        } else {
            if (kernel_has_bias)
                fused_bias_inorm_add_mul_smem_kernel<false, true><<<B, threads, smem, stream>>>(
                    zp, eff_bias, yp, op, Fdim, feps);
            else
                fused_bias_inorm_add_mul_smem_kernel<false, false><<<B, threads, smem, stream>>>(
                    zp, nullptr, yp, op, Fdim, feps);
        }
    } else {
        if (vec) {
            if (kernel_has_bias)
                fused_bias_inorm_add_mul_kernel<true, true><<<B, threads, 0, stream>>>(
                    zp, eff_bias, yp, op, Fdim, feps);
            else
                fused_bias_inorm_add_mul_kernel<true, false><<<B, threads, 0, stream>>>(
                    zp, nullptr, yp, op, Fdim, feps);
        } else {
            if (kernel_has_bias)
                fused_bias_inorm_add_mul_kernel<false, true><<<B, threads, 0, stream>>>(
                    zp, eff_bias, yp, op, Fdim, feps);
            else
                fused_bias_inorm_add_mul_kernel<false, false><<<B, threads, 0, stream>>>(
                    zp, nullptr, yp, op, Fdim, feps);
        }
    }
    return out;
}
"""

fused_cpp_source = (
    "torch::Tensor forward_fused(torch::Tensor x, torch::Tensor w_half, "
    "torch::Tensor bias, torch::Tensor y, double eps);"
)

fused_ext = load_inline(
    name="fused_linear_inorm_fp16_v4",
    cpp_sources=fused_cpp_source,
    cuda_sources=fused_cuda_source,
    functions=["forward_fused"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
    extra_ldflags=["-lcublas", "-lcublasLt"],
)


class ModelNew(nn.Module):
    """
    Optimized model:
      * Linear GEMM runs on FP16 tensor cores with FP32 accumulation via an
        autotuned cuBLASLt call with a fused BIAS epilogue (~2x the throughput
        of TF32, same 10-bit input mantissa rounding). The FP16 copy of the
        weight is cached and rebuilt only if the weight tensor changes.
      * Instance normalization + residual add + elementwise multiply are fused
        into a single custom CUDA kernel with the row cached in shared memory
        and streaming loads/stores.
    """
    def __init__(self, in_features, out_features, eps=1e-5, momentum=0.1):
        super(ModelNew, self).__init__()
        _enable_tf32()
        self.bmm = nn.Linear(in_features, out_features)
        self.eps = float(eps)
        self._wh = None
        self._wh_key = None
        self._empty_bias = None

    def _half_weight(self, w):
        key = (w._version, w.data_ptr(), w.device)
        if self._wh is None or self._wh_key != key:
            with torch.no_grad():
                self._wh = w.detach().to(torch.float16).contiguous()
            self._wh_key = key
        return self._wh

    def forward(self, x, y):
        w = self.bmm.weight
        b = self.bmm.bias
        if (x.is_cuda and y.is_cuda and x.dim() == 2 and y.dim() == 2
                and x.dtype == torch.float32 and y.dtype == torch.float32
                and w.is_cuda and w.dtype == torch.float32):
            wh = self._half_weight(w)
            if b is not None:
                return fused_ext.forward_fused(x, wh, b, y, self.eps)
            if self._empty_bias is None or self._empty_bias.device != x.device:
                self._empty_bias = torch.empty(0, device=x.device)
            return fused_ext.forward_fused(x, wh, self._empty_bias, y, self.eps)
        # fallback path (CPU / other dtypes) replicating the original ops
        z = F.linear(x, w, b)
        mean = z.mean(dim=1, keepdim=True)
        var = z.var(dim=1, unbiased=False, keepdim=True)
        z = (z - mean) / torch.sqrt(var + self.eps)
        return (z + y) * y
```