import torch
from torch._inductor.select_algorithm import extern_kernels
import triton
import triton.language as tl
from torch._inductor.runtime.triton_heuristics import grid
from torch._C import _cuda_getCurrentRawStream as get_raw_stream
from torch._inductor.runtime import triton_helpers
from torch._inductor.runtime.triton_helpers import libdevice, math as tl_math
import math
import torch.nn as nn
import torch.cuda
import torch.distributed
assert_size_stride = torch._C._dynamo.guards.assert_size_stride
empty_strided_cuda = torch._C._dynamo.guards._empty_strided_cuda
reinterpret_tensor = torch._C._dynamo.guards._reinterpret_tensor


@triton.jit
def triton_poi_fused_native_layer_norm_0(in_ptr0, out_ptr0, out_ptr1,
    xnumel, XBLOCK: tl.constexpr):
    xnumel = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex
    tmp0 = tl.load(in_ptr0 + 4 * x0, xmask, eviction_policy='evict_last')
    tmp1 = tl.load(in_ptr0 + (1 + 4 * x0), xmask, eviction_policy='evict_last')
    tmp3 = tl.load(in_ptr0 + (2 + 4 * x0), xmask, eviction_policy='evict_last')
    tmp5 = tl.load(in_ptr0 + (3 + 4 * x0), xmask, eviction_policy='evict_last')
    tmp2 = tmp0 + tmp1
    tmp4 = tmp2 + tmp3
    tmp6 = tmp4 + tmp5
    tmp7 = 4.0
    tmp8 = tmp6 / tmp7
    tmp9 = tmp0 - tmp8
    tmp10 = tmp9 * tmp9
    tmp11 = tmp1 - tmp8
    tmp12 = tmp11 * tmp11
    tmp13 = tmp10 + tmp12
    tmp14 = tmp3 - tmp8
    tmp15 = tmp14 * tmp14
    tmp16 = tmp13 + tmp15
    tmp17 = tmp5 - tmp8
    tmp18 = tmp17 * tmp17
    tmp19 = tmp16 + tmp18
    tmp20 = tmp19 / tmp7
    tmp21 = 1e-05
    tmp22 = tmp20 + tmp21
    tmp23 = libdevice.rsqrt(tmp22)
    tl.store(out_ptr0 + x0, tmp8, xmask)
    tl.store(out_ptr1 + x0, tmp23, xmask)


@triton.jit
def triton_poi_fused_native_layer_norm_1(in_ptr0, in_ptr1, in_ptr2, in_ptr3,
    in_ptr4, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 64
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x2 = xindex
    x1 = xindex // 4
    x0 = xindex % 4
    tmp0 = tl.load(in_ptr0 + x2, xmask)
    tmp1 = tl.load(in_ptr1 + x1, xmask, eviction_policy='evict_last')
    tmp3 = tl.load(in_ptr2 + x1, xmask, eviction_policy='evict_last')
    tmp5 = tl.load(in_ptr3 + x0, xmask, eviction_policy='evict_last')
    tmp7 = tl.load(in_ptr4 + x0, xmask, eviction_policy='evict_last')
    tmp2 = tmp0 - tmp1
    tmp4 = tmp2 * tmp3
    tmp6 = tmp4 * tmp5
    tmp8 = tmp6 + tmp7
    tl.store(out_ptr0 + x2, tmp8, xmask)


@triton.jit
def triton_poi_fused_clone_div_2(in_ptr0, in_ptr1, out_ptr0, ynumel, xnumel,
    YBLOCK: tl.constexpr, XBLOCK: tl.constexpr):
    ynumel = 16
    xnumel = 4
    yoffset = tl.program_id(1) * YBLOCK
    yindex = yoffset + tl.arange(0, YBLOCK)[None, :]
    ymask = yindex < ynumel
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:, None]
    xmask = xindex < xnumel
    x2 = xindex
    y0 = yindex % 4
    y1 = yindex // 4
    y3 = yindex
    tmp0 = tl.load(in_ptr0 + (y0 + 4 * x2 + 16 * y1), xmask & ymask,
        eviction_policy='evict_last')
    tmp1 = tl.load(in_ptr1 + y0, ymask, eviction_policy='evict_last')
    tmp2 = tmp0 + tmp1
    tmp3 = 1.0
    tmp4 = tmp2 * tmp3
    tl.store(out_ptr0 + (x2 + 4 * y3), tmp4, xmask & ymask)


