import torch
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.triton_helpers import math as tl_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_eye_masked_fill_ne_sum_0(in_ptr0, out_ptr0, xnumel,
    XBLOCK: tl.constexpr):
    xnumel = 4
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex
    tmp7 = tl.load(in_ptr0 + x0, xmask)
    tmp16 = tl.load(in_ptr0 + (4 + x0), xmask)
    tmp25 = tl.load(in_ptr0 + (8 + x0), xmask)
    tmp34 = tl.load(in_ptr0 + (12 + x0), xmask)
    tmp0 = tl.full([1], 0, tl.int64)
    tmp1 = x0
    tmp2 = tmp0 == tmp1
    tmp3 = 1.0
    tmp4 = 0.0
    tmp5 = tl.where(tmp2, tmp3, tmp4)
    tmp6 = tmp5 != tmp4
    tmp8 = tl_math.exp(tmp7)
    tmp9 = 1e-05
    tmp10 = tmp8 + tmp9
    tmp11 = tl.where(tmp6, tmp4, tmp10)
    tmp12 = tl.full([1], 1, tl.int64)
    tmp13 = tmp12 == tmp1
    tmp14 = tl.where(tmp13, tmp3, tmp4)
    tmp15 = tmp14 != tmp4
    tmp17 = tl_math.exp(tmp16)
    tmp18 = tmp17 + tmp9
    tmp19 = tl.where(tmp15, tmp4, tmp18)
    tmp20 = tmp11 + tmp19
    tmp21 = tl.full([1], 2, tl.int64)
    tmp22 = tmp21 == tmp1
    tmp23 = tl.where(tmp22, tmp3, tmp4)
    tmp24 = tmp23 != tmp4
    tmp26 = tl_math.exp(tmp25)
    tmp27 = tmp26 + tmp9
    tmp28 = tl.where(tmp24, tmp4, tmp27)
    tmp29 = tmp20 + tmp28
    tmp30 = tl.full([1], 3, tl.int64)
    tmp31 = tmp30 == tmp1
    tmp32 = tl.where(tmp31, tmp3, tmp4)
    tmp33 = tmp32 != tmp4
    tmp35 = tl_math.exp(tmp34)
    tmp36 = tmp35 + tmp9
    tmp37 = tl.where(tmp33, tmp4, tmp36)
    tmp38 = tmp29 + tmp37
    tl.store(out_ptr0 + x0, tmp38, xmask)