@triton.jit
def triton_poi_fused_clone_3(in_ptr0, in_ptr1, out_ptr0, ynumel, xnumel,
    YBLOCK: tl.constexpr, XBLOCK: tl.constexpr):
    ynumel = 16
    xnumel = 4
    yoffset = tl.program_id(1) * YBLOCK
    yindex = yoffset + tl.arange(0, YBLOCK)[None, :]
    ymask = yindex < ynumel
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:, None]
    xmask = xindex < xnumel
    x2 = xindex
    y0 = yindex % 4
    y1 = yindex // 4
    y3 = yindex
    tmp0 = tl.load(in_ptr0 + (y0 + 4 * x2 + 16 * y1), xmask & ymask,
        eviction_policy='evict_last')
    tmp1 = tl.load(in_ptr1 + y0, ymask, eviction_policy='evict_last')
    tmp2 = tmp0 + tmp1
    tl.store(out_ptr0 + (x2 + 4 * y3), tmp2, xmask & ymask)


@triton.jit
def triton_poi_fused_add_gt_4(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 64
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x3 = xindex
    x0 = xindex % 4
    x1 = xindex // 4 % 4
    tmp0 = tl.load(in_ptr0 + x3, xmask)
    tmp1 = x0 + -1 * x1
    tmp2 = tl.full([1], 1, tl.int64)
    tmp3 = tmp1 >= tmp2
    tmp4 = tl.full([1], 1, tl.uint8)
    tmp5 = tl.full([1], 0, tl.uint8)
    tmp6 = tl.where(tmp3, tmp4, tmp5)
    tmp7 = tmp6.to(tl.float32)
    tmp8 = tmp0 + tmp7
    tmp9 = 0.0
    tmp10 = tmp8 > tmp9
    tl.store(out_ptr0 + x3, tmp10, xmask)


@triton.jit
def triton_poi_fused__softmax_masked_fill_5(in_ptr0, in_ptr1, out_ptr0,
    out_ptr1, xnumel, XBLOCK: tl.constexpr):
    xnumel = 64
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex % 4
    x2 = xindex // 16
    x3 = xindex
    tmp0 = tl.load(in_ptr0 + (4 * x0 + 16 * x2), xmask, eviction_policy=
        'evict_last').to(tl.int1)
    tmp1 = tl.load(in_ptr1 + 4 * x3, xmask, eviction_policy='evict_last')
    tmp4 = tl.load(in_ptr0 + (1 + 4 * x0 + 16 * x2), xmask, eviction_policy
        ='evict_last').to(tl.int1)
    tmp5 = tl.load(in_ptr1 + (1 + 4 * x3), xmask, eviction_policy='evict_last')
    tmp8 = tl.load(in_ptr0 + (2 + 4 * x0 + 16 * x2), xmask, eviction_policy
        ='evict_last').to(tl.int1)
    tmp9 = tl.load(in_ptr1 + (2 + 4 * x3), xmask, eviction_policy='evict_last')
    tmp12 = tl.load(in_ptr0 + (3 + 4 * x0 + 16 * x2), xmask,
        eviction_policy='evict_last').to(tl.int1)
    tmp13 = tl.load(in_ptr1 + (3 + 4 * x3), xmask, eviction_policy='evict_last'
        )
    tmp2 = -9.999999843067494e+17
    tmp3 = tl.where(tmp0, tmp2, tmp1)
    tmp6 = tl.where(tmp4, tmp2, tmp5)
    tmp7 = triton_helpers.maximum(tmp3, tmp6)
    tmp10 = tl.where(tmp8, tmp2, tmp9)
    tmp11 = triton_helpers.maximum(tmp7, tmp10)
    tmp14 = tl.where(tmp12, tmp2, tmp13)
    tmp15 = triton_helpers.maximum(tmp11, tmp14)
    tmp16 = tmp3 - tmp15
    tmp17 = tl_math.exp(tmp16)
    tmp18 = tmp6 - tmp15
    tmp19 = tl_math.exp(tmp18)
    tmp20 = tmp17 + tmp19
    tmp21 = tmp10 - tmp15
    tmp22 = tl_math.exp(tmp21)
    tmp23 = tmp20 + tmp22
    tmp24 = tmp14 - tmp15
    tmp25 = tl_math.exp(tmp24)
    tmp26 = tmp23 + tmp25
    tl.store(out_ptr0 + x3, tmp15, xmask)
    tl.store(out_ptr1 + x3, tmp26, xmask)


@triton.jit
def triton_poi_fused__softmax_masked_fill_6(in_out_ptr0, in_ptr0, in_ptr1,
    in_ptr2, xnumel, XBLOCK: tl.constexpr):
    xnumel = 256
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x3 = xindex // 64
    x4 = xindex % 16
    x5 = xindex
    x6 = xindex // 4
    tmp0 = tl.load(in_ptr0 + (x4 + 16 * x3), xmask, eviction_policy=
        'evict_last').to(tl.int1)
    tmp1 = tl.load(in_out_ptr0 + x5, xmask)
    tmp4 = tl.load(in_ptr1 + x6, xmask, eviction_policy='evict_last')
    tmp7 = tl.load(in_ptr2 + x6, xmask, eviction_policy='evict_last')
    tmp2 = -9.999999843067494e+17
    tmp3 = tl.where(tmp0, tmp2, tmp1)
    tmp5 = tmp3 - tmp4
    tmp6 = tl_math.exp(tmp5)
    tmp8 = tmp6 / tmp7
    tl.store(in_out_ptr0 + x5, tmp8, xmask)


@triton.jit
def triton_poi_fused_clone_7(in_ptr0, out_ptr0, ynumel, xnumel, YBLOCK: tl.
    constexpr, XBLOCK: tl.constexpr):
    ynumel = 16
    xnumel = 4
    yoffset = tl.program_id(1) * YBLOCK
    yindex = yoffset + tl.arange(0, YBLOCK)[None, :]
    ymask = yindex < ynumel
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:, None]
    xmask = xindex < xnumel
    x2 = xindex
    y0 = yindex % 4
    y1 = yindex // 4
    y3 = yindex
    tmp0 = tl.load(in_ptr0 + (y0 + 4 * x2 + 16 * y1), xmask & ymask,
        eviction_policy='evict_last')
    tl.store(out_ptr0 + (x2 + 4 * y3), tmp0, xmask & ymask)


@triton.jit
def triton_poi_fused_add_native_layer_norm_8(in_ptr0, in_ptr1, out_ptr0,
    out_ptr1, xnumel, XBLOCK: tl.constexpr):
    xnumel = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex
    tmp0 = tl.load(in_ptr0 + 4 * x0, xmask, eviction_policy='evict_last')
    tmp1 = tl.load(in_ptr1 + 4 * x0, xmask, eviction_policy='evict_last')
    tmp3 = tl.load(in_ptr0 + (1 + 4 * x0), xmask, eviction_policy='evict_last')
    tmp4 = tl.load(in_ptr1 + (1 + 4 * x0), xmask, eviction_policy='evict_last')
    tmp7 = tl.load(in_ptr0 + (2 + 4 * x0), xmask, eviction_policy='evict_last')
    tmp8 = tl.load(in_ptr1 + (2 + 4 * x0), xmask, eviction_policy='evict_last')
    tmp11 = tl.load(in_ptr0 + (3 + 4 * x0), xmask, eviction_policy='evict_last'
        )
    tmp12 = tl.load(in_ptr1 + (3 + 4 * x0), xmask, eviction_policy='evict_last'
        )
    tmp2 = tmp0 + tmp1
    tmp5 = tmp3 + tmp4
    tmp6 = tmp2 + tmp5
    tmp9 = tmp7 + tmp8
    tmp10 = tmp6 + tmp9
    tmp13 = tmp11 + tmp12
    tmp14 = tmp10 + tmp13
    tmp15 = 4.0
    tmp16 = tmp14 / tmp15
    tmp17 = tmp2 - tmp16
    tmp18 = tmp17 * tmp17
    tmp19 = tmp5 - tmp16
    tmp20 = tmp19 * tmp19
    tmp21 = tmp18 + tmp20
    tmp22 = tmp9 - tmp16
    tmp23 = tmp22 * tmp22
    tmp24 = tmp21 + tmp23
    tmp25 = tmp13 - tmp16
    tmp26 = tmp25 * tmp25
    tmp27 = tmp24 + tmp26
    tmp28 = tmp27 / tmp15
    tl.store(out_ptr0 + x0, tmp16, xmask)
    tl.store(out_ptr1 + x0, tmp28, xmask)