@triton.jit
def triton_poi_fused_add_diag_embed_diagonal_copy_exp_eye_masked_fill_ne_neg_1(
    in_ptr0, in_ptr1, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x1 = xindex // 4
    x0 = xindex % 4
    x2 = xindex
    tmp3 = tl.load(in_ptr0 + 5 * x0, xmask, eviction_policy='evict_last')
    tmp11 = tl.load(in_ptr0 + x2, xmask)
    tmp18 = tl.load(in_ptr1 + x0, xmask, eviction_policy='evict_last')
    tmp0 = x1
    tmp1 = tl.full([1], 0, tl.int32)
    tmp2 = tmp0 == tmp1
    tmp4 = tl_math.exp(tmp3)
    tmp5 = x0
    tmp6 = tmp0 == tmp5
    tmp7 = 1.0
    tmp8 = 0.0
    tmp9 = tl.where(tmp6, tmp7, tmp8)
    tmp10 = tmp9 != tmp8
    tmp12 = tl_math.exp(tmp11)
    tmp13 = 1e-05
    tmp14 = tmp12 + tmp13
    tmp15 = tl.where(tmp10, tmp8, tmp14)
    tmp16 = -tmp15
    tmp17 = tmp5 == tmp0
    tmp19 = tl.where(tmp17, tmp18, tmp8)
    tmp20 = tmp16 + tmp19
    tmp21 = tl.where(tmp2, tmp4, tmp20)
    tl.store(out_ptr0 + x2, tmp21, xmask)


@triton.jit
def triton_poi_fused_eye_masked_fill_ne_sum_2(in_ptr0, out_ptr0, xnumel,
    XBLOCK: tl.constexpr):
    xnumel = 4
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex
    tmp7 = tl.load(in_ptr0 + (16 + x0), xmask)
    tmp16 = tl.load(in_ptr0 + (20 + x0), xmask)
    tmp25 = tl.load(in_ptr0 + (24 + x0), xmask)
    tmp34 = tl.load(in_ptr0 + (28 + x0), xmask)
    tmp0 = tl.full([1], 0, tl.int64)
    tmp1 = x0
    tmp2 = tmp0 == tmp1
    tmp3 = 1.0
    tmp4 = 0.0
    tmp5 = tl.where(tmp2, tmp3, tmp4)
    tmp6 = tmp5 != tmp4
    tmp8 = tl_math.exp(tmp7)
    tmp9 = 1e-05
    tmp10 = tmp8 + tmp9
    tmp11 = tl.where(tmp6, tmp4, tmp10)
    tmp12 = tl.full([1], 1, tl.int64)
    tmp13 = tmp12 == tmp1
    tmp14 = tl.where(tmp13, tmp3, tmp4)
    tmp15 = tmp14 != tmp4
    tmp17 = tl_math.exp(tmp16)
    tmp18 = tmp17 + tmp9
    tmp19 = tl.where(tmp15, tmp4, tmp18)
    tmp20 = tmp11 + tmp19
    tmp21 = tl.full([1], 2, tl.int64)
    tmp22 = tmp21 == tmp1
    tmp23 = tl.where(tmp22, tmp3, tmp4)
    tmp24 = tmp23 != tmp4
    tmp26 = tl_math.exp(tmp25)
    tmp27 = tmp26 + tmp9
    tmp28 = tl.where(tmp24, tmp4, tmp27)
    tmp29 = tmp20 + tmp28
    tmp30 = tl.full([1], 3, tl.int64)
    tmp31 = tmp30 == tmp1
    tmp32 = tl.where(tmp31, tmp3, tmp4)
    tmp33 = tmp32 != tmp4
    tmp35 = tl_math.exp(tmp34)
    tmp36 = tmp35 + tmp9
    tmp37 = tl.where(tmp33, tmp4, tmp36)
    tmp38 = tmp29 + tmp37
    tl.store(out_ptr0 + x0, tmp38, xmask)


@triton.jit
def triton_poi_fused_add_diag_embed_diagonal_copy_exp_eye_masked_fill_ne_neg_3(
    in_ptr0, in_ptr1, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x1 = xindex // 4
    x0 = xindex % 4
    x2 = xindex
    tmp3 = tl.load(in_ptr0 + (16 + 5 * x0), xmask, eviction_policy='evict_last'
        )
    tmp11 = tl.load(in_ptr0 + (16 + x2), xmask)
    tmp18 = tl.load(in_ptr1 + x0, xmask, eviction_policy='evict_last')
    tmp0 = x1
    tmp1 = tl.full([1], 0, tl.int32)
    tmp2 = tmp0 == tmp1
    tmp4 = tl_math.exp(tmp3)
    tmp5 = x0
    tmp6 = tmp0 == tmp5
    tmp7 = 1.0
    tmp8 = 0.0
    tmp9 = tl.where(tmp6, tmp7, tmp8)
    tmp10 = tmp9 != tmp8
    tmp12 = tl_math.exp(tmp11)
    tmp13 = 1e-05
    tmp14 = tmp12 + tmp13
    tmp15 = tl.where(tmp10, tmp8, tmp14)
    tmp16 = -tmp15
    tmp17 = tmp5 == tmp0
    tmp19 = tl.where(tmp17, tmp18, tmp8)
    tmp20 = tmp16 + tmp19
    tmp21 = tl.where(tmp2, tmp4, tmp20)
    tl.store(out_ptr0 + x2, tmp21, xmask)


@triton.jit
def triton_poi_fused_add_diag_embed_4(in_ptr0, in_ptr1, out_ptr0, xnumel,
    XBLOCK: tl.constexpr):
    xnumel = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex % 4
    x2 = xindex
    x1 = xindex // 4
    tmp4 = tl.load(in_ptr0 + x2, xmask)
    tmp6 = tl.load(in_ptr1 + 5 * x0, xmask, eviction_policy='evict_last')
    tmp12 = tl.load(in_ptr1 + x2, xmask)
    tmp18 = tl.load(in_ptr0 + 5 * x0, xmask, eviction_policy='evict_last')
    tmp20 = tl.load(in_ptr1 + x0, xmask, eviction_policy='evict_last')
    tmp0 = tl.full([1], 0, tl.int32)
    tmp1 = tmp0 == tmp0
    tmp2 = x0
    tmp3 = tmp2 == tmp0
    tmp5 = tl_math.exp(tmp4)
    tmp7 = tmp5 * tmp6
    tmp8 = 0.0
    tmp9 = tl.where(tmp3, tmp8, tmp7)
    tmp10 = x1
    tmp11 = tmp10 == tmp0
    tmp13 = tmp5 * tmp12
    tmp14 = tl.where(tmp11, tmp8, tmp13)
    tmp15 = tmp9 - tmp14
    tmp16 = tl.where(tmp1, tmp15, tmp4)
    tmp17 = tmp2 == tmp10
    tmp19 = tl_math.exp(tmp18)
    tmp21 = tmp19 * tmp20
    tmp22 = tl.where(tmp17, tmp21, tmp8)
    tmp23 = tmp16 + tmp22
    tl.store(out_ptr0 + x2, tmp23, xmask)


@triton.jit
def triton_poi_fused_add_diag_embed_exp_fill_lift_fresh_mul_sub_5(in_ptr0,
    in_ptr1, in_ptr2, 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 // 16
    x3 = xindex % 16
    x0 = xindex % 4
    x1 = xindex // 4 % 4
    x5 = xindex
    tmp3 = tl.load(in_ptr0 + x3, xmask, eviction_policy='evict_last')
    tmp6 = tl.load(in_ptr1 + x3, xmask, eviction_policy='evict_last')
    tmp8 = tl.load(in_ptr2 + 5 * x0, xmask, eviction_policy='evict_last')
    tmp14 = tl.load(in_ptr2 + x3, xmask, eviction_policy='evict_last')
    tmp18 = tl.load(in_ptr1 + x5, xmask)
    tmp0 = x2
    tmp1 = tl.full([1], 0, tl.int32)
    tmp2 = tmp0 == tmp1
    tmp4 = x0
    tmp5 = tmp4 == tmp1
    tmp7 = tl_math.exp(tmp6)
    tmp9 = tmp7 * tmp8
    tmp10 = 0.0
    tmp11 = tl.where(tmp5, tmp10, tmp9)
    tmp12 = x1
    tmp13 = tmp12 == tmp1
    tmp15 = tmp7 * tmp14
    tmp16 = tl.where(tmp13, tmp10, tmp15)
    tmp17 = tmp11 - tmp16
    tmp19 = tl.where(tmp2, tmp17, tmp18)
    tmp20 = tl.where(tmp2, tmp3, tmp19)
    tl.store(out_ptr0 + x5, tmp20, xmask)


@triton.jit
def triton_poi_fused_eye_masked_fill_ne_sum_6(in_ptr0, out_ptr0, xnumel,
    XBLOCK: tl.constexpr):
    xnumel = 4
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex
    tmp7 = tl.load(in_ptr0 + (32 + x0), xmask)
    tmp16 = tl.load(in_ptr0 + (36 + x0), xmask)
    tmp25 = tl.load(in_ptr0 + (40 + x0), xmask)
    tmp34 = tl.load(in_ptr0 + (44 + x0), xmask)
    tmp0 = tl.full([1], 0, tl.int64)
    tmp1 = x0
    tmp2 = tmp0 == tmp1
    tmp3 = 1.0
    tmp4 = 0.0
    tmp5 = tl.where(tmp2, tmp3, tmp4)
    tmp6 = tmp5 != tmp4
    tmp8 = tl_math.exp(tmp7)
    tmp9 = 1e-05
    tmp10 = tmp8 + tmp9
    tmp11 = tl.where(tmp6, tmp4, tmp10)
    tmp12 = tl.full([1], 1, tl.int64)
    tmp13 = tmp12 == tmp1
    tmp14 = tl.where(tmp13, tmp3, tmp4)
    tmp15 = tmp14 != tmp4
    tmp17 = tl_math.exp(tmp16)
    tmp18 = tmp17 + tmp9
    tmp19 = tl.where(tmp15, tmp4, tmp18)
    tmp20 = tmp11 + tmp19
    tmp21 = tl.full([1], 2, tl.int64)
    tmp22 = tmp21 == tmp1
    tmp23 = tl.where(tmp22, tmp3, tmp4)
    tmp24 = tmp23 != tmp4
    tmp26 = tl_math.exp(tmp25)
    tmp27 = tmp26 + tmp9
    tmp28 = tl.where(tmp24, tmp4, tmp27)
    tmp29 = tmp20 + tmp28
    tmp30 = tl.full([1], 3, tl.int64)
    tmp31 = tmp30 == tmp1
    tmp32 = tl.where(tmp31, tmp3, tmp4)
    tmp33 = tmp32 != tmp4
    tmp35 = tl_math.exp(tmp34)
    tmp36 = tmp35 + tmp9
    tmp37 = tl.where(tmp33, tmp4, tmp36)
    tmp38 = tmp29 + tmp37
    tl.store(out_ptr0 + x0, tmp38, xmask)


@triton.jit
def triton_poi_fused_add_diag_embed_diagonal_copy_exp_eye_masked_fill_ne_neg_7(
    in_ptr0, in_ptr1, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x1 = xindex // 4
    x0 = xindex % 4
    x2 = xindex
    tmp3 = tl.load(in_ptr0 + (32 + 5 * x0), xmask, eviction_policy='evict_last'
        )
    tmp11 = tl.load(in_ptr0 + (32 + x2), xmask)
    tmp18 = tl.load(in_ptr1 + x0, xmask, eviction_policy='evict_last')
    tmp0 = x1
    tmp1 = tl.full([1], 0, tl.int32)
    tmp2 = tmp0 == tmp1
    tmp4 = tl_math.exp(tmp3)
    tmp5 = x0
    tmp6 = tmp0 == tmp5
    tmp7 = 1.0
    tmp8 = 0.0
    tmp9 = tl.where(tmp6, tmp7, tmp8)
    tmp10 = tmp9 != tmp8
    tmp12 = tl_math.exp(tmp11)
    tmp13 = 1e-05
    tmp14 = tmp12 + tmp13
    tmp15 = tl.where(tmp10, tmp8, tmp14)
    tmp16 = -tmp15
    tmp17 = tmp5 == tmp0
    tmp19 = tl.where(tmp17, tmp18, tmp8)
    tmp20 = tmp16 + tmp19
    tmp21 = tl.where(tmp2, tmp4, tmp20)
    tl.store(out_ptr0 + x2, tmp21, xmask)


@triton.jit
def triton_poi_fused_add_diag_embed_8(in_ptr0, in_ptr1, in_ptr2, out_ptr0,
    xnumel, XBLOCK: tl.constexpr):
    xnumel = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex % 4
    x2 = xindex
    x1 = xindex // 4
    tmp5 = tl.load(in_ptr0 + (16 + x2), xmask)
    tmp7 = tl.load(in_ptr1 + 5 * x0, xmask, eviction_policy='evict_last')
    tmp13 = tl.load(in_ptr1 + x2, xmask)
    tmp17 = tl.load(in_ptr2 + (16 + x2), xmask)
    tmp20 = tl.load(in_ptr0 + (16 + 5 * x0), xmask, eviction_policy=
        'evict_last')
    tmp22 = tl.load(in_ptr1 + x0, xmask, eviction_policy='evict_last')
    tmp0 = tl.full([1], 1, tl.int32)
    tmp1 = tmp0 == tmp0
    tmp2 = x0
    tmp3 = tl.full([1], 0, tl.int32)
    tmp4 = tmp2 == tmp3
    tmp6 = tl_math.exp(tmp5)
    tmp8 = tmp6 * tmp7
    tmp9 = 0.0
    tmp10 = tl.where(tmp4, tmp9, tmp8)
    tmp11 = x1
    tmp12 = tmp11 == tmp3
    tmp14 = tmp6 * tmp13
    tmp15 = tl.where(tmp12, tmp9, tmp14)
    tmp16 = tmp10 - tmp15
    tmp18 = tl.where(tmp1, tmp16, tmp17)
    tmp19 = tmp2 == tmp11
    tmp21 = tl_math.exp(tmp20)
    tmp23 = tmp21 * tmp22
    tmp24 = tl.where(tmp19, tmp23, tmp9)
    tmp25 = tmp18 + tmp24
    tl.store(out_ptr0 + x2, tmp25, xmask)


@triton.jit
def triton_poi_fused_add_diag_embed_exp_fill_lift_fresh_mul_sub_9(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 // 16
    x3 = xindex % 16
    x0 = xindex % 4
    x1 = xindex // 4 % 4
    x5 = xindex
    tmp3 = tl.load(in_ptr0 + x3, xmask, eviction_policy='evict_last')
    tmp7 = tl.load(in_ptr1 + (16 + x3), xmask, eviction_policy='evict_last')
    tmp9 = tl.load(in_ptr2 + 5 * x0, xmask, eviction_policy='evict_last')
    tmp15 = tl.load(in_ptr2 + x3, xmask, eviction_policy='evict_last')
    tmp19 = tl.load(in_out_ptr0 + x5, xmask)
    tmp0 = x2
    tmp1 = tl.full([1], 1, tl.int32)
    tmp2 = tmp0 == tmp1
    tmp4 = x0
    tmp5 = tl.full([1], 0, tl.int32)
    tmp6 = tmp4 == tmp5
    tmp8 = tl_math.exp(tmp7)
    tmp10 = tmp8 * tmp9
    tmp11 = 0.0
    tmp12 = tl.where(tmp6, tmp11, tmp10)
    tmp13 = x1
    tmp14 = tmp13 == tmp5
    tmp16 = tmp8 * tmp15
    tmp17 = tl.where(tmp14, tmp11, tmp16)
    tmp18 = tmp12 - tmp17
    tmp20 = tl.where(tmp2, tmp18, tmp19)
    tmp21 = tl.where(tmp2, tmp3, tmp20)
    tl.store(in_out_ptr0 + x5, tmp21, xmask)


@triton.jit
def triton_poi_fused_eye_masked_fill_ne_sum_10(in_ptr0, out_ptr0, xnumel,
    XBLOCK: tl.constexpr):
    xnumel = 4
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex
    tmp7 = tl.load(in_ptr0 + (48 + x0), xmask)
    tmp16 = tl.load(in_ptr0 + (52 + x0), xmask)
    tmp25 = tl.load(in_ptr0 + (56 + x0), xmask)
    tmp34 = tl.load(in_ptr0 + (60 + x0), xmask)
    tmp0 = tl.full([1], 0, tl.int64)
    tmp1 = x0
    tmp2 = tmp0 == tmp1
    tmp3 = 1.0
    tmp4 = 0.0
    tmp5 = tl.where(tmp2, tmp3, tmp4)
    tmp6 = tmp5 != tmp4
    tmp8 = tl_math.exp(tmp7)
    tmp9 = 1e-05
    tmp10 = tmp8 + tmp9
    tmp11 = tl.where(tmp6, tmp4, tmp10)
    tmp12 = tl.full([1], 1, tl.int64)
    tmp13 = tmp12 == tmp1
    tmp14 = tl.where(tmp13, tmp3, tmp4)
    tmp15 = tmp14 != tmp4
    tmp17 = tl_math.exp(tmp16)
    tmp18 = tmp17 + tmp9
    tmp19 = tl.where(tmp15, tmp4, tmp18)
    tmp20 = tmp11 + tmp19
    tmp21 = tl.full([1], 2, tl.int64)
    tmp22 = tmp21 == tmp1
    tmp23 = tl.where(tmp22, tmp3, tmp4)
    tmp24 = tmp23 != tmp4
    tmp26 = tl_math.exp(tmp25)
    tmp27 = tmp26 + tmp9
    tmp28 = tl.where(tmp24, tmp4, tmp27)
    tmp29 = tmp20 + tmp28
    tmp30 = tl.full([1], 3, tl.int64)
    tmp31 = tmp30 == tmp1
    tmp32 = tl.where(tmp31, tmp3, tmp4)
    tmp33 = tmp32 != tmp4
    tmp35 = tl_math.exp(tmp34)
    tmp36 = tmp35 + tmp9
    tmp37 = tl.where(tmp33, tmp4, tmp36)
    tmp38 = tmp29 + tmp37
    tl.store(out_ptr0 + x0, tmp38, xmask)


@triton.jit
def triton_poi_fused_add_diag_embed_diagonal_copy_exp_eye_masked_fill_ne_neg_11(
    in_ptr0, in_ptr1, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x1 = xindex // 4
    x0 = xindex % 4
    x2 = xindex
    tmp3 = tl.load(in_ptr0 + (48 + 5 * x0), xmask, eviction_policy='evict_last'
        )
    tmp11 = tl.load(in_ptr0 + (48 + x2), xmask)
    tmp18 = tl.load(in_ptr1 + x0, xmask, eviction_policy='evict_last')
    tmp0 = x1
    tmp1 = tl.full([1], 0, tl.int32)
    tmp2 = tmp0 == tmp1
    tmp4 = tl_math.exp(tmp3)
    tmp5 = x0
    tmp6 = tmp0 == tmp5
    tmp7 = 1.0
    tmp8 = 0.0
    tmp9 = tl.where(tmp6, tmp7, tmp8)
    tmp10 = tmp9 != tmp8
    tmp12 = tl_math.exp(tmp11)
    tmp13 = 1e-05
    tmp14 = tmp12 + tmp13
    tmp15 = tl.where(tmp10, tmp8, tmp14)
    tmp16 = -tmp15
    tmp17 = tmp5 == tmp0
    tmp19 = tl.where(tmp17, tmp18, tmp8)
    tmp20 = tmp16 + tmp19
    tmp21 = tl.where(tmp2, tmp4, tmp20)
    tl.store(out_ptr0 + x2, tmp21, xmask)


@triton.jit
def triton_poi_fused_add_diag_embed_12(in_ptr0, in_ptr1, in_ptr2, out_ptr0,
    xnumel, XBLOCK: tl.constexpr):
    xnumel = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex % 4
    x2 = xindex
    x1 = xindex // 4
    tmp5 = tl.load(in_ptr0 + (32 + x2), xmask)
    tmp7 = tl.load(in_ptr1 + 5 * x0, xmask, eviction_policy='evict_last')
    tmp13 = tl.load(in_ptr1 + x2, xmask)
    tmp17 = tl.load(in_ptr2 + (32 + x2), xmask)
    tmp20 = tl.load(in_ptr0 + (32 + 5 * x0), xmask, eviction_policy=
        'evict_last')
    tmp22 = tl.load(in_ptr1 + x0, xmask, eviction_policy='evict_last')
    tmp0 = tl.full([1], 2, tl.int32)
    tmp1 = tmp0 == tmp0
    tmp2 = x0
    tmp3 = tl.full([1], 0, tl.int32)
    tmp4 = tmp2 == tmp3
    tmp6 = tl_math.exp(tmp5)
    tmp8 = tmp6 * tmp7
    tmp9 = 0.0
    tmp10 = tl.where(tmp4, tmp9, tmp8)
    tmp11 = x1
    tmp12 = tmp11 == tmp3
    tmp14 = tmp6 * tmp13
    tmp15 = tl.where(tmp12, tmp9, tmp14)
    tmp16 = tmp10 - tmp15
    tmp18 = tl.where(tmp1, tmp16, tmp17)
    tmp19 = tmp2 == tmp11
    tmp21 = tl_math.exp(tmp20)
    tmp23 = tmp21 * tmp22
    tmp24 = tl.where(tmp19, tmp23, tmp9)
    tmp25 = tmp18 + tmp24
    tl.store(out_ptr0 + x2, tmp25, xmask)


@triton.jit
def triton_poi_fused_add_diag_embed_exp_fill_lift_fresh_mul_sub_13(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 // 16
    x3 = xindex % 16
    x0 = xindex % 4
    x1 = xindex // 4 % 4
    x5 = xindex
    tmp3 = tl.load(in_ptr0 + x3, xmask, eviction_policy='evict_last')
    tmp7 = tl.load(in_ptr1 + (32 + x3), xmask, eviction_policy='evict_last')
    tmp9 = tl.load(in_ptr2 + 5 * x0, xmask, eviction_policy='evict_last')
    tmp15 = tl.load(in_ptr2 + x3, xmask, eviction_policy='evict_last')
    tmp19 = tl.load(in_out_ptr0 + x5, xmask)
    tmp0 = x2
    tmp1 = tl.full([1], 2, tl.int32)
    tmp2 = tmp0 == tmp1
    tmp4 = x0
    tmp5 = tl.full([1], 0, tl.int32)
    tmp6 = tmp4 == tmp5
    tmp8 = tl_math.exp(tmp7)
    tmp10 = tmp8 * tmp9
    tmp11 = 0.0
    tmp12 = tl.where(tmp6, tmp11, tmp10)
    tmp13 = x1
    tmp14 = tmp13 == tmp5
    tmp16 = tmp8 * tmp15
    tmp17 = tl.where(tmp14, tmp11, tmp16)
    tmp18 = tmp12 - tmp17
    tmp20 = tl.where(tmp2, tmp18, tmp19)
    tmp21 = tl.where(tmp2, tmp3, tmp20)
    tl.store(in_out_ptr0 + x5, tmp21, xmask)


@triton.jit
def triton_poi_fused_add_diag_embed_14(in_ptr0, in_ptr1, in_ptr2, out_ptr0,
    xnumel, XBLOCK: tl.constexpr):
    xnumel = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex % 4
    x2 = xindex
    x1 = xindex // 4
    tmp5 = tl.load(in_ptr0 + (48 + x2), xmask)
    tmp7 = tl.load(in_ptr1 + 5 * x0, xmask, eviction_policy='evict_last')
    tmp13 = tl.load(in_ptr1 + x2, xmask)
    tmp17 = tl.load(in_ptr2 + (48 + x2), xmask)
    tmp20 = tl.load(in_ptr0 + (48 + 5 * x0), xmask, eviction_policy=
        'evict_last')
    tmp22 = tl.load(in_ptr1 + x0, xmask, eviction_policy='evict_last')
    tmp0 = tl.full([1], 3, tl.int32)
    tmp1 = tmp0 == tmp0
    tmp2 = x0
    tmp3 = tl.full([1], 0, tl.int32)
    tmp4 = tmp2 == tmp3
    tmp6 = tl_math.exp(tmp5)
    tmp8 = tmp6 * tmp7
    tmp9 = 0.0
    tmp10 = tl.where(tmp4, tmp9, tmp8)
    tmp11 = x1
    tmp12 = tmp11 == tmp3
    tmp14 = tmp6 * tmp13
    tmp15 = tl.where(tmp12, tmp9, tmp14)
    tmp16 = tmp10 - tmp15
    tmp18 = tl.where(tmp1, tmp16, tmp17)
    tmp19 = tmp2 == tmp11
    tmp21 = tl_math.exp(tmp20)
    tmp23 = tmp21 * tmp22
    tmp24 = tl.where(tmp19, tmp23, tmp9)
    tmp25 = tmp18 + tmp24
    tl.store(out_ptr0 + x2, tmp25, xmask)


@triton.jit
def triton_poi_fused_add_diag_embed_exp_fill_lift_fresh_mul_sub_15(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 // 16
    x3 = xindex % 16
    x0 = xindex % 4
    x1 = xindex // 4 % 4
    x5 = xindex
    tmp3 = tl.load(in_ptr0 + x3, xmask, eviction_policy='evict_last')
    tmp7 = tl.load(in_ptr1 + (48 + x3), xmask, eviction_policy='evict_last')
    tmp9 = tl.load(in_ptr2 + 5 * x0, xmask, eviction_policy='evict_last')
    tmp15 = tl.load(in_ptr2 + x3, xmask, eviction_policy='evict_last')
    tmp19 = tl.load(in_out_ptr0 + x5, xmask)
    tmp0 = x2
    tmp1 = tl.full([1], 3, tl.int32)
    tmp2 = tmp0 == tmp1
    tmp4 = x0
    tmp5 = tl.full([1], 0, tl.int32)
    tmp6 = tmp4 == tmp5
    tmp8 = tl_math.exp(tmp7)
    tmp10 = tmp8 * tmp9
    tmp11 = 0.0
    tmp12 = tl.where(tmp6, tmp11, tmp10)
    tmp13 = x1
    tmp14 = tmp13 == tmp5
    tmp16 = tmp8 * tmp15
    tmp17 = tl.where(tmp14, tmp11, tmp16)
    tmp18 = tmp12 - tmp17
    tmp20 = tl.where(tmp2, tmp18, tmp19)
    tmp21 = tl.where(tmp2, tmp3, tmp20)
    tl.store(in_out_ptr0 + x5, tmp21, xmask)


def call(args):
    arg0_1, = args
    args.clear()
    assert_size_stride(arg0_1, (4, 4, 4), (16, 4, 1))
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        buf0 = empty_strided_cuda((4,), (1,), torch.float32)
        get_raw_stream(0)
        triton_poi_fused_eye_masked_fill_ne_sum_0[grid(4)](arg0_1, buf0, 4,
            XBLOCK=4, num_warps=1, num_stages=1)
        buf1 = empty_strided_cuda((4, 4), (4, 1), torch.float32)
        triton_poi_fused_add_diag_embed_diagonal_copy_exp_eye_masked_fill_ne_neg_1[
            grid(16)](arg0_1, buf0, buf1, 16, XBLOCK=16, num_warps=1,
            num_stages=1)
        buf2 = torch.ops.aten.linalg_inv_ex.default(buf1)
        buf3 = buf2[0]
        del buf2
        buf5 = buf0
        del buf0
        triton_poi_fused_eye_masked_fill_ne_sum_2[grid(4)](arg0_1, buf5, 4,
            XBLOCK=4, num_warps=1, num_stages=1)
        buf6 = buf1
        del buf1
        triton_poi_fused_add_diag_embed_diagonal_copy_exp_eye_masked_fill_ne_neg_3[
            grid(16)](arg0_1, buf5, buf6, 16, XBLOCK=16, num_warps=1,
            num_stages=1)
        buf7 = torch.ops.aten.linalg_inv_ex.default(buf6)
        buf8 = buf7[0]
        del buf7
        buf10 = buf6
        del buf6
        triton_poi_fused_add_diag_embed_4[grid(16)](arg0_1, buf3, buf10, 16,
            XBLOCK=16, num_warps=1, num_stages=1)
        buf11 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        triton_poi_fused_add_diag_embed_exp_fill_lift_fresh_mul_sub_5[grid(64)
            ](buf10, arg0_1, buf3, buf11, 64, XBLOCK=64, num_warps=1,
            num_stages=1)
        del buf10
        buf12 = buf5
        del buf5
        triton_poi_fused_eye_masked_fill_ne_sum_6[grid(4)](arg0_1, buf12, 4,
            XBLOCK=4, num_warps=1, num_stages=1)
        buf13 = reinterpret_tensor(buf3, (4, 4), (4, 1), 0)
        del buf3
        triton_poi_fused_add_diag_embed_diagonal_copy_exp_eye_masked_fill_ne_neg_7[
            grid(16)](arg0_1, buf12, buf13, 16, XBLOCK=16, num_warps=1,
            num_stages=1)
        buf14 = torch.ops.aten.linalg_inv_ex.default(buf13)
        buf15 = buf14[0]
        del buf14
        buf17 = buf13
        del buf13
        triton_poi_fused_add_diag_embed_8[grid(16)](arg0_1, buf8, buf11,
            buf17, 16, XBLOCK=16, num_warps=1, num_stages=1)
        buf18 = buf11
        del buf11
        triton_poi_fused_add_diag_embed_exp_fill_lift_fresh_mul_sub_9[grid(64)
            ](buf18, buf17, arg0_1, buf8, 64, XBLOCK=64, num_warps=1,
            num_stages=1)
        del buf17
        buf19 = buf12
        del buf12
        triton_poi_fused_eye_masked_fill_ne_sum_10[grid(4)](arg0_1, buf19, 
            4, XBLOCK=4, num_warps=1, num_stages=1)
        buf20 = reinterpret_tensor(buf8, (4, 4), (4, 1), 0)
        del buf8
        triton_poi_fused_add_diag_embed_diagonal_copy_exp_eye_masked_fill_ne_neg_11[
            grid(16)](arg0_1, buf19, buf20, 16, XBLOCK=16, num_warps=1,
            num_stages=1)
        del buf19
        buf21 = torch.ops.aten.linalg_inv_ex.default(buf20)
        buf22 = buf21[0]
        del buf21
        buf24 = buf20
        del buf20
        triton_poi_fused_add_diag_embed_12[grid(16)](arg0_1, buf15, buf18,
            buf24, 16, XBLOCK=16, num_warps=1, num_stages=1)
        buf25 = buf18
        del buf18
        triton_poi_fused_add_diag_embed_exp_fill_lift_fresh_mul_sub_13[grid(64)
            ](buf25, buf24, arg0_1, buf15, 64, XBLOCK=64, num_warps=1,
            num_stages=1)
        del buf15
        buf26 = buf24
        del buf24
        triton_poi_fused_add_diag_embed_14[grid(16)](arg0_1, buf22, buf25,
            buf26, 16, XBLOCK=16, num_warps=1, num_stages=1)
        buf27 = buf25
        del buf25
        triton_poi_fused_add_diag_embed_exp_fill_lift_fresh_mul_sub_15[grid(64)
            ](buf27, buf26, arg0_1, buf22, 64, XBLOCK=64, num_warps=1,
            num_stages=1)
        del arg0_1
        del buf22
        del buf26
    return buf27,


class MatrixTreeNew(nn.Module):
    """Implementation of the matrix-tree theorem for computing marginals
    of non-projective dependency parsing. This attention layer is used
    in the paper "Learning Structured Text Representations"
    :cite:`DBLP:journals/corr/LiuL17d`.
    """

    def __init__(self, eps=1e-05):
        self.eps = eps
        super(MatrixTreeNew, self).__init__()

    def forward(self, input_0):
        arg0_1 = input_0
        output = call([arg0_1])
        return output[0]


Measured latency_ms: 0.06096