@triton.jit
def triton_poi_fused_add_native_layer_norm_9(in_ptr0, in_ptr1, in_ptr2,
    in_ptr3, in_ptr4, in_ptr5, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 64
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x2 = xindex
    x1 = xindex // 4
    x0 = xindex % 4
    tmp0 = tl.load(in_ptr0 + x2, xmask)
    tmp1 = tl.load(in_ptr1 + x2, xmask)
    tmp3 = tl.load(in_ptr2 + x1, xmask, eviction_policy='evict_last')
    tmp5 = tl.load(in_ptr3 + x1, xmask, eviction_policy='evict_last')
    tmp10 = tl.load(in_ptr4 + x0, xmask, eviction_policy='evict_last')
    tmp12 = tl.load(in_ptr5 + x0, xmask, eviction_policy='evict_last')
    tmp2 = tmp0 + tmp1
    tmp4 = tmp2 - tmp3
    tmp6 = 1e-05
    tmp7 = tmp5 + tmp6
    tmp8 = libdevice.rsqrt(tmp7)
    tmp9 = tmp4 * tmp8
    tmp11 = tmp9 * tmp10
    tmp13 = tmp11 + tmp12
    tl.store(out_ptr0 + x2, tmp13, xmask)


@triton.jit
def triton_poi_fused_add_mul_pow_tanh_10(in_ptr0, out_ptr0, xnumel, XBLOCK:
    tl.constexpr):
    xnumel = 256
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex
    tmp0 = tl.load(in_ptr0 + x0, xmask)
    tmp1 = 0.5
    tmp2 = tmp0 * tmp1
    tmp3 = tmp0 * tmp0
    tmp4 = tmp3 * tmp0
    tmp5 = 0.044715
    tmp6 = tmp4 * tmp5
    tmp7 = tmp0 + tmp6
    tmp8 = 0.7978845608028654
    tmp9 = tmp7 * tmp8
    tmp10 = libdevice.tanh(tmp9)
    tmp11 = 1.0
    tmp12 = tmp10 + tmp11
    tmp13 = tmp2 * tmp12
    tl.store(out_ptr0 + x0, tmp13, xmask)


@triton.jit
def triton_poi_fused_add_11(in_out_ptr0, in_ptr0, in_ptr1, in_ptr2, xnumel,
    XBLOCK: tl.constexpr):
    xnumel = 64
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x2 = xindex
    x0 = xindex % 4
    tmp0 = tl.load(in_out_ptr0 + x2, xmask)
    tmp1 = tl.load(in_ptr0 + x0, xmask, eviction_policy='evict_last')
    tmp3 = tl.load(in_ptr1 + x2, xmask)
    tmp4 = tl.load(in_ptr2 + x2, xmask)
    tmp2 = tmp0 + tmp1
    tmp5 = tmp3 + tmp4
    tmp6 = tmp2 + tmp5
    tl.store(in_out_ptr0 + x2, tmp6, xmask)


def call(args):
    (primals_1, primals_2, primals_3, primals_4, primals_5, primals_6,
        primals_7, primals_8, primals_9, primals_10, primals_11, primals_12,
        primals_13, primals_14, primals_15, primals_16, primals_17, primals_18
        ) = args
    args.clear()
    assert_size_stride(primals_1, (4, 4, 4), (16, 4, 1))
    assert_size_stride(primals_2, (4,), (1,))
    assert_size_stride(primals_3, (4,), (1,))
    assert_size_stride(primals_4, (4, 4, 4), (16, 4, 1))
    assert_size_stride(primals_5, (4, 4), (4, 1))
    assert_size_stride(primals_6, (4,), (1,))
    assert_size_stride(primals_7, (4, 4), (4, 1))
    assert_size_stride(primals_8, (4,), (1,))
    assert_size_stride(primals_9, (4, 4), (4, 1))
    assert_size_stride(primals_10, (4,), (1,))
    assert_size_stride(primals_11, (4, 4), (4, 1))
    assert_size_stride(primals_12, (4,), (1,))
    assert_size_stride(primals_13, (4,), (1,))
    assert_size_stride(primals_14, (4,), (1,))
    assert_size_stride(primals_15, (16, 4), (4, 1))
    assert_size_stride(primals_16, (16,), (1,))
    assert_size_stride(primals_17, (4, 16), (16, 1))
    assert_size_stride(primals_18, (4,), (1,))
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        buf0 = empty_strided_cuda((4, 4, 1), (4, 1, 16), torch.float32)
        buf1 = empty_strided_cuda((4, 4, 1), (4, 1, 16), torch.float32)
        get_raw_stream(0)
        triton_poi_fused_native_layer_norm_0[grid(16)](primals_4, buf0,
            buf1, 16, XBLOCK=16, num_warps=1, num_stages=1)
        buf2 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        triton_poi_fused_native_layer_norm_1[grid(64)](primals_4, buf0,
            buf1, primals_2, primals_3, buf2, 64, XBLOCK=64, num_warps=1,
            num_stages=1)
        del primals_2
        del primals_3
        buf3 = empty_strided_cuda((16, 4), (4, 1), torch.float32)
        extern_kernels.mm(reinterpret_tensor(buf2, (16, 4), (4, 1), 0),
            reinterpret_tensor(primals_5, (4, 4), (1, 4), 0), out=buf3)
        buf4 = empty_strided_cuda((16, 4), (4, 1), torch.float32)
        extern_kernels.mm(reinterpret_tensor(buf2, (16, 4), (4, 1), 0),
            reinterpret_tensor(primals_7, (4, 4), (1, 4), 0), out=buf4)
        buf5 = empty_strided_cuda((16, 4), (4, 1), torch.float32)
        extern_kernels.mm(reinterpret_tensor(buf2, (16, 4), (4, 1), 0),
            reinterpret_tensor(primals_9, (4, 4), (1, 4), 0), out=buf5)
        buf6 = empty_strided_cuda((4, 4, 4, 1), (16, 4, 1, 1), torch.float32)
        triton_poi_fused_clone_div_2[grid(16, 4)](buf5, primals_10, buf6, 
            16, 4, XBLOCK=4, YBLOCK=8, num_warps=1, num_stages=1)
        del primals_10
        buf7 = reinterpret_tensor(buf5, (4, 4, 1, 4), (16, 4, 4, 1), 0)
        del buf5
        triton_poi_fused_clone_3[grid(16, 4)](buf3, primals_6, buf7, 16, 4,
            XBLOCK=4, YBLOCK=16, num_warps=1, num_stages=1)
        del primals_6
        buf8 = empty_strided_cuda((16, 4, 4), (16, 4, 1), torch.float32)
        extern_kernels.bmm(reinterpret_tensor(buf6, (16, 4, 1), (4, 1, 0), 
            0), reinterpret_tensor(buf7, (16, 1, 4), (4, 0, 1), 0), out=buf8)
        buf9 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.bool)
        triton_poi_fused_add_gt_4[grid(64)](primals_1, buf9, 64, XBLOCK=64,
            num_warps=1, num_stages=1)
        del primals_1
        buf10 = reinterpret_tensor(buf3, (4, 4, 4, 1), (16, 4, 1, 64), 0)
        del buf3
        buf11 = empty_strided_cuda((4, 4, 4, 1), (16, 4, 1, 64), torch.float32)
        triton_poi_fused__softmax_masked_fill_5[grid(64)](buf9, buf8, buf10,
            buf11, 64, XBLOCK=64, num_warps=1, num_stages=1)
        buf12 = reinterpret_tensor(buf8, (4, 4, 4, 4), (64, 16, 4, 1), 0)
        del buf8
        triton_poi_fused__softmax_masked_fill_6[grid(256)](buf12, buf9,
            buf10, buf11, 256, XBLOCK=256, num_warps=4, num_stages=1)
        buf13 = reinterpret_tensor(buf11, (4, 4, 4, 1), (16, 4, 1, 1), 0)
        del buf11
        triton_poi_fused_clone_3[grid(16, 4)](buf4, primals_8, buf13, 16, 4,
            XBLOCK=4, YBLOCK=16, num_warps=1, num_stages=1)
        del primals_8
        buf14 = reinterpret_tensor(buf4, (16, 4, 1), (4, 1, 1), 0)
        del buf4
        extern_kernels.bmm(reinterpret_tensor(buf12, (16, 4, 4), (16, 4, 1),
            0), reinterpret_tensor(buf13, (16, 4, 1), (4, 1, 0), 0), out=buf14)
        buf15 = reinterpret_tensor(buf10, (4, 4, 4, 1), (16, 4, 1, 1), 0)
        del buf10
        triton_poi_fused_clone_7[grid(16, 4)](buf14, buf15, 16, 4, XBLOCK=4,
            YBLOCK=16, num_warps=1, num_stages=1)
        buf16 = reinterpret_tensor(buf14, (16, 4), (4, 1), 0)
        del buf14
        extern_kernels.addmm(primals_12, reinterpret_tensor(buf15, (16, 4),
            (4, 1), 0), reinterpret_tensor(primals_11, (4, 4), (1, 4), 0),
            alpha=1, beta=1, out=buf16)
        del primals_12
        buf17 = buf1
        del buf1
        buf18 = buf0
        del buf0
        triton_poi_fused_add_native_layer_norm_8[grid(16)](buf16, primals_4,
            buf17, buf18, 16, XBLOCK=16, num_warps=1, num_stages=1)
        buf19 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        triton_poi_fused_add_native_layer_norm_9[grid(64)](buf16, primals_4,
            buf17, buf18, primals_13, primals_14, buf19, 64, XBLOCK=64,
            num_warps=1, num_stages=1)
        del buf17
        del buf18
        del primals_14
        buf20 = empty_strided_cuda((16, 16), (16, 1), torch.float32)
        extern_kernels.addmm(primals_16, reinterpret_tensor(buf19, (16, 4),
            (4, 1), 0), reinterpret_tensor(primals_15, (4, 16), (1, 4), 0),
            alpha=1, beta=1, out=buf20)
        del primals_16
        buf21 = empty_strided_cuda((4, 4, 16), (64, 16, 1), torch.float32)
        triton_poi_fused_add_mul_pow_tanh_10[grid(256)](buf20, buf21, 256,
            XBLOCK=256, num_warps=4, num_stages=1)
        buf22 = empty_strided_cuda((16, 4), (4, 1), torch.float32)
        extern_kernels.mm(reinterpret_tensor(buf21, (16, 16), (16, 1), 0),
            reinterpret_tensor(primals_17, (16, 4), (1, 16), 0), out=buf22)
        buf23 = reinterpret_tensor(buf22, (4, 4, 4), (16, 4, 1), 0)
        del buf22
        triton_poi_fused_add_11[grid(64)](buf23, primals_18, buf16,
            primals_4, 64, XBLOCK=64, num_warps=1, num_stages=1)
        del primals_18
    return buf23, primals_4, primals_13, reinterpret_tensor(buf2, (16, 4),
        (4, 1), 0), reinterpret_tensor(buf9, (4, 1, 4, 4), (16, 16, 4, 1), 0
        ), buf12, reinterpret_tensor(buf15, (16, 4), (4, 1), 0
        ), buf16, reinterpret_tensor(buf19, (16, 4), (4, 1), 0
        ), buf20, reinterpret_tensor(buf21, (16, 16), (16, 1), 0
        ), primals_17, primals_15, primals_11, reinterpret_tensor(buf13, (
        16, 1, 4), (4, 1, 1), 0), reinterpret_tensor(buf6, (16, 1, 4), (4, 
        1, 1), 0), reinterpret_tensor(buf7, (16, 4, 1), (4, 1, 4), 0
        ), primals_9, primals_7, primals_5


def gelu(x):
    return 0.5 * x * (1 + torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 *
        torch.pow(x, 3))))


def generate_relative_positions_matrix(length, max_relative_positions,
    cache=False):
    """Generate the clipped relative positions matrix
       for a given length and maximum relative positions"""
    if cache:
        distance_mat = torch.arange(-length + 1, 1, 1).unsqueeze(0)
    else:
        range_vec = torch.arange(length)
        range_mat = range_vec.unsqueeze(-1).expand(-1, length).transpose(0, 1)
        distance_mat = range_mat - range_mat.transpose(0, 1)
    distance_mat_clipped = torch.clamp(distance_mat, min=-
        max_relative_positions, max=max_relative_positions)
    final_mat = distance_mat_clipped + max_relative_positions
    return final_mat


def relative_matmul(x, z, transpose):
    """Helper function for relative positions attention."""
    batch_size = x.shape[0]
    heads = x.shape[1]
    length = x.shape[2]
    x_t = x.permute(2, 0, 1, 3)
    x_t_r = x_t.reshape(length, heads * batch_size, -1)
    if transpose:
        z_t = z.transpose(1, 2)
        x_tz_matmul = torch.matmul(x_t_r, z_t)
    else:
        x_tz_matmul = torch.matmul(x_t_r, z)
    x_tz_matmul_r = x_tz_matmul.reshape(length, batch_size, heads, -1)
    x_tz_matmul_r_t = x_tz_matmul_r.permute(1, 2, 0, 3)
    return x_tz_matmul_r_t


class MLP(nn.Module):

    def __init__(self, n_embd, n_state, dropout):
        super(MLP, self).__init__()
        self.c_fc = nn.Linear(n_embd, n_state)
        self.c_proj = nn.Linear(n_state, n_embd)
        self.act = gelu
        self.dropout_1 = nn.Dropout(dropout)
        self.dropout_2 = nn.Dropout(dropout)
        self.reset_parameters()

    def reset_parameters(self):
        self.c_fc.weight.data.normal_(std=0.02)
        self.c_fc.bias.data.zero_()
        self.c_proj.weight.data.normal_(std=0.02)
        self.c_proj.bias.data.zero_()

    def forward(self, x):
        """
            x is input, [T, B, n_state]
        """
        h = self.dropout_1(self.act(self.c_fc(x)))
        h2 = self.dropout_2(self.c_proj(h))
        return h2


class MultiHeadedAttention(nn.Module):
    """Multi-Head Attention module from "Attention is All You Need"
    :cite:`DBLP:journals/corr/VaswaniSPUJGKP17`.

    Similar to standard `dot` attention but uses
    multiple attention distributions simulataneously
    to select relevant items.

    .. mermaid::

       graph BT
          A[key]
          B[value]
          C[query]
          O[output]
          subgraph Attn
            D[Attn 1]
            E[Attn 2]
            F[Attn N]
          end
          A --> D
          C --> D
          A --> E
          C --> E
          A --> F
          C --> F
          D --> O
          E --> O
          F --> O
          B --> O

    Also includes several additional tricks.

    Args:
       head_count (int): number of parallel heads
       model_dim (int): the dimension of keys/values/queries,
           must be divisible by head_count
       dropout (float): dropout parameter
    """

    def __init__(self, head_count, model_dim, dropout=0.1,
        max_relative_positions=0):
        assert model_dim % head_count == 0
        self.dim_per_head = model_dim // head_count
        self.model_dim = model_dim
        super(MultiHeadedAttention, self).__init__()
        self.head_count = head_count
        self.linear_keys = nn.Linear(model_dim, head_count * self.dim_per_head)
        self.linear_values = nn.Linear(model_dim, head_count * self.
            dim_per_head)
        self.linear_query = nn.Linear(model_dim, head_count * self.dim_per_head
            )
        self.softmax = nn.Softmax(dim=-1)
        self.dropout = nn.Dropout(dropout)
        self.final_linear = nn.Linear(model_dim, model_dim)
        self.max_relative_positions = max_relative_positions
        if max_relative_positions > 0:
            vocab_size = max_relative_positions * 2 + 1
            self.relative_positions_embeddings = nn.Embedding(vocab_size,
                self.dim_per_head)

    def forward(self, key, value, query, mask=None, layer_cache=None, type=None
        ):
        """
        Compute the context vector and the attention vectors.

        Args:
           key (FloatTensor): set of `key_len`
               key vectors ``(batch, key_len, dim)``
           value (FloatTensor): set of `key_len`
               value vectors ``(batch, key_len, dim)``
           query (FloatTensor): set of `query_len`
               query vectors  ``(batch, query_len, dim)``
           mask: binary mask indicating which keys have
               non-zero attention ``(batch, query_len, key_len)``
        Returns:
           (FloatTensor, FloatTensor):

           * output context vectors ``(batch, query_len, dim)``
           * one of the attention vectors ``(batch, query_len, key_len)``
        """
        batch_size = key.size(0)
        dim_per_head = self.dim_per_head
        head_count = self.head_count
        key_len = key.size(1)
        query_len = query.size(1)

        def shape(x):
            """Projection."""
            return x.view(batch_size, -1, head_count, dim_per_head).transpose(
                1, 2)

        def unshape(x):
            """Compute context."""
            return x.transpose(1, 2).contiguous().view(batch_size, -1, 
                head_count * dim_per_head)
        if layer_cache is not None:
            if type == 'self':
                query, key, value = self.linear_query(query), self.linear_keys(
                    query), self.linear_values(query)
                key = shape(key)
                value = shape(value)
                if layer_cache['self_keys'] is not None:
                    key = torch.cat((layer_cache['self_keys'], key), dim=2)
                if layer_cache['self_values'] is not None:
                    value = torch.cat((layer_cache['self_values'], value),
                        dim=2)
                layer_cache['self_keys'] = key
                layer_cache['self_values'] = value
            elif type == 'context':
                query = self.linear_query(query)
                if layer_cache['memory_keys'] is None:
                    key, value = self.linear_keys(key), self.linear_values(
                        value)
                    key = shape(key)
                    value = shape(value)
                else:
                    key, value = layer_cache['memory_keys'], layer_cache[
                        'memory_values']
                layer_cache['memory_keys'] = key
                layer_cache['memory_values'] = value
        else:
            key = self.linear_keys(key)
            value = self.linear_values(value)
            query = self.linear_query(query)
            key = shape(key)
            value = shape(value)
        if self.max_relative_positions > 0 and type == 'self':
            key_len = key.size(2)
            relative_positions_matrix = generate_relative_positions_matrix(
                key_len, self.max_relative_positions, cache=True if 
                layer_cache is not None else False)
            relations_keys = self.relative_positions_embeddings(
                relative_positions_matrix)
            relations_values = self.relative_positions_embeddings(
                relative_positions_matrix)
        query = shape(query)
        key_len = key.size(2)
        query_len = query.size(2)
        query = query / math.sqrt(dim_per_head)
        query_key = torch.matmul(query, key.transpose(2, 3))
        if self.max_relative_positions > 0 and type == 'self':
            scores = query_key + relative_matmul(query, relations_keys, True)
        else:
            scores = query_key
        scores = scores.float()
        if mask is not None:
            mask = mask.unsqueeze(1)
            scores = scores.masked_fill(mask, -1e+18)
        attn = self.softmax(scores)
        drop_attn = self.dropout(attn)
        context_original = torch.matmul(drop_attn, value)
        if self.max_relative_positions > 0 and type == 'self':
            context = unshape(context_original + relative_matmul(drop_attn,
                relations_values, False))
        else:
            context = unshape(context_original)
        output = self.final_linear(context)
        top_attn = attn.view(batch_size, head_count, query_len, key_len)[:,
            0, :, :].contiguous()
        return output, top_attn


class TransformerGPTEncoderLayerNew(nn.Module):
    """
    A single layer of the transformer encoder.

    Args:
        d_model (int): the dimension of keys/values/queries in
                   MultiHeadedAttention, also the input size of
                   the first-layer of the PositionwiseFeedForward.
        heads (int): the number of head for MultiHeadedAttention.
        d_ff (int): the second-layer of the PositionwiseFeedForward.
        dropout (float): dropout probability(0-1.0).
    """

    def __init__(self, d_model, heads, d_ff, dropout, attn_dropout,
        max_relative_positions=0):
        super(TransformerGPTEncoderLayerNew, self).__init__()
        self.self_attn = MultiHeadedAttention(heads, d_model, dropout=
            attn_dropout, max_relative_positions=max_relative_positions)
        self.feed_forward = MLP(d_model, d_model * 4, dropout)
        self.layer_norm_1 = nn.LayerNorm(d_model, eps=1e-05)
        self.layer_norm_2 = nn.LayerNorm(d_model, eps=1e-05)
        self.dropout = nn.Dropout(dropout)

    def forward(self, input_0, input_1):
        primals_5 = self.self_attn.linear_keys.weight
        primals_2 = self.self_attn.linear_keys.bias
        primals_7 = self.self_attn.linear_values.weight
        primals_3 = self.self_attn.linear_values.bias
        primals_9 = self.self_attn.linear_query.weight
        primals_6 = self.self_attn.linear_query.bias
        primals_11 = self.self_attn.final_linear.weight
        primals_8 = self.self_attn.final_linear.bias
        primals_15 = self.feed_forward.c_fc.weight
        primals_16 = self.feed_forward.c_fc.bias
        primals_17 = self.feed_forward.c_proj.weight
        primals_10 = self.feed_forward.c_proj.bias
        primals_12 = self.layer_norm_1.weight
        primals_13 = self.layer_norm_1.bias
        primals_14 = self.layer_norm_2.weight
        primals_18 = self.layer_norm_2.bias
        primals_1 = input_0
        primals_4 = input_1
        output = call([primals_1, primals_2, primals_3, primals_4,
            primals_5, primals_6, primals_7, primals_8, primals_9,
            primals_10, primals_11, primals_12, primals_13, primals_14,
            primals_15, primals_16, primals_17, primals_18])
        return output[0]


Measured latency_ms: 0.026451