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 import triton_helpers
from torch._inductor.runtime.triton_helpers import math as tl_math
import torch.nn as nn
import torch.nn.functional as F
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_per_fused_sum_0(in_ptr0, out_ptr0, xnumel, rnumel, XBLOCK: tl.
    constexpr):
    RBLOCK: tl.constexpr = 16
    xoffset = tl.program_id(0) * XBLOCK
    xoffset + tl.arange(0, XBLOCK)[:, None]
    tl.full([XBLOCK, RBLOCK], True, tl.int1)
    rindex = tl.arange(0, RBLOCK)[None, :]
    tl.full([XBLOCK, RBLOCK], True, tl.int1)
    r0 = rindex
    tmp0 = tl.load(in_ptr0 + r0, None)
    tmp1 = 20.0
    tmp2 = tmp0 * tmp1
    tmp3 = tl_math.exp(tmp2)
    tmp4 = tl.broadcast_to(tmp3, [XBLOCK, RBLOCK])
    tmp6 = tl.sum(tmp4, 1)[:, None]
    tl.store(out_ptr0 + tl.full([XBLOCK, 1], 0, tl.int32), tmp6, None)


@triton.jit
def triton_poi_fused_sum_1(in_ptr0, in_ptr1, 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
    tmp0 = tl.load(in_ptr0 + x0, xmask)
    tmp4 = tl.load(in_ptr1 + 0)
    tmp5 = tl.broadcast_to(tmp4, [XBLOCK])
    tmp7 = tl.load(in_ptr0 + (4 + x0), xmask)
    tmp12 = tl.load(in_ptr0 + (8 + x0), xmask)
    tmp17 = tl.load(in_ptr0 + (12 + x0), xmask)
    tmp1 = 20.0
    tmp2 = tmp0 * tmp1
    tmp3 = tl_math.exp(tmp2)
    tmp6 = tmp3 / tmp5
    tmp8 = tmp7 * tmp1
    tmp9 = tl_math.exp(tmp8)
    tmp10 = tmp9 / tmp5
    tmp11 = tmp6 + tmp10
    tmp13 = tmp12 * tmp1
    tmp14 = tl_math.exp(tmp13)
    tmp15 = tmp14 / tmp5
    tmp16 = tmp11 + tmp15
    tmp18 = tmp17 * tmp1
    tmp19 = tl_math.exp(tmp18)
    tmp20 = tmp19 / tmp5
    tmp21 = tmp16 + tmp20
    tl.store(out_ptr0 + x0, tmp21, xmask)


@triton.jit
def triton_poi_fused_sum_2(in_ptr0, in_ptr1, in_ptr2, 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
    tmp0 = tl.load(in_ptr0 + 4 * x0, xmask, eviction_policy='evict_last')
    tmp4 = tl.load(in_ptr1 + 0)
    tmp5 = tl.broadcast_to(tmp4, [XBLOCK])
    tmp7 = tl.load(in_ptr2 + 0)
    tmp8 = tl.broadcast_to(tmp7, [XBLOCK])
    tmp12 = tl.load(in_ptr0 + (1 + 4 * x0), xmask, eviction_policy='evict_last'
        )
    tmp16 = tl.load(in_ptr2 + 1)
    tmp17 = tl.broadcast_to(tmp16, [XBLOCK])
    tmp21 = tl.load(in_ptr0 + (2 + 4 * x0), xmask, eviction_policy='evict_last'
        )
    tmp25 = tl.load(in_ptr2 + 2)
    tmp26 = tl.broadcast_to(tmp25, [XBLOCK])
    tmp30 = tl.load(in_ptr0 + (3 + 4 * x0), xmask, eviction_policy='evict_last'
        )
    tmp34 = tl.load(in_ptr2 + 3)
    tmp35 = tl.broadcast_to(tmp34, [XBLOCK])
    tmp1 = 20.0
    tmp2 = tmp0 * tmp1
    tmp3 = tl_math.exp(tmp2)
    tmp6 = tmp3 / tmp5
    tmp9 = tmp6 / tmp8
    tmp10 = 0.25
    tmp11 = tmp9 * tmp10
    tmp13 = tmp12 * tmp1
    tmp14 = tl_math.exp(tmp13)
    tmp15 = tmp14 / tmp5
    tmp18 = tmp15 / tmp17
    tmp19 = tmp18 * tmp10
    tmp20 = tmp11 + tmp19
    tmp22 = tmp21 * tmp1
    tmp23 = tl_math.exp(tmp22)
    tmp24 = tmp23 / tmp5
    tmp27 = tmp24 / tmp26
    tmp28 = tmp27 * tmp10
    tmp29 = tmp20 + tmp28
    tmp31 = tmp30 * tmp1
    tmp32 = tl_math.exp(tmp31)
    tmp33 = tmp32 / tmp5
    tmp36 = tmp33 / tmp35
    tmp37 = tmp36 * tmp10
    tmp38 = tmp29 + tmp37
    tl.store(out_ptr0 + x0, tmp38, xmask)


@triton.jit
def triton_poi_fused_sum_3(in_ptr0, in_ptr1, in_ptr2, in_ptr3, 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
    tmp0 = tl.load(in_ptr0 + x0, xmask)
    tmp4 = tl.load(in_ptr1 + 0)
    tmp5 = tl.broadcast_to(tmp4, [XBLOCK])
    tmp7 = tl.load(in_ptr2 + x0, xmask)
    tmp11 = tl.load(in_ptr3 + 0)
    tmp12 = tl.broadcast_to(tmp11, [XBLOCK])
    tmp15 = tl.load(in_ptr0 + (4 + x0), xmask)
    tmp21 = tl.load(in_ptr3 + 1)
    tmp22 = tl.broadcast_to(tmp21, [XBLOCK])
    tmp26 = tl.load(in_ptr0 + (8 + x0), xmask)
    tmp32 = tl.load(in_ptr3 + 2)
    tmp33 = tl.broadcast_to(tmp32, [XBLOCK])
    tmp37 = tl.load(in_ptr0 + (12 + x0), xmask)
    tmp43 = tl.load(in_ptr3 + 3)
    tmp44 = tl.broadcast_to(tmp43, [XBLOCK])
    tmp1 = 20.0
    tmp2 = tmp0 * tmp1
    tmp3 = tl_math.exp(tmp2)
    tmp6 = tmp3 / tmp5
    tmp8 = tmp6 / tmp7
    tmp9 = 0.25
    tmp10 = tmp8 * tmp9
    tmp13 = tmp10 / tmp12
    tmp14 = tmp13 * tmp9
    tmp16 = tmp15 * tmp1
    tmp17 = tl_math.exp(tmp16)
    tmp18 = tmp17 / tmp5
    tmp19 = tmp18 / tmp7
    tmp20 = tmp19 * tmp9
    tmp23 = tmp20 / tmp22
    tmp24 = tmp23 * tmp9
    tmp25 = tmp14 + tmp24
    tmp27 = tmp26 * tmp1
    tmp28 = tl_math.exp(tmp27)
    tmp29 = tmp28 / tmp5
    tmp30 = tmp29 / tmp7
    tmp31 = tmp30 * tmp9
    tmp34 = tmp31 / tmp33
    tmp35 = tmp34 * tmp9
    tmp36 = tmp25 + tmp35
    tmp38 = tmp37 * tmp1
    tmp39 = tl_math.exp(tmp38)
    tmp40 = tmp39 / tmp5
    tmp41 = tmp40 / tmp7
    tmp42 = tmp41 * tmp9
    tmp45 = tmp42 / tmp44
    tmp46 = tmp45 * tmp9
    tmp47 = tmp36 + tmp46
    tl.store(out_ptr0 + x0, tmp47, xmask)


@triton.jit
def triton_poi_fused_div_4(in_ptr0, in_ptr1, in_ptr2, in_ptr3, in_ptr4,
    out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x2 = xindex
    x0 = xindex % 4
    x1 = xindex // 4
    tmp0 = tl.load(in_ptr0 + x2, xmask)
    tmp4 = tl.load(in_ptr1 + 0)
    tmp5 = tl.broadcast_to(tmp4, [XBLOCK])
    tmp7 = tl.load(in_ptr2 + x0, xmask, eviction_policy='evict_last')
    tmp11 = tl.load(in_ptr3 + x1, xmask, eviction_policy='evict_last')
    tmp14 = tl.load(in_ptr4 + x0, xmask, eviction_policy='evict_last')
    tmp1 = 20.0
    tmp2 = tmp0 * tmp1
    tmp3 = tl_math.exp(tmp2)
    tmp6 = tmp3 / tmp5
    tmp8 = tmp6 / tmp7
    tmp9 = 0.25
    tmp10 = tmp8 * tmp9
    tmp12 = tmp10 / tmp11
    tmp13 = tmp12 * tmp9
    tmp15 = tmp13 / tmp14
    tmp16 = tmp15 * tmp9
    tl.store(out_ptr0 + x2, tmp16, xmask)


@triton.jit
def triton_poi_fused_div_5(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x2 = xindex
    x1 = xindex // 4
    tmp0 = tl.load(in_ptr0 + x2, xmask)
    tmp1 = tl.load(in_ptr0 + 4 * x1, xmask, eviction_policy='evict_last')
    tmp2 = tl.load(in_ptr0 + (1 + 4 * x1), xmask, eviction_policy='evict_last')
    tmp4 = tl.load(in_ptr0 + (2 + 4 * x1), xmask, eviction_policy='evict_last')
    tmp6 = tl.load(in_ptr0 + (3 + 4 * x1), xmask, eviction_policy='evict_last')
    tmp3 = tmp1 + tmp2
    tmp5 = tmp3 + tmp4
    tmp7 = tmp5 + tmp6
    tmp8 = tmp0 / tmp7
    tmp9 = 0.25
    tmp10 = tmp8 * tmp9
    tl.store(out_ptr0 + x2, tmp10, xmask)


@triton.jit
def triton_poi_fused_div_6(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x2 = xindex
    x0 = xindex % 4
    tmp0 = tl.load(in_ptr0 + x2, xmask)
    tmp1 = tl.load(in_ptr0 + x0, xmask, eviction_policy='evict_last')
    tmp2 = tl.load(in_ptr0 + (4 + x0), xmask, eviction_policy='evict_last')
    tmp4 = tl.load(in_ptr0 + (8 + x0), xmask, eviction_policy='evict_last')
    tmp6 = tl.load(in_ptr0 + (12 + x0), xmask, eviction_policy='evict_last')
    tmp3 = tmp1 + tmp2
    tmp5 = tmp3 + tmp4
    tmp7 = tmp5 + tmp6
    tmp8 = tmp0 / tmp7
    tmp9 = 0.25
    tmp10 = tmp8 * tmp9
    tl.store(out_ptr0 + x2, tmp10, xmask)


@triton.jit
def triton_poi_fused_mul_7(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x2 = xindex
    x1 = xindex // 4
    tmp0 = tl.load(in_ptr0 + x2, xmask)
    tmp1 = tl.load(in_ptr0 + 4 * x1, xmask, eviction_policy='evict_last')
    tmp2 = tl.load(in_ptr0 + (1 + 4 * x1), xmask, eviction_policy='evict_last')
    tmp4 = tl.load(in_ptr0 + (2 + 4 * x1), xmask, eviction_policy='evict_last')
    tmp6 = tl.load(in_ptr0 + (3 + 4 * x1), xmask, eviction_policy='evict_last')
    tmp3 = tmp1 + tmp2
    tmp5 = tmp3 + tmp4
    tmp7 = tmp5 + tmp6
    tmp8 = tmp0 / tmp7
    tmp9 = 0.25
    tmp10 = tmp8 * tmp9
    tmp11 = 4.0
    tmp12 = tmp10 * tmp11
    tl.store(out_ptr0 + x2, tmp12, xmask)


@triton.jit
def triton_per_fused_sum_8(in_ptr0, out_ptr0, out_ptr1, out_ptr2, xnumel,
    rnumel, XBLOCK: tl.constexpr):
    RBLOCK: tl.constexpr = 16
    xoffset = tl.program_id(0) * XBLOCK
    xoffset + tl.arange(0, XBLOCK)[:, None]
    tl.full([XBLOCK, RBLOCK], True, tl.int1)
    rindex = tl.arange(0, RBLOCK)[None, :]
    tl.full([XBLOCK, RBLOCK], True, tl.int1)
    r0 = rindex
    r2 = rindex // 4
    tmp0 = tl.load(in_ptr0 + (16 + r0), None)
    tmp9 = tl.load(in_ptr0 + (16 + 4 * r2), None, eviction_policy='evict_last')
    tmp11 = tl.load(in_ptr0 + (17 + 4 * r2), None, eviction_policy='evict_last'
        )
    tmp14 = tl.load(in_ptr0 + (18 + 4 * r2), None, eviction_policy='evict_last'
        )
    tmp17 = tl.load(in_ptr0 + (19 + 4 * r2), None, eviction_policy='evict_last'
        )
    tmp1 = 20.0
    tmp2 = tmp0 * tmp1
    tmp3 = tl_math.exp(tmp2)
    tmp4 = tl.broadcast_to(tmp3, [XBLOCK, RBLOCK])
    tmp6 = tl.sum(tmp4, 1)[:, None]
    tmp7 = 1.0
    tmp8 = tmp0 * tmp7
    tmp10 = tmp9 * tmp7
    tmp12 = tmp11 * tmp7
    tmp13 = triton_helpers.maximum(tmp10, tmp12)
    tmp15 = tmp14 * tmp7
    tmp16 = triton_helpers.maximum(tmp13, tmp15)
    tmp18 = tmp17 * tmp7
    tmp19 = triton_helpers.maximum(tmp16, tmp18)
    tmp20 = tmp8 - tmp19
    tmp21 = 10.0
    tmp22 = tmp20 * tmp21
    tl.store(out_ptr1 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp22, None)
    tl.store(out_ptr2 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp22, None)
    tl.store(out_ptr0 + tl.full([XBLOCK, 1], 0, tl.int32), tmp6, None)


@triton.jit
def triton_poi_fused_sum_9(in_ptr0, in_ptr1, 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
    tmp0 = tl.load(in_ptr0 + (16 + x0), xmask)
    tmp4 = tl.load(in_ptr1 + 0)
    tmp5 = tl.broadcast_to(tmp4, [XBLOCK])
    tmp7 = tl.load(in_ptr0 + (20 + x0), xmask)
    tmp12 = tl.load(in_ptr0 + (24 + x0), xmask)
    tmp17 = tl.load(in_ptr0 + (28 + x0), xmask)
    tmp1 = 20.0
    tmp2 = tmp0 * tmp1
    tmp3 = tl_math.exp(tmp2)
    tmp6 = tmp3 / tmp5
    tmp8 = tmp7 * tmp1
    tmp9 = tl_math.exp(tmp8)
    tmp10 = tmp9 / tmp5
    tmp11 = tmp6 + tmp10
    tmp13 = tmp12 * tmp1
    tmp14 = tl_math.exp(tmp13)
    tmp15 = tmp14 / tmp5
    tmp16 = tmp11 + tmp15
    tmp18 = tmp17 * tmp1
    tmp19 = tl_math.exp(tmp18)
    tmp20 = tmp19 / tmp5
    tmp21 = tmp16 + tmp20
    tl.store(out_ptr0 + x0, tmp21, xmask)


@triton.jit
def triton_poi_fused_sum_10(in_ptr0, in_ptr1, in_ptr2, 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
    tmp0 = tl.load(in_ptr0 + (16 + 4 * x0), xmask, eviction_policy='evict_last'
        )
    tmp4 = tl.load(in_ptr1 + 0)
    tmp5 = tl.broadcast_to(tmp4, [XBLOCK])
    tmp7 = tl.load(in_ptr2 + 0)
    tmp8 = tl.broadcast_to(tmp7, [XBLOCK])
    tmp12 = tl.load(in_ptr0 + (17 + 4 * x0), xmask, eviction_policy=
        'evict_last')
    tmp16 = tl.load(in_ptr2 + 1)
    tmp17 = tl.broadcast_to(tmp16, [XBLOCK])
    tmp21 = tl.load(in_ptr0 + (18 + 4 * x0), xmask, eviction_policy=
        'evict_last')
    tmp25 = tl.load(in_ptr2 + 2)
    tmp26 = tl.broadcast_to(tmp25, [XBLOCK])
    tmp30 = tl.load(in_ptr0 + (19 + 4 * x0), xmask, eviction_policy=
        'evict_last')
    tmp34 = tl.load(in_ptr2 + 3)
    tmp35 = tl.broadcast_to(tmp34, [XBLOCK])
    tmp1 = 20.0
    tmp2 = tmp0 * tmp1
    tmp3 = tl_math.exp(tmp2)
    tmp6 = tmp3 / tmp5
    tmp9 = tmp6 / tmp8
    tmp10 = 0.25
    tmp11 = tmp9 * tmp10
    tmp13 = tmp12 * tmp1
    tmp14 = tl_math.exp(tmp13)
    tmp15 = tmp14 / tmp5
    tmp18 = tmp15 / tmp17
    tmp19 = tmp18 * tmp10
    tmp20 = tmp11 + tmp19
    tmp22 = tmp21 * tmp1
    tmp23 = tl_math.exp(tmp22)
    tmp24 = tmp23 / tmp5
    tmp27 = tmp24 / tmp26
    tmp28 = tmp27 * tmp10
    tmp29 = tmp20 + tmp28
    tmp31 = tmp30 * tmp1
    tmp32 = tl_math.exp(tmp31)
    tmp33 = tmp32 / tmp5
    tmp36 = tmp33 / tmp35
    tmp37 = tmp36 * tmp10
    tmp38 = tmp29 + tmp37
    tl.store(out_ptr0 + x0, tmp38, xmask)


@triton.jit
def triton_poi_fused_sum_11(in_ptr0, in_ptr1, in_ptr2, in_ptr3, 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
    tmp0 = tl.load(in_ptr0 + (16 + x0), xmask)
    tmp4 = tl.load(in_ptr1 + 0)
    tmp5 = tl.broadcast_to(tmp4, [XBLOCK])
    tmp7 = tl.load(in_ptr2 + x0, xmask)
    tmp11 = tl.load(in_ptr3 + 0)
    tmp12 = tl.broadcast_to(tmp11, [XBLOCK])
    tmp15 = tl.load(in_ptr0 + (20 + x0), xmask)
    tmp21 = tl.load(in_ptr3 + 1)
    tmp22 = tl.broadcast_to(tmp21, [XBLOCK])
    tmp26 = tl.load(in_ptr0 + (24 + x0), xmask)
    tmp32 = tl.load(in_ptr3 + 2)
    tmp33 = tl.broadcast_to(tmp32, [XBLOCK])
    tmp37 = tl.load(in_ptr0 + (28 + x0), xmask)
    tmp43 = tl.load(in_ptr3 + 3)
    tmp44 = tl.broadcast_to(tmp43, [XBLOCK])
    tmp1 = 20.0
    tmp2 = tmp0 * tmp1
    tmp3 = tl_math.exp(tmp2)
    tmp6 = tmp3 / tmp5
    tmp8 = tmp6 / tmp7
    tmp9 = 0.25
    tmp10 = tmp8 * tmp9
    tmp13 = tmp10 / tmp12
    tmp14 = tmp13 * tmp9
    tmp16 = tmp15 * tmp1
    tmp17 = tl_math.exp(tmp16)
    tmp18 = tmp17 / tmp5
    tmp19 = tmp18 / tmp7
    tmp20 = tmp19 * tmp9
    tmp23 = tmp20 / tmp22
    tmp24 = tmp23 * tmp9
    tmp25 = tmp14 + tmp24
    tmp27 = tmp26 * tmp1
    tmp28 = tl_math.exp(tmp27)
    tmp29 = tmp28 / tmp5
    tmp30 = tmp29 / tmp7
    tmp31 = tmp30 * tmp9
    tmp34 = tmp31 / tmp33
    tmp35 = tmp34 * tmp9
    tmp36 = tmp25 + tmp35
    tmp38 = tmp37 * tmp1
    tmp39 = tl_math.exp(tmp38)
    tmp40 = tmp39 / tmp5
    tmp41 = tmp40 / tmp7
    tmp42 = tmp41 * tmp9
    tmp45 = tmp42 / tmp44
    tmp46 = tmp45 * tmp9
    tmp47 = tmp36 + tmp46
    tl.store(out_ptr0 + x0, tmp47, xmask)


@triton.jit
def triton_poi_fused_div_12(in_ptr0, in_ptr1, in_ptr2, in_ptr3, in_ptr4,
    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
    x2 = xindex
    x1 = xindex // 4
    x0 = xindex % 4
    tmp0 = tl.load(in_ptr0 + (16 + x2), xmask)
    tmp3 = tl.load(in_ptr0 + (16 + 4 * x1), xmask, eviction_policy='evict_last'
        )
    tmp5 = tl.load(in_ptr0 + (17 + 4 * x1), xmask, eviction_policy='evict_last'
        )
    tmp8 = tl.load(in_ptr0 + (18 + 4 * x1), xmask, eviction_policy='evict_last'
        )
    tmp11 = tl.load(in_ptr0 + (19 + 4 * x1), xmask, eviction_policy=
        'evict_last')
    tmp20 = tl.load(in_ptr1 + 0)
    tmp21 = tl.broadcast_to(tmp20, [XBLOCK])
    tmp23 = tl.load(in_ptr2 + x0, xmask, eviction_policy='evict_last')
    tmp27 = tl.load(in_ptr3 + x1, xmask, eviction_policy='evict_last')
    tmp30 = tl.load(in_ptr4 + x0, xmask, eviction_policy='evict_last')
    tmp1 = 1.0
    tmp2 = tmp0 * tmp1
    tmp4 = tmp3 * tmp1
    tmp6 = tmp5 * tmp1
    tmp7 = triton_helpers.maximum(tmp4, tmp6)
    tmp9 = tmp8 * tmp1
    tmp10 = triton_helpers.maximum(tmp7, tmp9)
    tmp12 = tmp11 * tmp1
    tmp13 = triton_helpers.maximum(tmp10, tmp12)
    tmp14 = tmp2 - tmp13
    tmp15 = 10.0
    tmp16 = tmp14 * tmp15
    tmp17 = 20.0
    tmp18 = tmp0 * tmp17
    tmp19 = tl_math.exp(tmp18)
    tmp22 = tmp19 / tmp21
    tmp24 = tmp22 / tmp23
    tmp25 = 0.25
    tmp26 = tmp24 * tmp25
    tmp28 = tmp26 / tmp27
    tmp29 = tmp28 * tmp25
    tmp31 = tmp29 / tmp30
    tmp32 = tmp31 * tmp25
    tl.store(out_ptr0 + x2, tmp16, xmask)
    tl.store(out_ptr1 + x2, tmp32, xmask)


@triton.jit
def triton_per_fused_sum_13(in_ptr0, out_ptr0, out_ptr1, xnumel, rnumel,
    XBLOCK: tl.constexpr):
    RBLOCK: tl.constexpr = 16
    xoffset = tl.program_id(0) * XBLOCK
    xoffset + tl.arange(0, XBLOCK)[:, None]
    tl.full([XBLOCK, RBLOCK], True, tl.int1)
    rindex = tl.arange(0, RBLOCK)[None, :]
    tl.full([XBLOCK, RBLOCK], True, tl.int1)
    r0 = rindex
    r2 = rindex // 4
    tmp0 = tl.load(in_ptr0 + (32 + r0), None)
    tmp9 = tl.load(in_ptr0 + (32 + 4 * r2), None, eviction_policy='evict_last')
    tmp11 = tl.load(in_ptr0 + (33 + 4 * r2), None, eviction_policy='evict_last'
        )
    tmp14 = tl.load(in_ptr0 + (34 + 4 * r2), None, eviction_policy='evict_last'
        )
    tmp17 = tl.load(in_ptr0 + (35 + 4 * r2), None, eviction_policy='evict_last'
        )
    tmp1 = 20.0
    tmp2 = tmp0 * tmp1
    tmp3 = tl_math.exp(tmp2)
    tmp4 = tl.broadcast_to(tmp3, [XBLOCK, RBLOCK])
    tmp6 = tl.sum(tmp4, 1)[:, None]
    tmp7 = 1.0
    tmp8 = tmp0 * tmp7
    tmp10 = tmp9 * tmp7
    tmp12 = tmp11 * tmp7
    tmp13 = triton_helpers.maximum(tmp10, tmp12)
    tmp15 = tmp14 * tmp7
    tmp16 = triton_helpers.maximum(tmp13, tmp15)
    tmp18 = tmp17 * tmp7
    tmp19 = triton_helpers.maximum(tmp16, tmp18)
    tmp20 = tmp8 - tmp19
    tmp21 = 10.0
    tmp22 = tmp20 * tmp21
    tl.store(out_ptr1 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp22, None)
    tl.store(out_ptr0 + tl.full([XBLOCK, 1], 0, tl.int32), tmp6, None)


@triton.jit
def triton_poi_fused_sum_14(in_ptr0, in_ptr1, 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
    tmp0 = tl.load(in_ptr0 + (32 + x0), xmask)
    tmp4 = tl.load(in_ptr1 + 0)
    tmp5 = tl.broadcast_to(tmp4, [XBLOCK])
    tmp7 = tl.load(in_ptr0 + (36 + x0), xmask)
    tmp12 = tl.load(in_ptr0 + (40 + x0), xmask)
    tmp17 = tl.load(in_ptr0 + (44 + x0), xmask)
    tmp1 = 20.0
    tmp2 = tmp0 * tmp1
    tmp3 = tl_math.exp(tmp2)
    tmp6 = tmp3 / tmp5
    tmp8 = tmp7 * tmp1
    tmp9 = tl_math.exp(tmp8)
    tmp10 = tmp9 / tmp5
    tmp11 = tmp6 + tmp10
    tmp13 = tmp12 * tmp1
    tmp14 = tl_math.exp(tmp13)
    tmp15 = tmp14 / tmp5
    tmp16 = tmp11 + tmp15
    tmp18 = tmp17 * tmp1
    tmp19 = tl_math.exp(tmp18)
    tmp20 = tmp19 / tmp5
    tmp21 = tmp16 + tmp20
    tl.store(out_ptr0 + x0, tmp21, xmask)


@triton.jit
def triton_poi_fused_sum_15(in_ptr0, in_ptr1, in_ptr2, 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
    tmp0 = tl.load(in_ptr0 + (32 + 4 * x0), xmask, eviction_policy='evict_last'
        )
    tmp4 = tl.load(in_ptr1 + 0)
    tmp5 = tl.broadcast_to(tmp4, [XBLOCK])
    tmp7 = tl.load(in_ptr2 + 0)
    tmp8 = tl.broadcast_to(tmp7, [XBLOCK])
    tmp12 = tl.load(in_ptr0 + (33 + 4 * x0), xmask, eviction_policy=
        'evict_last')
    tmp16 = tl.load(in_ptr2 + 1)
    tmp17 = tl.broadcast_to(tmp16, [XBLOCK])
    tmp21 = tl.load(in_ptr0 + (34 + 4 * x0), xmask, eviction_policy=
        'evict_last')
    tmp25 = tl.load(in_ptr2 + 2)
    tmp26 = tl.broadcast_to(tmp25, [XBLOCK])
    tmp30 = tl.load(in_ptr0 + (35 + 4 * x0), xmask, eviction_policy=
        'evict_last')
    tmp34 = tl.load(in_ptr2 + 3)
    tmp35 = tl.broadcast_to(tmp34, [XBLOCK])
    tmp1 = 20.0
    tmp2 = tmp0 * tmp1
    tmp3 = tl_math.exp(tmp2)
    tmp6 = tmp3 / tmp5
    tmp9 = tmp6 / tmp8
    tmp10 = 0.25
    tmp11 = tmp9 * tmp10
    tmp13 = tmp12 * tmp1
    tmp14 = tl_math.exp(tmp13)
    tmp15 = tmp14 / tmp5
    tmp18 = tmp15 / tmp17
    tmp19 = tmp18 * tmp10
    tmp20 = tmp11 + tmp19
    tmp22 = tmp21 * tmp1
    tmp23 = tl_math.exp(tmp22)
    tmp24 = tmp23 / tmp5
    tmp27 = tmp24 / tmp26
    tmp28 = tmp27 * tmp10
    tmp29 = tmp20 + tmp28
    tmp31 = tmp30 * tmp1
    tmp32 = tl_math.exp(tmp31)
    tmp33 = tmp32 / tmp5
    tmp36 = tmp33 / tmp35
    tmp37 = tmp36 * tmp10
    tmp38 = tmp29 + tmp37
    tl.store(out_ptr0 + x0, tmp38, xmask)


@triton.jit
def triton_poi_fused_sum_16(in_ptr0, in_ptr1, in_ptr2, in_ptr3, 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
    tmp0 = tl.load(in_ptr0 + (32 + x0), xmask)
    tmp4 = tl.load(in_ptr1 + 0)
    tmp5 = tl.broadcast_to(tmp4, [XBLOCK])
    tmp7 = tl.load(in_ptr2 + x0, xmask)
    tmp11 = tl.load(in_ptr3 + 0)
    tmp12 = tl.broadcast_to(tmp11, [XBLOCK])
    tmp15 = tl.load(in_ptr0 + (36 + x0), xmask)
    tmp21 = tl.load(in_ptr3 + 1)
    tmp22 = tl.broadcast_to(tmp21, [XBLOCK])
    tmp26 = tl.load(in_ptr0 + (40 + x0), xmask)
    tmp32 = tl.load(in_ptr3 + 2)
    tmp33 = tl.broadcast_to(tmp32, [XBLOCK])
    tmp37 = tl.load(in_ptr0 + (44 + x0), xmask)
    tmp43 = tl.load(in_ptr3 + 3)
    tmp44 = tl.broadcast_to(tmp43, [XBLOCK])
    tmp1 = 20.0
    tmp2 = tmp0 * tmp1
    tmp3 = tl_math.exp(tmp2)
    tmp6 = tmp3 / tmp5
    tmp8 = tmp6 / tmp7
    tmp9 = 0.25
    tmp10 = tmp8 * tmp9
    tmp13 = tmp10 / tmp12
    tmp14 = tmp13 * tmp9
    tmp16 = tmp15 * tmp1
    tmp17 = tl_math.exp(tmp16)
    tmp18 = tmp17 / tmp5
    tmp19 = tmp18 / tmp7
    tmp20 = tmp19 * tmp9
    tmp23 = tmp20 / tmp22
    tmp24 = tmp23 * tmp9
    tmp25 = tmp14 + tmp24
    tmp27 = tmp26 * tmp1
    tmp28 = tl_math.exp(tmp27)
    tmp29 = tmp28 / tmp5
    tmp30 = tmp29 / tmp7
    tmp31 = tmp30 * tmp9
    tmp34 = tmp31 / tmp33
    tmp35 = tmp34 * tmp9
    tmp36 = tmp25 + tmp35
    tmp38 = tmp37 * tmp1
    tmp39 = tl_math.exp(tmp38)
    tmp40 = tmp39 / tmp5
    tmp41 = tmp40 / tmp7
    tmp42 = tmp41 * tmp9
    tmp45 = tmp42 / tmp44
    tmp46 = tmp45 * tmp9
    tmp47 = tmp36 + tmp46
    tl.store(out_ptr0 + x0, tmp47, xmask)


@triton.jit
def triton_poi_fused_div_17(in_ptr0, in_ptr1, in_ptr2, in_ptr3, in_ptr4,
    out_ptr0, out_ptr1, out_ptr2, xnumel, XBLOCK: tl.constexpr):
    xnumel = 16
    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 + (32 + x2), xmask)
    tmp3 = tl.load(in_ptr0 + (32 + 4 * x1), xmask, eviction_policy='evict_last'
        )
    tmp5 = tl.load(in_ptr0 + (33 + 4 * x1), xmask, eviction_policy='evict_last'
        )
    tmp8 = tl.load(in_ptr0 + (34 + 4 * x1), xmask, eviction_policy='evict_last'
        )
    tmp11 = tl.load(in_ptr0 + (35 + 4 * x1), xmask, eviction_policy=
        'evict_last')
    tmp20 = tl.load(in_ptr1 + 0)
    tmp21 = tl.broadcast_to(tmp20, [XBLOCK])
    tmp23 = tl.load(in_ptr2 + x0, xmask, eviction_policy='evict_last')
    tmp27 = tl.load(in_ptr3 + x1, xmask, eviction_policy='evict_last')
    tmp30 = tl.load(in_ptr4 + x0, xmask, eviction_policy='evict_last')
    tmp1 = 1.0
    tmp2 = tmp0 * tmp1
    tmp4 = tmp3 * tmp1
    tmp6 = tmp5 * tmp1
    tmp7 = triton_helpers.maximum(tmp4, tmp6)
    tmp9 = tmp8 * tmp1
    tmp10 = triton_helpers.maximum(tmp7, tmp9)
    tmp12 = tmp11 * tmp1
    tmp13 = triton_helpers.maximum(tmp10, tmp12)
    tmp14 = tmp2 - tmp13
    tmp15 = 10.0
    tmp16 = tmp14 * tmp15
    tmp17 = 20.0
    tmp18 = tmp0 * tmp17
    tmp19 = tl_math.exp(tmp18)
    tmp22 = tmp19 / tmp21
    tmp24 = tmp22 / tmp23
    tmp25 = 0.25
    tmp26 = tmp24 * tmp25
    tmp28 = tmp26 / tmp27
    tmp29 = tmp28 * tmp25
    tmp31 = tmp29 / tmp30
    tmp32 = tmp31 * tmp25
    tl.store(out_ptr0 + x2, tmp16, xmask)
    tl.store(out_ptr1 + x2, tmp16, xmask)
    tl.store(out_ptr2 + x2, tmp32, xmask)


@triton.jit
def triton_per_fused_sum_18(in_ptr0, out_ptr0, xnumel, rnumel, XBLOCK: tl.
    constexpr):
    RBLOCK: tl.constexpr = 16
    xoffset = tl.program_id(0) * XBLOCK
    xoffset + tl.arange(0, XBLOCK)[:, None]
    tl.full([XBLOCK, RBLOCK], True, tl.int1)
    rindex = tl.arange(0, RBLOCK)[None, :]
    tl.full([XBLOCK, RBLOCK], True, tl.int1)
    r0 = rindex
    tmp0 = tl.load(in_ptr0 + (48 + r0), None)
    tmp1 = 20.0
    tmp2 = tmp0 * tmp1
    tmp3 = tl_math.exp(tmp2)
    tmp4 = tl.broadcast_to(tmp3, [XBLOCK, RBLOCK])
    tmp6 = tl.sum(tmp4, 1)[:, None]
    tl.store(out_ptr0 + tl.full([XBLOCK, 1], 0, tl.int32), tmp6, None)


@triton.jit
def triton_poi_fused_sum_19(in_ptr0, in_ptr1, 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
    tmp0 = tl.load(in_ptr0 + (48 + x0), xmask)
    tmp4 = tl.load(in_ptr1 + 0)
    tmp5 = tl.broadcast_to(tmp4, [XBLOCK])
    tmp7 = tl.load(in_ptr0 + (52 + x0), xmask)
    tmp12 = tl.load(in_ptr0 + (56 + x0), xmask)
    tmp17 = tl.load(in_ptr0 + (60 + x0), xmask)
    tmp1 = 20.0
    tmp2 = tmp0 * tmp1
    tmp3 = tl_math.exp(tmp2)
    tmp6 = tmp3 / tmp5
    tmp8 = tmp7 * tmp1
    tmp9 = tl_math.exp(tmp8)
    tmp10 = tmp9 / tmp5
    tmp11 = tmp6 + tmp10
    tmp13 = tmp12 * tmp1
    tmp14 = tl_math.exp(tmp13)
    tmp15 = tmp14 / tmp5
    tmp16 = tmp11 + tmp15
    tmp18 = tmp17 * tmp1
    tmp19 = tl_math.exp(tmp18)
    tmp20 = tmp19 / tmp5
    tmp21 = tmp16 + tmp20
    tl.store(out_ptr0 + x0, tmp21, xmask)


@triton.jit
def triton_poi_fused_sum_20(in_ptr0, in_ptr1, in_ptr2, 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
    tmp0 = tl.load(in_ptr0 + (48 + 4 * x0), xmask, eviction_policy='evict_last'
        )
    tmp4 = tl.load(in_ptr1 + 0)
    tmp5 = tl.broadcast_to(tmp4, [XBLOCK])
    tmp7 = tl.load(in_ptr2 + 0)
    tmp8 = tl.broadcast_to(tmp7, [XBLOCK])
    tmp12 = tl.load(in_ptr0 + (49 + 4 * x0), xmask, eviction_policy=
        'evict_last')
    tmp16 = tl.load(in_ptr2 + 1)
    tmp17 = tl.broadcast_to(tmp16, [XBLOCK])
    tmp21 = tl.load(in_ptr0 + (50 + 4 * x0), xmask, eviction_policy=
        'evict_last')
    tmp25 = tl.load(in_ptr2 + 2)
    tmp26 = tl.broadcast_to(tmp25, [XBLOCK])
    tmp30 = tl.load(in_ptr0 + (51 + 4 * x0), xmask, eviction_policy=
        'evict_last')
    tmp34 = tl.load(in_ptr2 + 3)
    tmp35 = tl.broadcast_to(tmp34, [XBLOCK])
    tmp1 = 20.0
    tmp2 = tmp0 * tmp1
    tmp3 = tl_math.exp(tmp2)
    tmp6 = tmp3 / tmp5
    tmp9 = tmp6 / tmp8
    tmp10 = 0.25
    tmp11 = tmp9 * tmp10
    tmp13 = tmp12 * tmp1
    tmp14 = tl_math.exp(tmp13)
    tmp15 = tmp14 / tmp5
    tmp18 = tmp15 / tmp17
    tmp19 = tmp18 * tmp10
    tmp20 = tmp11 + tmp19
    tmp22 = tmp21 * tmp1
    tmp23 = tl_math.exp(tmp22)
    tmp24 = tmp23 / tmp5
    tmp27 = tmp24 / tmp26
    tmp28 = tmp27 * tmp10
    tmp29 = tmp20 + tmp28
    tmp31 = tmp30 * tmp1
    tmp32 = tl_math.exp(tmp31)
    tmp33 = tmp32 / tmp5
    tmp36 = tmp33 / tmp35
    tmp37 = tmp36 * tmp10
    tmp38 = tmp29 + tmp37
    tl.store(out_ptr0 + x0, tmp38, xmask)


@triton.jit
def triton_poi_fused_sum_21(in_ptr0, in_ptr1, in_ptr2, in_ptr3, 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
    tmp0 = tl.load(in_ptr0 + (48 + x0), xmask)
    tmp4 = tl.load(in_ptr1 + 0)
    tmp5 = tl.broadcast_to(tmp4, [XBLOCK])
    tmp7 = tl.load(in_ptr2 + x0, xmask)
    tmp11 = tl.load(in_ptr3 + 0)
    tmp12 = tl.broadcast_to(tmp11, [XBLOCK])
    tmp15 = tl.load(in_ptr0 + (52 + x0), xmask)
    tmp21 = tl.load(in_ptr3 + 1)
    tmp22 = tl.broadcast_to(tmp21, [XBLOCK])
    tmp26 = tl.load(in_ptr0 + (56 + x0), xmask)
    tmp32 = tl.load(in_ptr3 + 2)
    tmp33 = tl.broadcast_to(tmp32, [XBLOCK])
    tmp37 = tl.load(in_ptr0 + (60 + x0), xmask)
    tmp43 = tl.load(in_ptr3 + 3)
    tmp44 = tl.broadcast_to(tmp43, [XBLOCK])
    tmp1 = 20.0
    tmp2 = tmp0 * tmp1
    tmp3 = tl_math.exp(tmp2)
    tmp6 = tmp3 / tmp5
    tmp8 = tmp6 / tmp7
    tmp9 = 0.25
    tmp10 = tmp8 * tmp9
    tmp13 = tmp10 / tmp12
    tmp14 = tmp13 * tmp9
    tmp16 = tmp15 * tmp1
    tmp17 = tl_math.exp(tmp16)
    tmp18 = tmp17 / tmp5
    tmp19 = tmp18 / tmp7
    tmp20 = tmp19 * tmp9
    tmp23 = tmp20 / tmp22
    tmp24 = tmp23 * tmp9
    tmp25 = tmp14 + tmp24
    tmp27 = tmp26 * tmp1
    tmp28 = tl_math.exp(tmp27)
    tmp29 = tmp28 / tmp5
    tmp30 = tmp29 / tmp7
    tmp31 = tmp30 * tmp9
    tmp34 = tmp31 / tmp33
    tmp35 = tmp34 * tmp9
    tmp36 = tmp25 + tmp35
    tmp38 = tmp37 * tmp1
    tmp39 = tl_math.exp(tmp38)
    tmp40 = tmp39 / tmp5
    tmp41 = tmp40 / tmp7
    tmp42 = tmp41 * tmp9
    tmp45 = tmp42 / tmp44
    tmp46 = tmp45 * tmp9
    tmp47 = tmp36 + tmp46
    tl.store(out_ptr0 + x0, tmp47, xmask)


@triton.jit
def triton_poi_fused_div_22(in_ptr0, in_ptr1, in_ptr2, in_ptr3, in_ptr4,
    out_ptr0, out_ptr1, out_ptr2, out_ptr3, xnumel, XBLOCK: tl.constexpr):
    xnumel = 16
    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 + (48 + x2), xmask)
    tmp3 = tl.load(in_ptr0 + (48 + 4 * x1), xmask, eviction_policy='evict_last'
        )
    tmp5 = tl.load(in_ptr0 + (49 + 4 * x1), xmask, eviction_policy='evict_last'
        )
    tmp8 = tl.load(in_ptr0 + (50 + 4 * x1), xmask, eviction_policy='evict_last'
        )
    tmp11 = tl.load(in_ptr0 + (51 + 4 * x1), xmask, eviction_policy=
        'evict_last')
    tmp20 = tl.load(in_ptr1 + 0)
    tmp21 = tl.broadcast_to(tmp20, [XBLOCK])
    tmp23 = tl.load(in_ptr2 + x0, xmask, eviction_policy='evict_last')
    tmp27 = tl.load(in_ptr3 + x1, xmask, eviction_policy='evict_last')
    tmp30 = tl.load(in_ptr4 + x0, xmask, eviction_policy='evict_last')
    tmp1 = 1.0
    tmp2 = tmp0 * tmp1
    tmp4 = tmp3 * tmp1
    tmp6 = tmp5 * tmp1
    tmp7 = triton_helpers.maximum(tmp4, tmp6)
    tmp9 = tmp8 * tmp1
    tmp10 = triton_helpers.maximum(tmp7, tmp9)
    tmp12 = tmp11 * tmp1
    tmp13 = triton_helpers.maximum(tmp10, tmp12)
    tmp14 = tmp2 - tmp13
    tmp15 = 10.0
    tmp16 = tmp14 * tmp15
    tmp17 = 20.0
    tmp18 = tmp0 * tmp17
    tmp19 = tl_math.exp(tmp18)
    tmp22 = tmp19 / tmp21
    tmp24 = tmp22 / tmp23
    tmp25 = 0.25
    tmp26 = tmp24 * tmp25
    tmp28 = tmp26 / tmp27
    tmp29 = tmp28 * tmp25
    tmp31 = tmp29 / tmp30
    tmp32 = tmp31 * tmp25
    tl.store(out_ptr0 + x2, tmp16, xmask)
    tl.store(out_ptr1 + x2, tmp16, xmask)
    tl.store(out_ptr2 + x2, tmp16, xmask)
    tl.store(out_ptr3 + x2, tmp32, xmask)


@triton.jit
def triton_poi_fused_23(in_ptr0, out_ptr0, out_ptr1, out_ptr2, xnumel,
    XBLOCK: tl.constexpr):
    xnumel = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x2 = xindex
    x1 = xindex // 4
    tmp0 = tl.load(in_ptr0 + x2, xmask)
    tmp3 = tl.load(in_ptr0 + 4 * x1, xmask, eviction_policy='evict_last')
    tmp5 = tl.load(in_ptr0 + (1 + 4 * x1), xmask, eviction_policy='evict_last')
    tmp8 = tl.load(in_ptr0 + (2 + 4 * x1), xmask, eviction_policy='evict_last')
    tmp11 = tl.load(in_ptr0 + (3 + 4 * x1), xmask, eviction_policy='evict_last'
        )
    tmp1 = 1.0
    tmp2 = tmp0 * tmp1
    tmp4 = tmp3 * tmp1
    tmp6 = tmp5 * tmp1
    tmp7 = triton_helpers.maximum(tmp4, tmp6)
    tmp9 = tmp8 * tmp1
    tmp10 = triton_helpers.maximum(tmp7, tmp9)
    tmp12 = tmp11 * tmp1
    tmp13 = triton_helpers.maximum(tmp10, tmp12)
    tmp14 = tmp2 - tmp13
    tmp15 = 10.0
    tmp16 = tmp14 * tmp15
    tl.store(out_ptr0 + x2, tmp16, xmask)
    tl.store(out_ptr1 + x2, tmp16, xmask)
    tl.store(out_ptr2 + x2, tmp16, xmask)


@triton.jit
def triton_poi_fused_24(in_ptr0, out_ptr0, out_ptr1, out_ptr2, xnumel,
    XBLOCK: tl.constexpr):
    xnumel = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x2 = xindex
    x1 = xindex // 4
    tmp0 = tl.load(in_ptr0 + (16 + x2), xmask)
    tmp3 = tl.load(in_ptr0 + (16 + 4 * x1), xmask, eviction_policy='evict_last'
        )
    tmp5 = tl.load(in_ptr0 + (17 + 4 * x1), xmask, eviction_policy='evict_last'
        )
    tmp8 = tl.load(in_ptr0 + (18 + 4 * x1), xmask, eviction_policy='evict_last'
        )
    tmp11 = tl.load(in_ptr0 + (19 + 4 * x1), xmask, eviction_policy=
        'evict_last')
    tmp1 = 1.0
    tmp2 = tmp0 * tmp1
    tmp4 = tmp3 * tmp1
    tmp6 = tmp5 * tmp1
    tmp7 = triton_helpers.maximum(tmp4, tmp6)
    tmp9 = tmp8 * tmp1
    tmp10 = triton_helpers.maximum(tmp7, tmp9)
    tmp12 = tmp11 * tmp1
    tmp13 = triton_helpers.maximum(tmp10, tmp12)
    tmp14 = tmp2 - tmp13
    tmp15 = 10.0
    tmp16 = tmp14 * tmp15
    tl.store(out_ptr0 + x2, tmp16, xmask)
    tl.store(out_ptr1 + x2, tmp16, xmask)
    tl.store(out_ptr2 + x2, tmp16, xmask)


@triton.jit
def triton_poi_fused_25(in_ptr0, out_ptr0, out_ptr1, out_ptr2, xnumel,
    XBLOCK: tl.constexpr):
    xnumel = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x2 = xindex
    x1 = xindex // 4
    tmp0 = tl.load(in_ptr0 + (32 + x2), xmask)
    tmp3 = tl.load(in_ptr0 + (32 + 4 * x1), xmask, eviction_policy='evict_last'
        )
    tmp5 = tl.load(in_ptr0 + (33 + 4 * x1), xmask, eviction_policy='evict_last'
        )
    tmp8 = tl.load(in_ptr0 + (34 + 4 * x1), xmask, eviction_policy='evict_last'
        )
    tmp11 = tl.load(in_ptr0 + (35 + 4 * x1), xmask, eviction_policy=
        'evict_last')
    tmp1 = 1.0
    tmp2 = tmp0 * tmp1
    tmp4 = tmp3 * tmp1
    tmp6 = tmp5 * tmp1
    tmp7 = triton_helpers.maximum(tmp4, tmp6)
    tmp9 = tmp8 * tmp1
    tmp10 = triton_helpers.maximum(tmp7, tmp9)
    tmp12 = tmp11 * tmp1
    tmp13 = triton_helpers.maximum(tmp10, tmp12)
    tmp14 = tmp2 - tmp13
    tmp15 = 10.0
    tmp16 = tmp14 * tmp15
    tl.store(out_ptr0 + x2, tmp16, xmask)
    tl.store(out_ptr1 + x2, tmp16, xmask)
    tl.store(out_ptr2 + x2, tmp16, xmask)


@triton.jit
def triton_poi_fused_26(in_ptr0, out_ptr0, out_ptr1, out_ptr2, xnumel,
    XBLOCK: tl.constexpr):
    xnumel = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x2 = xindex
    x1 = xindex // 4
    tmp0 = tl.load(in_ptr0 + (48 + x2), xmask)
    tmp3 = tl.load(in_ptr0 + (48 + 4 * x1), xmask, eviction_policy='evict_last'
        )
    tmp5 = tl.load(in_ptr0 + (49 + 4 * x1), xmask, eviction_policy='evict_last'
        )
    tmp8 = tl.load(in_ptr0 + (50 + 4 * x1), xmask, eviction_policy='evict_last'
        )
    tmp11 = tl.load(in_ptr0 + (51 + 4 * x1), xmask, eviction_policy=
        'evict_last')
    tmp1 = 1.0
    tmp2 = tmp0 * tmp1
    tmp4 = tmp3 * tmp1
    tmp6 = tmp5 * tmp1
    tmp7 = triton_helpers.maximum(tmp4, tmp6)
    tmp9 = tmp8 * tmp1
    tmp10 = triton_helpers.maximum(tmp7, tmp9)
    tmp12 = tmp11 * tmp1
    tmp13 = triton_helpers.maximum(tmp10, tmp12)
    tmp14 = tmp2 - tmp13
    tmp15 = 10.0
    tmp16 = tmp14 * tmp15
    tl.store(out_ptr0 + x2, tmp16, xmask)
    tl.store(out_ptr1 + x2, tmp16, xmask)
    tl.store(out_ptr2 + x2, tmp16, xmask)


@triton.jit
def triton_poi_fused_27(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x2 = xindex
    x1 = xindex // 4
    tmp0 = tl.load(in_ptr0 + x2, xmask)
    tmp3 = tl.load(in_ptr0 + 4 * x1, xmask, eviction_policy='evict_last')
    tmp5 = tl.load(in_ptr0 + (1 + 4 * x1), xmask, eviction_policy='evict_last')
    tmp8 = tl.load(in_ptr0 + (2 + 4 * x1), xmask, eviction_policy='evict_last')
    tmp11 = tl.load(in_ptr0 + (3 + 4 * x1), xmask, eviction_policy='evict_last'
        )
    tmp1 = 1.0
    tmp2 = tmp0 * tmp1
    tmp4 = tmp3 * tmp1
    tmp6 = tmp5 * tmp1
    tmp7 = triton_helpers.maximum(tmp4, tmp6)
    tmp9 = tmp8 * tmp1
    tmp10 = triton_helpers.maximum(tmp7, tmp9)
    tmp12 = tmp11 * tmp1
    tmp13 = triton_helpers.maximum(tmp10, tmp12)
    tmp14 = tmp2 - tmp13
    tmp15 = 10.0
    tmp16 = tmp14 * tmp15
    tl.store(out_ptr0 + x2, tmp16, xmask)


@triton.jit
def triton_poi_fused_28(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x2 = xindex
    x1 = xindex // 4
    tmp0 = tl.load(in_ptr0 + (16 + x2), xmask)
    tmp3 = tl.load(in_ptr0 + (16 + 4 * x1), xmask, eviction_policy='evict_last'
        )
    tmp5 = tl.load(in_ptr0 + (17 + 4 * x1), xmask, eviction_policy='evict_last'
        )
    tmp8 = tl.load(in_ptr0 + (18 + 4 * x1), xmask, eviction_policy='evict_last'
        )
    tmp11 = tl.load(in_ptr0 + (19 + 4 * x1), xmask, eviction_policy=
        'evict_last')
    tmp1 = 1.0
    tmp2 = tmp0 * tmp1
    tmp4 = tmp3 * tmp1
    tmp6 = tmp5 * tmp1
    tmp7 = triton_helpers.maximum(tmp4, tmp6)
    tmp9 = tmp8 * tmp1
    tmp10 = triton_helpers.maximum(tmp7, tmp9)
    tmp12 = tmp11 * tmp1
    tmp13 = triton_helpers.maximum(tmp10, tmp12)
    tmp14 = tmp2 - tmp13
    tmp15 = 10.0
    tmp16 = tmp14 * tmp15
    tl.store(out_ptr0 + x2, tmp16, xmask)


@triton.jit
def triton_poi_fused_29(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x2 = xindex
    x1 = xindex // 4
    tmp0 = tl.load(in_ptr0 + (32 + x2), xmask)
    tmp3 = tl.load(in_ptr0 + (32 + 4 * x1), xmask, eviction_policy='evict_last'
        )
    tmp5 = tl.load(in_ptr0 + (33 + 4 * x1), xmask, eviction_policy='evict_last'
        )
    tmp8 = tl.load(in_ptr0 + (34 + 4 * x1), xmask, eviction_policy='evict_last'
        )
    tmp11 = tl.load(in_ptr0 + (35 + 4 * x1), xmask, eviction_policy=
        'evict_last')
    tmp1 = 1.0
    tmp2 = tmp0 * tmp1
    tmp4 = tmp3 * tmp1
    tmp6 = tmp5 * tmp1
    tmp7 = triton_helpers.maximum(tmp4, tmp6)
    tmp9 = tmp8 * tmp1
    tmp10 = triton_helpers.maximum(tmp7, tmp9)
    tmp12 = tmp11 * tmp1
    tmp13 = triton_helpers.maximum(tmp10, tmp12)
    tmp14 = tmp2 - tmp13
    tmp15 = 10.0
    tmp16 = tmp14 * tmp15
    tl.store(out_ptr0 + x2, tmp16, xmask)


@triton.jit
def triton_poi_fused_30(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x2 = xindex
    x1 = xindex // 4
    tmp0 = tl.load(in_ptr0 + (48 + x2), xmask)
    tmp3 = tl.load(in_ptr0 + (48 + 4 * x1), xmask, eviction_policy='evict_last'
        )
    tmp5 = tl.load(in_ptr0 + (49 + 4 * x1), xmask, eviction_policy='evict_last'
        )
    tmp8 = tl.load(in_ptr0 + (50 + 4 * x1), xmask, eviction_policy='evict_last'
        )
    tmp11 = tl.load(in_ptr0 + (51 + 4 * x1), xmask, eviction_policy=
        'evict_last')
    tmp1 = 1.0
    tmp2 = tmp0 * tmp1
    tmp4 = tmp3 * tmp1
    tmp6 = tmp5 * tmp1
    tmp7 = triton_helpers.maximum(tmp4, tmp6)
    tmp9 = tmp8 * tmp1
    tmp10 = triton_helpers.maximum(tmp7, tmp9)
    tmp12 = tmp11 * tmp1
    tmp13 = triton_helpers.maximum(tmp10, tmp12)
    tmp14 = tmp2 - tmp13
    tmp15 = 10.0
    tmp16 = tmp14 * tmp15
    tl.store(out_ptr0 + x2, tmp16, xmask)


@triton.jit
def triton_per_fused__log_softmax_add_div_mean_mul_neg_sum_31(in_out_ptr0,
    in_ptr0, in_ptr1, in_ptr2, in_ptr3, in_ptr4, in_ptr5, in_ptr6, in_ptr7,
    in_ptr8, in_ptr9, in_ptr10, in_ptr11, in_ptr12, in_ptr13, in_ptr14,
    in_ptr15, in_ptr16, in_ptr17, in_ptr18, in_ptr19, in_ptr20, in_ptr21,
    in_ptr22, in_ptr23, in_ptr24, in_ptr25, in_ptr26, in_ptr27, in_ptr28,
    in_ptr29, in_ptr30, in_ptr31, xnumel, rnumel, XBLOCK: tl.constexpr):
    RBLOCK: tl.constexpr = 4
    xoffset = tl.program_id(0) * XBLOCK
    xoffset + tl.arange(0, XBLOCK)[:, None]
    tl.full([XBLOCK, RBLOCK], True, tl.int1)
    rindex = tl.arange(0, RBLOCK)[None, :]
    tl.full([XBLOCK, RBLOCK], True, tl.int1)
    r0 = rindex
    tmp0 = tl.load(in_ptr0 + 4 * r0, None, eviction_policy='evict_last')
    tmp1 = tl.load(in_ptr1 + 4 * r0, None, eviction_policy='evict_last')
    tmp3 = tl.load(in_ptr1 + (1 + 4 * r0), None, eviction_policy='evict_last')
    tmp6 = tl.load(in_ptr1 + (2 + 4 * r0), None, eviction_policy='evict_last')
    tmp9 = tl.load(in_ptr1 + (3 + 4 * r0), None, eviction_policy='evict_last')
    tmp15 = tl.load(in_ptr0 + (1 + 4 * r0), None, eviction_policy='evict_last')
    tmp19 = tl.load(in_ptr0 + (2 + 4 * r0), None, eviction_policy='evict_last')
    tmp23 = tl.load(in_ptr0 + (3 + 4 * r0), None, eviction_policy='evict_last')
    tmp30 = tl.load(in_ptr2 + 4 * r0, None, eviction_policy='evict_last')
    tmp32 = tl.load(in_ptr2 + (1 + 4 * r0), None, eviction_policy='evict_last')
    tmp35 = tl.load(in_ptr2 + (2 + 4 * r0), None, eviction_policy='evict_last')
    tmp38 = tl.load(in_ptr2 + (3 + 4 * r0), None, eviction_policy='evict_last')
    tmp56 = tl.load(in_ptr3 + 4 * r0, None, eviction_policy='evict_last')
    tmp58 = tl.load(in_ptr3 + (1 + 4 * r0), None, eviction_policy='evict_last')
    tmp61 = tl.load(in_ptr3 + (2 + 4 * r0), None, eviction_policy='evict_last')
    tmp64 = tl.load(in_ptr3 + (3 + 4 * r0), None, eviction_policy='evict_last')
    tmp82 = tl.load(in_ptr4 + 4 * r0, None, eviction_policy='evict_last')
    tmp84 = tl.load(in_ptr4 + (1 + 4 * r0), None, eviction_policy='evict_last')
    tmp87 = tl.load(in_ptr4 + (2 + 4 * r0), None, eviction_policy='evict_last')
    tmp90 = tl.load(in_ptr4 + (3 + 4 * r0), None, eviction_policy='evict_last')
    tmp108 = tl.load(in_ptr5 + 4 * r0, None, eviction_policy='evict_last')
    tmp110 = tl.load(in_ptr5 + (1 + 4 * r0), None, eviction_policy='evict_last'
        )
    tmp113 = tl.load(in_ptr5 + (2 + 4 * r0), None, eviction_policy='evict_last'
        )
    tmp116 = tl.load(in_ptr5 + (3 + 4 * r0), None, eviction_policy='evict_last'
        )
    tmp134 = tl.load(in_ptr6 + 4 * r0, None, eviction_policy='evict_last')
    tmp136 = tl.load(in_ptr6 + (1 + 4 * r0), None, eviction_policy='evict_last'
        )
    tmp139 = tl.load(in_ptr6 + (2 + 4 * r0), None, eviction_policy='evict_last'
        )
    tmp142 = tl.load(in_ptr6 + (3 + 4 * r0), None, eviction_policy='evict_last'
        )
    tmp160 = tl.load(in_ptr7 + 4 * r0, None, eviction_policy='evict_last')
    tmp162 = tl.load(in_ptr7 + (1 + 4 * r0), None, eviction_policy='evict_last'
        )
    tmp165 = tl.load(in_ptr7 + (2 + 4 * r0), None, eviction_policy='evict_last'
        )
    tmp168 = tl.load(in_ptr7 + (3 + 4 * r0), None, eviction_policy='evict_last'
        )
    tmp186 = tl.load(in_ptr8 + 4 * r0, None, eviction_policy='evict_last')
    tmp187 = tl.load(in_ptr9 + 4 * r0, None, eviction_policy='evict_last')
    tmp189 = tl.load(in_ptr9 + (1 + 4 * r0), None, eviction_policy='evict_last'
        )
    tmp192 = tl.load(in_ptr9 + (2 + 4 * r0), None, eviction_policy='evict_last'
        )
    tmp195 = tl.load(in_ptr9 + (3 + 4 * r0), None, eviction_policy='evict_last'
        )
    tmp201 = tl.load(in_ptr8 + (1 + 4 * r0), None, eviction_policy='evict_last'
        )
    tmp205 = tl.load(in_ptr8 + (2 + 4 * r0), None, eviction_policy='evict_last'
        )
    tmp209 = tl.load(in_ptr8 + (3 + 4 * r0), None, eviction_policy='evict_last'
        )
    tmp216 = tl.load(in_ptr10 + 4 * r0, None, eviction_policy='evict_last')
    tmp218 = tl.load(in_ptr10 + (1 + 4 * r0), None, eviction_policy=
        'evict_last')
    tmp221 = tl.load(in_ptr10 + (2 + 4 * r0), None, eviction_policy=
        'evict_last')
    tmp224 = tl.load(in_ptr10 + (3 + 4 * r0), None, eviction_policy=
        'evict_last')
    tmp242 = tl.load(in_ptr11 + 4 * r0, None, eviction_policy='evict_last')
    tmp244 = tl.load(in_ptr11 + (1 + 4 * r0), None, eviction_policy=
        'evict_last')
    tmp247 = tl.load(in_ptr11 + (2 + 4 * r0), None, eviction_policy=
        'evict_last')
    tmp250 = tl.load(in_ptr11 + (3 + 4 * r0), None, eviction_policy=
        'evict_last')
    tmp268 = tl.load(in_ptr12 + 4 * r0, None, eviction_policy='evict_last')
    tmp270 = tl.load(in_ptr12 + (1 + 4 * r0), None, eviction_policy=
        'evict_last')
    tmp273 = tl.load(in_ptr12 + (2 + 4 * r0), None, eviction_policy=
        'evict_last')
    tmp276 = tl.load(in_ptr12 + (3 + 4 * r0), None, eviction_policy=
        'evict_last')
    tmp294 = tl.load(in_ptr13 + 4 * r0, None, eviction_policy='evict_last')
    tmp296 = tl.load(in_ptr13 + (1 + 4 * r0), None, eviction_policy=
        'evict_last')
    tmp299 = tl.load(in_ptr13 + (2 + 4 * r0), None, eviction_policy=
        'evict_last')
    tmp302 = tl.load(in_ptr13 + (3 + 4 * r0), None, eviction_policy=
        'evict_last')
    tmp320 = tl.load(in_ptr14 + 4 * r0, None, eviction_policy='evict_last')
    tmp322 = tl.load(in_ptr14 + (1 + 4 * r0), None, eviction_policy=
        'evict_last')
    tmp325 = tl.load(in_ptr14 + (2 + 4 * r0), None, eviction_policy=
        'evict_last')
    tmp328 = tl.load(in_ptr14 + (3 + 4 * r0), None, eviction_policy=
        'evict_last')
    tmp346 = tl.load(in_ptr15 + 4 * r0, None, eviction_policy='evict_last')
    tmp348 = tl.load(in_ptr15 + (1 + 4 * r0), None, eviction_policy=
        'evict_last')
    tmp351 = tl.load(in_ptr15 + (2 + 4 * r0), None, eviction_policy=
        'evict_last')
    tmp354 = tl.load(in_ptr15 + (3 + 4 * r0), None, eviction_policy=
        'evict_last')
    tmp372 = tl.load(in_ptr16 + 4 * r0, None, eviction_policy='evict_last')
    tmp373 = tl.load(in_ptr17 + 4 * r0, None, eviction_policy='evict_last')
    tmp375 = tl.load(in_ptr17 + (1 + 4 * r0), None, eviction_policy=
        'evict_last')
    tmp378 = tl.load(in_ptr17 + (2 + 4 * r0), None, eviction_policy=
        'evict_last')
    tmp381 = tl.load(in_ptr17 + (3 + 4 * r0), None, eviction_policy=
        'evict_last')
    tmp387 = tl.load(in_ptr16 + (1 + 4 * r0), None, eviction_policy=
        'evict_last')
    tmp391 = tl.load(in_ptr16 + (2 + 4 * r0), None, eviction_policy=
        'evict_last')
    tmp395 = tl.load(in_ptr16 + (3 + 4 * r0), None, eviction_policy=
        'evict_last')
    tmp402 = tl.load(in_ptr18 + 4 * r0, None, eviction_policy='evict_last')
    tmp404 = tl.load(in_ptr18 + (1 + 4 * r0), None, eviction_policy=
        'evict_last')
    tmp407 = tl.load(in_ptr18 + (2 + 4 * r0), None, eviction_policy=
        'evict_last')
    tmp410 = tl.load(in_ptr18 + (3 + 4 * r0), None, eviction_policy=
        'evict_last')
    tmp428 = tl.load(in_ptr19 + 4 * r0, None, eviction_policy='evict_last')
    tmp430 = tl.load(in_ptr19 + (1 + 4 * r0), None, eviction_policy=
        'evict_last')
    tmp433 = tl.load(in_ptr19 + (2 + 4 * r0), None, eviction_policy=
        'evict_last')
    tmp436 = tl.load(in_ptr19 + (3 + 4 * r0), None, eviction_policy=
        'evict_last')
    tmp454 = tl.load(in_ptr20 + 4 * r0, None, eviction_policy='evict_last')
    tmp456 = tl.load(in_ptr20 + (1 + 4 * r0), None, eviction_policy=
        'evict_last')
    tmp459 = tl.load(in_ptr20 + (2 + 4 * r0), None, eviction_policy=
        'evict_last')
    tmp462 = tl.load(in_ptr20 + (3 + 4 * r0), None, eviction_policy=
        'evict_last')
    tmp480 = tl.load(in_ptr21 + 4 * r0, None, eviction_policy='evict_last')
    tmp482 = tl.load(in_ptr21 + (1 + 4 * r0), None, eviction_policy=
        'evict_last')
    tmp485 = tl.load(in_ptr21 + (2 + 4 * r0), None, eviction_policy=
        'evict_last')
    tmp488 = tl.load(in_ptr21 + (3 + 4 * r0), None, eviction_policy=
        'evict_last')
    tmp506 = tl.load(in_ptr22 + 4 * r0, None, eviction_policy='evict_last')
    tmp508 = tl.load(in_ptr22 + (1 + 4 * r0), None, eviction_policy=
        'evict_last')
    tmp511 = tl.load(in_ptr22 + (2 + 4 * r0), None, eviction_policy=
        'evict_last')
    tmp514 = tl.load(in_ptr22 + (3 + 4 * r0), None, eviction_policy=
        'evict_last')
    tmp532 = tl.load(in_ptr23 + 4 * r0, None, eviction_policy='evict_last')
    tmp534 = tl.load(in_ptr23 + (1 + 4 * r0), None, eviction_policy=
        'evict_last')
    tmp537 = tl.load(in_ptr23 + (2 + 4 * r0), None, eviction_policy=
        'evict_last')
    tmp540 = tl.load(in_ptr23 + (3 + 4 * r0), None, eviction_policy=
        'evict_last')
    tmp558 = tl.load(in_ptr24 + 4 * r0, None, eviction_policy='evict_last')
    tmp559 = tl.load(in_ptr25 + 4 * r0, None, eviction_policy='evict_last')
    tmp561 = tl.load(in_ptr25 + (1 + 4 * r0), None, eviction_policy=
        'evict_last')
    tmp564 = tl.load(in_ptr25 + (2 + 4 * r0), None, eviction_policy=
        'evict_last')
    tmp567 = tl.load(in_ptr25 + (3 + 4 * r0), None, eviction_policy=
        'evict_last')
    tmp573 = tl.load(in_ptr24 + (1 + 4 * r0), None, eviction_policy=
        'evict_last')
    tmp577 = tl.load(in_ptr24 + (2 + 4 * r0), None, eviction_policy=
        'evict_last')
    tmp581 = tl.load(in_ptr24 + (3 + 4 * r0), None, eviction_policy=
        'evict_last')
    tmp588 = tl.load(in_ptr26 + 4 * r0, None, eviction_policy='evict_last')
    tmp590 = tl.load(in_ptr26 + (1 + 4 * r0), None, eviction_policy=
        'evict_last')
    tmp593 = tl.load(in_ptr26 + (2 + 4 * r0), None, eviction_policy=
        'evict_last')
    tmp596 = tl.load(in_ptr26 + (3 + 4 * r0), None, eviction_policy=
        'evict_last')
    tmp614 = tl.load(in_ptr27 + 4 * r0, None, eviction_policy='evict_last')
    tmp616 = tl.load(in_ptr27 + (1 + 4 * r0), None, eviction_policy=
        'evict_last')
    tmp619 = tl.load(in_ptr27 + (2 + 4 * r0), None, eviction_policy=
        'evict_last')
    tmp622 = tl.load(in_ptr27 + (3 + 4 * r0), None, eviction_policy=
        'evict_last')
    tmp640 = tl.load(in_ptr28 + 4 * r0, None, eviction_policy='evict_last')
    tmp642 = tl.load(in_ptr28 + (1 + 4 * r0), None, eviction_policy=
        'evict_last')
    tmp645 = tl.load(in_ptr28 + (2 + 4 * r0), None, eviction_policy=
        'evict_last')
    tmp648 = tl.load(in_ptr28 + (3 + 4 * r0), None, eviction_policy=
        'evict_last')
    tmp666 = tl.load(in_ptr29 + 4 * r0, None, eviction_policy='evict_last')
    tmp668 = tl.load(in_ptr29 + (1 + 4 * r0), None, eviction_policy=
        'evict_last')
    tmp671 = tl.load(in_ptr29 + (2 + 4 * r0), None, eviction_policy=
        'evict_last')
    tmp674 = tl.load(in_ptr29 + (3 + 4 * r0), None, eviction_policy=
        'evict_last')
    tmp692 = tl.load(in_ptr30 + 4 * r0, None, eviction_policy='evict_last')
    tmp694 = tl.load(in_ptr30 + (1 + 4 * r0), None, eviction_policy=
        'evict_last')
    tmp697 = tl.load(in_ptr30 + (2 + 4 * r0), None, eviction_policy=
        'evict_last')
    tmp700 = tl.load(in_ptr30 + (3 + 4 * r0), None, eviction_policy=
        'evict_last')
    tmp718 = tl.load(in_ptr31 + 4 * r0, None, eviction_policy='evict_last')
    tmp720 = tl.load(in_ptr31 + (1 + 4 * r0), None, eviction_policy=
        'evict_last')
    tmp723 = tl.load(in_ptr31 + (2 + 4 * r0), None, eviction_policy=
        'evict_last')
    tmp726 = tl.load(in_ptr31 + (3 + 4 * r0), None, eviction_policy=
        'evict_last')
    tmp2 = tl_math.exp(tmp1)
    tmp4 = tl_math.exp(tmp3)
    tmp5 = tmp2 + tmp4
    tmp7 = tl_math.exp(tmp6)
    tmp8 = tmp5 + tmp7
    tmp10 = tl_math.exp(tmp9)
    tmp11 = tmp8 + tmp10
    tmp12 = tl_math.log(tmp11)
    tmp13 = tmp1 - tmp12
    tmp14 = tmp0 * tmp13
    tmp16 = tmp3 - tmp12
    tmp17 = tmp15 * tmp16
    tmp18 = tmp14 + tmp17
    tmp20 = tmp6 - tmp12
    tmp21 = tmp19 * tmp20
    tmp22 = tmp18 + tmp21
    tmp24 = tmp9 - tmp12
    tmp25 = tmp23 * tmp24
    tmp26 = tmp22 + tmp25
    tmp27 = tl.broadcast_to(tmp26, [XBLOCK, RBLOCK])
    tmp29 = tl.sum(tmp27, 1)[:, None]
    tmp31 = tl_math.exp(tmp30)
    tmp33 = tl_math.exp(tmp32)
    tmp34 = tmp31 + tmp33
    tmp36 = tl_math.exp(tmp35)
    tmp37 = tmp34 + tmp36
    tmp39 = tl_math.exp(tmp38)
    tmp40 = tmp37 + tmp39
    tmp41 = tl_math.log(tmp40)
    tmp42 = tmp30 - tmp41
    tmp43 = tmp0 * tmp42
    tmp44 = tmp32 - tmp41
    tmp45 = tmp15 * tmp44
    tmp46 = tmp43 + tmp45
    tmp47 = tmp35 - tmp41
    tmp48 = tmp19 * tmp47
    tmp49 = tmp46 + tmp48
    tmp50 = tmp38 - tmp41
    tmp51 = tmp23 * tmp50
    tmp52 = tmp49 + tmp51
    tmp53 = tl.broadcast_to(tmp52, [XBLOCK, RBLOCK])
    tmp55 = tl.sum(tmp53, 1)[:, None]
    tmp57 = tl_math.exp(tmp56)
    tmp59 = tl_math.exp(tmp58)
    tmp60 = tmp57 + tmp59
    tmp62 = tl_math.exp(tmp61)
    tmp63 = tmp60 + tmp62
    tmp65 = tl_math.exp(tmp64)
    tmp66 = tmp63 + tmp65
    tmp67 = tl_math.log(tmp66)
    tmp68 = tmp56 - tmp67
    tmp69 = tmp0 * tmp68
    tmp70 = tmp58 - tmp67
    tmp71 = tmp15 * tmp70
    tmp72 = tmp69 + tmp71
    tmp73 = tmp61 - tmp67
    tmp74 = tmp19 * tmp73
    tmp75 = tmp72 + tmp74
    tmp76 = tmp64 - tmp67
    tmp77 = tmp23 * tmp76
    tmp78 = tmp75 + tmp77
    tmp79 = tl.broadcast_to(tmp78, [XBLOCK, RBLOCK])
    tmp81 = tl.sum(tmp79, 1)[:, None]
    tmp83 = tl_math.exp(tmp82)
    tmp85 = tl_math.exp(tmp84)
    tmp86 = tmp83 + tmp85
    tmp88 = tl_math.exp(tmp87)
    tmp89 = tmp86 + tmp88
    tmp91 = tl_math.exp(tmp90)
    tmp92 = tmp89 + tmp91
    tmp93 = tl_math.log(tmp92)
    tmp94 = tmp82 - tmp93
    tmp95 = tmp0 * tmp94
    tmp96 = tmp84 - tmp93
    tmp97 = tmp15 * tmp96
    tmp98 = tmp95 + tmp97
    tmp99 = tmp87 - tmp93
    tmp100 = tmp19 * tmp99
    tmp101 = tmp98 + tmp100
    tmp102 = tmp90 - tmp93
    tmp103 = tmp23 * tmp102
    tmp104 = tmp101 + tmp103
    tmp105 = tl.broadcast_to(tmp104, [XBLOCK, RBLOCK])
    tmp107 = tl.sum(tmp105, 1)[:, None]
    tmp109 = tl_math.exp(tmp108)
    tmp111 = tl_math.exp(tmp110)
    tmp112 = tmp109 + tmp111
    tmp114 = tl_math.exp(tmp113)
    tmp115 = tmp112 + tmp114
    tmp117 = tl_math.exp(tmp116)
    tmp118 = tmp115 + tmp117
    tmp119 = tl_math.log(tmp118)
    tmp120 = tmp108 - tmp119
    tmp121 = tmp0 * tmp120
    tmp122 = tmp110 - tmp119
    tmp123 = tmp15 * tmp122
    tmp124 = tmp121 + tmp123
    tmp125 = tmp113 - tmp119
    tmp126 = tmp19 * tmp125
    tmp127 = tmp124 + tmp126
    tmp128 = tmp116 - tmp119
    tmp129 = tmp23 * tmp128
    tmp130 = tmp127 + tmp129
    tmp131 = tl.broadcast_to(tmp130, [XBLOCK, RBLOCK])
    tmp133 = tl.sum(tmp131, 1)[:, None]
    tmp135 = tl_math.exp(tmp134)
    tmp137 = tl_math.exp(tmp136)
    tmp138 = tmp135 + tmp137
    tmp140 = tl_math.exp(tmp139)
    tmp141 = tmp138 + tmp140
    tmp143 = tl_math.exp(tmp142)
    tmp144 = tmp141 + tmp143
    tmp145 = tl_math.log(tmp144)
    tmp146 = tmp134 - tmp145
    tmp147 = tmp0 * tmp146
    tmp148 = tmp136 - tmp145
    tmp149 = tmp15 * tmp148
    tmp150 = tmp147 + tmp149
    tmp151 = tmp139 - tmp145
    tmp152 = tmp19 * tmp151
    tmp153 = tmp150 + tmp152
    tmp154 = tmp142 - tmp145
    tmp155 = tmp23 * tmp154
    tmp156 = tmp153 + tmp155
    tmp157 = tl.broadcast_to(tmp156, [XBLOCK, RBLOCK])
    tmp159 = tl.sum(tmp157, 1)[:, None]
    tmp161 = tl_math.exp(tmp160)
    tmp163 = tl_math.exp(tmp162)
    tmp164 = tmp161 + tmp163
    tmp166 = tl_math.exp(tmp165)
    tmp167 = tmp164 + tmp166
    tmp169 = tl_math.exp(tmp168)
    tmp170 = tmp167 + tmp169
    tmp171 = tl_math.log(tmp170)
    tmp172 = tmp160 - tmp171
    tmp173 = tmp0 * tmp172
    tmp174 = tmp162 - tmp171
    tmp175 = tmp15 * tmp174
    tmp176 = tmp173 + tmp175
    tmp177 = tmp165 - tmp171
    tmp178 = tmp19 * tmp177
    tmp179 = tmp176 + tmp178
    tmp180 = tmp168 - tmp171
    tmp181 = tmp23 * tmp180
    tmp182 = tmp179 + tmp181
    tmp183 = tl.broadcast_to(tmp182, [XBLOCK, RBLOCK])
    tmp185 = tl.sum(tmp183, 1)[:, None]
    tmp188 = tl_math.exp(tmp187)
    tmp190 = tl_math.exp(tmp189)
    tmp191 = tmp188 + tmp190
    tmp193 = tl_math.exp(tmp192)
    tmp194 = tmp191 + tmp193
    tmp196 = tl_math.exp(tmp195)
    tmp197 = tmp194 + tmp196
    tmp198 = tl_math.log(tmp197)
    tmp199 = tmp187 - tmp198
    tmp200 = tmp186 * tmp199
    tmp202 = tmp189 - tmp198
    tmp203 = tmp201 * tmp202
    tmp204 = tmp200 + tmp203
    tmp206 = tmp192 - tmp198
    tmp207 = tmp205 * tmp206
    tmp208 = tmp204 + tmp207
    tmp210 = tmp195 - tmp198
    tmp211 = tmp209 * tmp210
    tmp212 = tmp208 + tmp211
    tmp213 = tl.broadcast_to(tmp212, [XBLOCK, RBLOCK])
    tmp215 = tl.sum(tmp213, 1)[:, None]
    tmp217 = tl_math.exp(tmp216)
    tmp219 = tl_math.exp(tmp218)
    tmp220 = tmp217 + tmp219
    tmp222 = tl_math.exp(tmp221)
    tmp223 = tmp220 + tmp222
    tmp225 = tl_math.exp(tmp224)
    tmp226 = tmp223 + tmp225
    tmp227 = tl_math.log(tmp226)
    tmp228 = tmp216 - tmp227
    tmp229 = tmp186 * tmp228
    tmp230 = tmp218 - tmp227
    tmp231 = tmp201 * tmp230
    tmp232 = tmp229 + tmp231
    tmp233 = tmp221 - tmp227
    tmp234 = tmp205 * tmp233
    tmp235 = tmp232 + tmp234
    tmp236 = tmp224 - tmp227
    tmp237 = tmp209 * tmp236
    tmp238 = tmp235 + tmp237
    tmp239 = tl.broadcast_to(tmp238, [XBLOCK, RBLOCK])
    tmp241 = tl.sum(tmp239, 1)[:, None]
    tmp243 = tl_math.exp(tmp242)
    tmp245 = tl_math.exp(tmp244)
    tmp246 = tmp243 + tmp245
    tmp248 = tl_math.exp(tmp247)
    tmp249 = tmp246 + tmp248
    tmp251 = tl_math.exp(tmp250)
    tmp252 = tmp249 + tmp251
    tmp253 = tl_math.log(tmp252)
    tmp254 = tmp242 - tmp253
    tmp255 = tmp186 * tmp254
    tmp256 = tmp244 - tmp253
    tmp257 = tmp201 * tmp256
    tmp258 = tmp255 + tmp257
    tmp259 = tmp247 - tmp253
    tmp260 = tmp205 * tmp259
    tmp261 = tmp258 + tmp260
    tmp262 = tmp250 - tmp253
    tmp263 = tmp209 * tmp262
    tmp264 = tmp261 + tmp263
    tmp265 = tl.broadcast_to(tmp264, [XBLOCK, RBLOCK])
    tmp267 = tl.sum(tmp265, 1)[:, None]
    tmp269 = tl_math.exp(tmp268)
    tmp271 = tl_math.exp(tmp270)
    tmp272 = tmp269 + tmp271
    tmp274 = tl_math.exp(tmp273)
    tmp275 = tmp272 + tmp274
    tmp277 = tl_math.exp(tmp276)
    tmp278 = tmp275 + tmp277
    tmp279 = tl_math.log(tmp278)
    tmp280 = tmp268 - tmp279
    tmp281 = tmp186 * tmp280
    tmp282 = tmp270 - tmp279
    tmp283 = tmp201 * tmp282
    tmp284 = tmp281 + tmp283
    tmp285 = tmp273 - tmp279
    tmp286 = tmp205 * tmp285
    tmp287 = tmp284 + tmp286
    tmp288 = tmp276 - tmp279
    tmp289 = tmp209 * tmp288
    tmp290 = tmp287 + tmp289
    tmp291 = tl.broadcast_to(tmp290, [XBLOCK, RBLOCK])
    tmp293 = tl.sum(tmp291, 1)[:, None]
    tmp295 = tl_math.exp(tmp294)
    tmp297 = tl_math.exp(tmp296)
    tmp298 = tmp295 + tmp297
    tmp300 = tl_math.exp(tmp299)
    tmp301 = tmp298 + tmp300
    tmp303 = tl_math.exp(tmp302)
    tmp304 = tmp301 + tmp303
    tmp305 = tl_math.log(tmp304)
    tmp306 = tmp294 - tmp305
    tmp307 = tmp186 * tmp306
    tmp308 = tmp296 - tmp305
    tmp309 = tmp201 * tmp308
    tmp310 = tmp307 + tmp309
    tmp311 = tmp299 - tmp305
    tmp312 = tmp205 * tmp311
    tmp313 = tmp310 + tmp312
    tmp314 = tmp302 - tmp305
    tmp315 = tmp209 * tmp314
    tmp316 = tmp313 + tmp315
    tmp317 = tl.broadcast_to(tmp316, [XBLOCK, RBLOCK])
    tmp319 = tl.sum(tmp317, 1)[:, None]
    tmp321 = tl_math.exp(tmp320)
    tmp323 = tl_math.exp(tmp322)
    tmp324 = tmp321 + tmp323
    tmp326 = tl_math.exp(tmp325)
    tmp327 = tmp324 + tmp326
    tmp329 = tl_math.exp(tmp328)
    tmp330 = tmp327 + tmp329
    tmp331 = tl_math.log(tmp330)
    tmp332 = tmp320 - tmp331
    tmp333 = tmp186 * tmp332
    tmp334 = tmp322 - tmp331
    tmp335 = tmp201 * tmp334
    tmp336 = tmp333 + tmp335
    tmp337 = tmp325 - tmp331
    tmp338 = tmp205 * tmp337
    tmp339 = tmp336 + tmp338
    tmp340 = tmp328 - tmp331
    tmp341 = tmp209 * tmp340
    tmp342 = tmp339 + tmp341
    tmp343 = tl.broadcast_to(tmp342, [XBLOCK, RBLOCK])
    tmp345 = tl.sum(tmp343, 1)[:, None]
    tmp347 = tl_math.exp(tmp346)
    tmp349 = tl_math.exp(tmp348)
    tmp350 = tmp347 + tmp349
    tmp352 = tl_math.exp(tmp351)
    tmp353 = tmp350 + tmp352
    tmp355 = tl_math.exp(tmp354)
    tmp356 = tmp353 + tmp355
    tmp357 = tl_math.log(tmp356)
    tmp358 = tmp346 - tmp357
    tmp359 = tmp186 * tmp358
    tmp360 = tmp348 - tmp357
    tmp361 = tmp201 * tmp360
    tmp362 = tmp359 + tmp361
    tmp363 = tmp351 - tmp357
    tmp364 = tmp205 * tmp363
    tmp365 = tmp362 + tmp364
    tmp366 = tmp354 - tmp357
    tmp367 = tmp209 * tmp366
    tmp368 = tmp365 + tmp367
    tmp369 = tl.broadcast_to(tmp368, [XBLOCK, RBLOCK])
    tmp371 = tl.sum(tmp369, 1)[:, None]
    tmp374 = tl_math.exp(tmp373)
    tmp376 = tl_math.exp(tmp375)
    tmp377 = tmp374 + tmp376
    tmp379 = tl_math.exp(tmp378)
    tmp380 = tmp377 + tmp379
    tmp382 = tl_math.exp(tmp381)
    tmp383 = tmp380 + tmp382
    tmp384 = tl_math.log(tmp383)
    tmp385 = tmp373 - tmp384
    tmp386 = tmp372 * tmp385
    tmp388 = tmp375 - tmp384
    tmp389 = tmp387 * tmp388
    tmp390 = tmp386 + tmp389
    tmp392 = tmp378 - tmp384
    tmp393 = tmp391 * tmp392
    tmp394 = tmp390 + tmp393
    tmp396 = tmp381 - tmp384
    tmp397 = tmp395 * tmp396
    tmp398 = tmp394 + tmp397
    tmp399 = tl.broadcast_to(tmp398, [XBLOCK, RBLOCK])
    tmp401 = tl.sum(tmp399, 1)[:, None]
    tmp403 = tl_math.exp(tmp402)
    tmp405 = tl_math.exp(tmp404)
    tmp406 = tmp403 + tmp405
    tmp408 = tl_math.exp(tmp407)
    tmp409 = tmp406 + tmp408
    tmp411 = tl_math.exp(tmp410)
    tmp412 = tmp409 + tmp411
    tmp413 = tl_math.log(tmp412)
    tmp414 = tmp402 - tmp413
    tmp415 = tmp372 * tmp414
    tmp416 = tmp404 - tmp413
    tmp417 = tmp387 * tmp416
    tmp418 = tmp415 + tmp417
    tmp419 = tmp407 - tmp413
    tmp420 = tmp391 * tmp419
    tmp421 = tmp418 + tmp420
    tmp422 = tmp410 - tmp413
    tmp423 = tmp395 * tmp422
    tmp424 = tmp421 + tmp423
    tmp425 = tl.broadcast_to(tmp424, [XBLOCK, RBLOCK])
    tmp427 = tl.sum(tmp425, 1)[:, None]
    tmp429 = tl_math.exp(tmp428)
    tmp431 = tl_math.exp(tmp430)
    tmp432 = tmp429 + tmp431
    tmp434 = tl_math.exp(tmp433)
    tmp435 = tmp432 + tmp434
    tmp437 = tl_math.exp(tmp436)
    tmp438 = tmp435 + tmp437
    tmp439 = tl_math.log(tmp438)
    tmp440 = tmp428 - tmp439
    tmp441 = tmp372 * tmp440
    tmp442 = tmp430 - tmp439
    tmp443 = tmp387 * tmp442
    tmp444 = tmp441 + tmp443
    tmp445 = tmp433 - tmp439
    tmp446 = tmp391 * tmp445
    tmp447 = tmp444 + tmp446
    tmp448 = tmp436 - tmp439
    tmp449 = tmp395 * tmp448
    tmp450 = tmp447 + tmp449
    tmp451 = tl.broadcast_to(tmp450, [XBLOCK, RBLOCK])
    tmp453 = tl.sum(tmp451, 1)[:, None]
    tmp455 = tl_math.exp(tmp454)
    tmp457 = tl_math.exp(tmp456)
    tmp458 = tmp455 + tmp457
    tmp460 = tl_math.exp(tmp459)
    tmp461 = tmp458 + tmp460
    tmp463 = tl_math.exp(tmp462)
    tmp464 = tmp461 + tmp463
    tmp465 = tl_math.log(tmp464)
    tmp466 = tmp454 - tmp465
    tmp467 = tmp372 * tmp466
    tmp468 = tmp456 - tmp465
    tmp469 = tmp387 * tmp468
    tmp470 = tmp467 + tmp469
    tmp471 = tmp459 - tmp465
    tmp472 = tmp391 * tmp471
    tmp473 = tmp470 + tmp472
    tmp474 = tmp462 - tmp465
    tmp475 = tmp395 * tmp474
    tmp476 = tmp473 + tmp475
    tmp477 = tl.broadcast_to(tmp476, [XBLOCK, RBLOCK])
    tmp479 = tl.sum(tmp477, 1)[:, None]
    tmp481 = tl_math.exp(tmp480)
    tmp483 = tl_math.exp(tmp482)
    tmp484 = tmp481 + tmp483
    tmp486 = tl_math.exp(tmp485)
    tmp487 = tmp484 + tmp486
    tmp489 = tl_math.exp(tmp488)
    tmp490 = tmp487 + tmp489
    tmp491 = tl_math.log(tmp490)
    tmp492 = tmp480 - tmp491
    tmp493 = tmp372 * tmp492
    tmp494 = tmp482 - tmp491
    tmp495 = tmp387 * tmp494
    tmp496 = tmp493 + tmp495
    tmp497 = tmp485 - tmp491
    tmp498 = tmp391 * tmp497
    tmp499 = tmp496 + tmp498
    tmp500 = tmp488 - tmp491
    tmp501 = tmp395 * tmp500
    tmp502 = tmp499 + tmp501
    tmp503 = tl.broadcast_to(tmp502, [XBLOCK, RBLOCK])
    tmp505 = tl.sum(tmp503, 1)[:, None]
    tmp507 = tl_math.exp(tmp506)
    tmp509 = tl_math.exp(tmp508)
    tmp510 = tmp507 + tmp509
    tmp512 = tl_math.exp(tmp511)
    tmp513 = tmp510 + tmp512
    tmp515 = tl_math.exp(tmp514)
    tmp516 = tmp513 + tmp515
    tmp517 = tl_math.log(tmp516)
    tmp518 = tmp506 - tmp517
    tmp519 = tmp372 * tmp518
    tmp520 = tmp508 - tmp517
    tmp521 = tmp387 * tmp520
    tmp522 = tmp519 + tmp521
    tmp523 = tmp511 - tmp517
    tmp524 = tmp391 * tmp523
    tmp525 = tmp522 + tmp524
    tmp526 = tmp514 - tmp517
    tmp527 = tmp395 * tmp526
    tmp528 = tmp525 + tmp527
    tmp529 = tl.broadcast_to(tmp528, [XBLOCK, RBLOCK])
    tmp531 = tl.sum(tmp529, 1)[:, None]
    tmp533 = tl_math.exp(tmp532)
    tmp535 = tl_math.exp(tmp534)
    tmp536 = tmp533 + tmp535
    tmp538 = tl_math.exp(tmp537)
    tmp539 = tmp536 + tmp538
    tmp541 = tl_math.exp(tmp540)
    tmp542 = tmp539 + tmp541
    tmp543 = tl_math.log(tmp542)
    tmp544 = tmp532 - tmp543
    tmp545 = tmp372 * tmp544
    tmp546 = tmp534 - tmp543
    tmp547 = tmp387 * tmp546
    tmp548 = tmp545 + tmp547
    tmp549 = tmp537 - tmp543
    tmp550 = tmp391 * tmp549
    tmp551 = tmp548 + tmp550
    tmp552 = tmp540 - tmp543
    tmp553 = tmp395 * tmp552
    tmp554 = tmp551 + tmp553
    tmp555 = tl.broadcast_to(tmp554, [XBLOCK, RBLOCK])
    tmp557 = tl.sum(tmp555, 1)[:, None]
    tmp560 = tl_math.exp(tmp559)
    tmp562 = tl_math.exp(tmp561)
    tmp563 = tmp560 + tmp562
    tmp565 = tl_math.exp(tmp564)
    tmp566 = tmp563 + tmp565
    tmp568 = tl_math.exp(tmp567)
    tmp569 = tmp566 + tmp568
    tmp570 = tl_math.log(tmp569)
    tmp571 = tmp559 - tmp570
    tmp572 = tmp558 * tmp571
    tmp574 = tmp561 - tmp570
    tmp575 = tmp573 * tmp574
    tmp576 = tmp572 + tmp575
    tmp578 = tmp564 - tmp570
    tmp579 = tmp577 * tmp578
    tmp580 = tmp576 + tmp579
    tmp582 = tmp567 - tmp570
    tmp583 = tmp581 * tmp582
    tmp584 = tmp580 + tmp583
    tmp585 = tl.broadcast_to(tmp584, [XBLOCK, RBLOCK])
    tmp587 = tl.sum(tmp585, 1)[:, None]
    tmp589 = tl_math.exp(tmp588)
    tmp591 = tl_math.exp(tmp590)
    tmp592 = tmp589 + tmp591
    tmp594 = tl_math.exp(tmp593)
    tmp595 = tmp592 + tmp594
    tmp597 = tl_math.exp(tmp596)
    tmp598 = tmp595 + tmp597
    tmp599 = tl_math.log(tmp598)
    tmp600 = tmp588 - tmp599
    tmp601 = tmp558 * tmp600
    tmp602 = tmp590 - tmp599
    tmp603 = tmp573 * tmp602
    tmp604 = tmp601 + tmp603
    tmp605 = tmp593 - tmp599
    tmp606 = tmp577 * tmp605
    tmp607 = tmp604 + tmp606
    tmp608 = tmp596 - tmp599
    tmp609 = tmp581 * tmp608
    tmp610 = tmp607 + tmp609
    tmp611 = tl.broadcast_to(tmp610, [XBLOCK, RBLOCK])
    tmp613 = tl.sum(tmp611, 1)[:, None]
    tmp615 = tl_math.exp(tmp614)
    tmp617 = tl_math.exp(tmp616)
    tmp618 = tmp615 + tmp617
    tmp620 = tl_math.exp(tmp619)
    tmp621 = tmp618 + tmp620
    tmp623 = tl_math.exp(tmp622)
    tmp624 = tmp621 + tmp623
    tmp625 = tl_math.log(tmp624)
    tmp626 = tmp614 - tmp625
    tmp627 = tmp558 * tmp626
    tmp628 = tmp616 - tmp625
    tmp629 = tmp573 * tmp628
    tmp630 = tmp627 + tmp629
    tmp631 = tmp619 - tmp625
    tmp632 = tmp577 * tmp631
    tmp633 = tmp630 + tmp632
    tmp634 = tmp622 - tmp625
    tmp635 = tmp581 * tmp634
    tmp636 = tmp633 + tmp635
    tmp637 = tl.broadcast_to(tmp636, [XBLOCK, RBLOCK])
    tmp639 = tl.sum(tmp637, 1)[:, None]
    tmp641 = tl_math.exp(tmp640)
    tmp643 = tl_math.exp(tmp642)
    tmp644 = tmp641 + tmp643
    tmp646 = tl_math.exp(tmp645)
    tmp647 = tmp644 + tmp646
    tmp649 = tl_math.exp(tmp648)
    tmp650 = tmp647 + tmp649
    tmp651 = tl_math.log(tmp650)
    tmp652 = tmp640 - tmp651
    tmp653 = tmp558 * tmp652
    tmp654 = tmp642 - tmp651
    tmp655 = tmp573 * tmp654
    tmp656 = tmp653 + tmp655
    tmp657 = tmp645 - tmp651
    tmp658 = tmp577 * tmp657
    tmp659 = tmp656 + tmp658
    tmp660 = tmp648 - tmp651
    tmp661 = tmp581 * tmp660
    tmp662 = tmp659 + tmp661
    tmp663 = tl.broadcast_to(tmp662, [XBLOCK, RBLOCK])
    tmp665 = tl.sum(tmp663, 1)[:, None]
    tmp667 = tl_math.exp(tmp666)
    tmp669 = tl_math.exp(tmp668)
    tmp670 = tmp667 + tmp669
    tmp672 = tl_math.exp(tmp671)
    tmp673 = tmp670 + tmp672
    tmp675 = tl_math.exp(tmp674)
    tmp676 = tmp673 + tmp675
    tmp677 = tl_math.log(tmp676)
    tmp678 = tmp666 - tmp677
    tmp679 = tmp558 * tmp678
    tmp680 = tmp668 - tmp677
    tmp681 = tmp573 * tmp680
    tmp682 = tmp679 + tmp681
    tmp683 = tmp671 - tmp677
    tmp684 = tmp577 * tmp683
    tmp685 = tmp682 + tmp684
    tmp686 = tmp674 - tmp677
    tmp687 = tmp581 * tmp686
    tmp688 = tmp685 + tmp687
    tmp689 = tl.broadcast_to(tmp688, [XBLOCK, RBLOCK])
    tmp691 = tl.sum(tmp689, 1)[:, None]
    tmp693 = tl_math.exp(tmp692)
    tmp695 = tl_math.exp(tmp694)
    tmp696 = tmp693 + tmp695
    tmp698 = tl_math.exp(tmp697)
    tmp699 = tmp696 + tmp698
    tmp701 = tl_math.exp(tmp700)
    tmp702 = tmp699 + tmp701
    tmp703 = tl_math.log(tmp702)
    tmp704 = tmp692 - tmp703
    tmp705 = tmp558 * tmp704
    tmp706 = tmp694 - tmp703
    tmp707 = tmp573 * tmp706
    tmp708 = tmp705 + tmp707
    tmp709 = tmp697 - tmp703
    tmp710 = tmp577 * tmp709
    tmp711 = tmp708 + tmp710
    tmp712 = tmp700 - tmp703
    tmp713 = tmp581 * tmp712
    tmp714 = tmp711 + tmp713
    tmp715 = tl.broadcast_to(tmp714, [XBLOCK, RBLOCK])
    tmp717 = tl.sum(tmp715, 1)[:, None]
    tmp719 = tl_math.exp(tmp718)
    tmp721 = tl_math.exp(tmp720)
    tmp722 = tmp719 + tmp721
    tmp724 = tl_math.exp(tmp723)
    tmp725 = tmp722 + tmp724
    tmp727 = tl_math.exp(tmp726)
    tmp728 = tmp725 + tmp727
    tmp729 = tl_math.log(tmp728)
    tmp730 = tmp718 - tmp729
    tmp731 = tmp558 * tmp730
    tmp732 = tmp720 - tmp729
    tmp733 = tmp573 * tmp732
    tmp734 = tmp731 + tmp733
    tmp735 = tmp723 - tmp729
    tmp736 = tmp577 * tmp735
    tmp737 = tmp734 + tmp736
    tmp738 = tmp726 - tmp729
    tmp739 = tmp581 * tmp738
    tmp740 = tmp737 + tmp739
    tmp741 = tl.broadcast_to(tmp740, [XBLOCK, RBLOCK])
    tmp743 = tl.sum(tmp741, 1)[:, None]
    tmp744 = 4.0
    tmp745 = tmp587 / tmp744
    tmp746 = -tmp745
    tmp747 = 0.0
    tmp748 = tmp746 + tmp747
    tmp749 = tmp613 / tmp744
    tmp750 = -tmp749
    tmp751 = tmp748 + tmp750
    tmp752 = tmp639 / tmp744
    tmp753 = -tmp752
    tmp754 = tmp751 + tmp753
    tmp755 = tmp665 / tmp744
    tmp756 = -tmp755
    tmp757 = tmp754 + tmp756
    tmp758 = tmp691 / tmp744
    tmp759 = -tmp758
    tmp760 = tmp757 + tmp759
    tmp761 = tmp717 / tmp744
    tmp762 = -tmp761
    tmp763 = tmp760 + tmp762
    tmp764 = tmp743 / tmp744
    tmp765 = -tmp764
    tmp766 = tmp763 + tmp765
    tmp767 = 0.14285714285714285
    tmp768 = tmp766 * tmp767
    tmp769 = tmp401 / tmp744
    tmp770 = -tmp769
    tmp771 = tmp770 + tmp747
    tmp772 = tmp427 / tmp744
    tmp773 = -tmp772
    tmp774 = tmp771 + tmp773
    tmp775 = tmp453 / tmp744
    tmp776 = -tmp775
    tmp777 = tmp774 + tmp776
    tmp778 = tmp479 / tmp744
    tmp779 = -tmp778
    tmp780 = tmp777 + tmp779
    tmp781 = tmp505 / tmp744
    tmp782 = -tmp781
    tmp783 = tmp780 + tmp782
    tmp784 = tmp531 / tmp744
    tmp785 = -tmp784
    tmp786 = tmp783 + tmp785
    tmp787 = tmp557 / tmp744
    tmp788 = -tmp787
    tmp789 = tmp786 + tmp788
    tmp790 = tmp789 * tmp767
    tmp791 = tmp215 / tmp744
    tmp792 = -tmp791
    tmp793 = tmp792 + tmp747
    tmp794 = tmp241 / tmp744
    tmp795 = -tmp794
    tmp796 = tmp793 + tmp795
    tmp797 = tmp267 / tmp744
    tmp798 = -tmp797
    tmp799 = tmp796 + tmp798
    tmp800 = tmp293 / tmp744
    tmp801 = -tmp800
    tmp802 = tmp799 + tmp801
    tmp803 = tmp319 / tmp744
    tmp804 = -tmp803
    tmp805 = tmp802 + tmp804
    tmp806 = tmp345 / tmp744
    tmp807 = -tmp806
    tmp808 = tmp805 + tmp807
    tmp809 = tmp371 / tmp744
    tmp810 = -tmp809
    tmp811 = tmp808 + tmp810
    tmp812 = tmp811 * tmp767
    tmp813 = tmp29 / tmp744
    tmp814 = -tmp813
    tmp815 = tmp814 + tmp747
    tmp816 = tmp55 / tmp744
    tmp817 = -tmp816
    tmp818 = tmp815 + tmp817
    tmp819 = tmp81 / tmp744
    tmp820 = -tmp819
    tmp821 = tmp818 + tmp820
    tmp822 = tmp107 / tmp744
    tmp823 = -tmp822
    tmp824 = tmp821 + tmp823
    tmp825 = tmp133 / tmp744
    tmp826 = -tmp825
    tmp827 = tmp824 + tmp826
    tmp828 = tmp159 / tmp744
    tmp829 = -tmp828
    tmp830 = tmp827 + tmp829
    tmp831 = tmp185 / tmp744
    tmp832 = -tmp831
    tmp833 = tmp830 + tmp832
    tmp834 = tmp833 * tmp767
    tmp835 = tmp768 + tmp747
    tmp836 = tmp835 + tmp790
    tmp837 = tmp836 + tmp812
    tmp838 = tmp837 + tmp834
    tmp839 = 0.25
    tmp840 = tmp838 * tmp839
    tl.debug_barrier()
    tl.store(in_out_ptr0 + tl.full([XBLOCK, 1], 0, tl.int32), tmp840, None)


def call(args):
    arg0_1, arg1_1 = args
    args.clear()
    assert_size_stride(arg0_1, (4, 4, 4), (16, 4, 1))
    assert_size_stride(arg1_1, (4, 4, 4), (16, 4, 1))
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        buf0 = empty_strided_cuda((), (), torch.float32)
        get_raw_stream(0)
        triton_per_fused_sum_0[grid(1)](arg0_1, buf0, 1, 16, XBLOCK=1,
            num_warps=2, num_stages=1)
        buf1 = empty_strided_cuda((4, 1), (1, 4), torch.float32)
        triton_poi_fused_sum_1[grid(4)](arg0_1, buf0, buf1, 4, XBLOCK=4,
            num_warps=1, num_stages=1)
        buf2 = empty_strided_cuda((1, 4), (4, 1), torch.float32)
        triton_poi_fused_sum_2[grid(4)](arg0_1, buf0, buf1, buf2, 4, XBLOCK
            =4, num_warps=1, num_stages=1)
        buf3 = empty_strided_cuda((4, 1), (1, 4), torch.float32)
        triton_poi_fused_sum_3[grid(4)](arg0_1, buf0, buf1, buf2, buf3, 4,
            XBLOCK=4, num_warps=1, num_stages=1)
        buf4 = empty_strided_cuda((4, 4), (1, 4), torch.float32)
        triton_poi_fused_div_4[grid(16)](arg0_1, buf0, buf1, buf2, buf3,
            buf4, 16, XBLOCK=16, num_warps=1, num_stages=1)
        buf5 = empty_strided_cuda((4, 4), (1, 4), torch.float32)
        triton_poi_fused_div_5[grid(16)](buf4, buf5, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf6 = buf4
        del buf4
        triton_poi_fused_div_6[grid(16)](buf5, buf6, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf7 = buf5
        del buf5
        triton_poi_fused_mul_7[grid(16)](buf6, buf7, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf23 = buf0
        del buf0
        buf56 = reinterpret_tensor(buf6, (4, 4), (4, 1), 0)
        del buf6
        buf79 = empty_strided_cuda((4, 4), (4, 1), torch.float32)
        triton_per_fused_sum_8[grid(1)](arg0_1, buf23, buf56, buf79, 1, 16,
            XBLOCK=1, num_warps=2, num_stages=1)
        buf24 = buf3
        del buf3
        triton_poi_fused_sum_9[grid(4)](arg0_1, buf23, buf24, 4, XBLOCK=4,
            num_warps=1, num_stages=1)
        buf25 = buf2
        del buf2
        triton_poi_fused_sum_10[grid(4)](arg0_1, buf23, buf24, buf25, 4,
            XBLOCK=4, num_warps=1, num_stages=1)
        buf26 = buf1
        del buf1
        triton_poi_fused_sum_11[grid(4)](arg0_1, buf23, buf24, buf25, buf26,
            4, XBLOCK=4, num_warps=1, num_stages=1)
        buf8 = empty_strided_cuda((4, 4), (4, 1), torch.float32)
        buf27 = empty_strided_cuda((4, 4), (1, 4), torch.float32)
        triton_poi_fused_div_12[grid(16)](arg0_1, buf23, buf24, buf25,
            buf26, buf8, buf27, 16, XBLOCK=16, num_warps=1, num_stages=1)
        buf46 = buf23
        del buf23
        buf81 = empty_strided_cuda((4, 4), (4, 1), torch.float32)
        triton_per_fused_sum_13[grid(1)](arg0_1, buf46, buf81, 1, 16,
            XBLOCK=1, num_warps=2, num_stages=1)
        buf47 = buf26
        del buf26
        triton_poi_fused_sum_14[grid(4)](arg0_1, buf46, buf47, 4, XBLOCK=4,
            num_warps=1, num_stages=1)
        buf48 = buf25
        del buf25
        triton_poi_fused_sum_15[grid(4)](arg0_1, buf46, buf47, buf48, 4,
            XBLOCK=4, num_warps=1, num_stages=1)
        buf49 = buf24
        del buf24
        triton_poi_fused_sum_16[grid(4)](arg0_1, buf46, buf47, buf48, buf49,
            4, XBLOCK=4, num_warps=1, num_stages=1)
        buf10 = empty_strided_cuda((4, 4), (4, 1), torch.float32)
        buf33 = empty_strided_cuda((4, 4), (4, 1), torch.float32)
        buf50 = empty_strided_cuda((4, 4), (1, 4), torch.float32)
        triton_poi_fused_div_17[grid(16)](arg0_1, buf46, buf47, buf48,
            buf49, buf10, buf33, buf50, 16, XBLOCK=16, num_warps=1,
            num_stages=1)
        buf69 = buf46
        del buf46
        triton_per_fused_sum_18[grid(1)](arg0_1, buf69, 1, 16, XBLOCK=1,
            num_warps=2, num_stages=1)
        buf70 = buf49
        del buf49
        triton_poi_fused_sum_19[grid(4)](arg0_1, buf69, buf70, 4, XBLOCK=4,
            num_warps=1, num_stages=1)
        buf71 = buf48
        del buf48
        triton_poi_fused_sum_20[grid(4)](arg0_1, buf69, buf70, buf71, 4,
            XBLOCK=4, num_warps=1, num_stages=1)
        buf72 = buf47
        del buf47
        triton_poi_fused_sum_21[grid(4)](arg0_1, buf69, buf70, buf71, buf72,
            4, XBLOCK=4, num_warps=1, num_stages=1)
        buf12 = empty_strided_cuda((4, 4), (4, 1), torch.float32)
        buf35 = empty_strided_cuda((4, 4), (4, 1), torch.float32)
        buf58 = empty_strided_cuda((4, 4), (4, 1), torch.float32)
        buf73 = empty_strided_cuda((4, 4), (1, 4), torch.float32)
        triton_poi_fused_div_22[grid(16)](arg0_1, buf69, buf70, buf71,
            buf72, buf12, buf35, buf58, buf73, 16, XBLOCK=16, num_warps=1,
            num_stages=1)
        del buf70
        del buf71
        del buf72
        buf14 = empty_strided_cuda((4, 4), (4, 1), torch.float32)
        buf37 = empty_strided_cuda((4, 4), (4, 1), torch.float32)
        buf60 = empty_strided_cuda((4, 4), (4, 1), torch.float32)
        triton_poi_fused_23[grid(16)](arg1_1, buf14, buf37, buf60, 16,
            XBLOCK=16, num_warps=1, num_stages=1)
        buf16 = empty_strided_cuda((4, 4), (4, 1), torch.float32)
        buf39 = empty_strided_cuda((4, 4), (4, 1), torch.float32)
        buf62 = empty_strided_cuda((4, 4), (4, 1), torch.float32)
        triton_poi_fused_24[grid(16)](arg1_1, buf16, buf39, buf62, 16,
            XBLOCK=16, num_warps=1, num_stages=1)
        buf18 = empty_strided_cuda((4, 4), (4, 1), torch.float32)
        buf41 = empty_strided_cuda((4, 4), (4, 1), torch.float32)
        buf64 = empty_strided_cuda((4, 4), (4, 1), torch.float32)
        triton_poi_fused_25[grid(16)](arg1_1, buf18, buf41, buf64, 16,
            XBLOCK=16, num_warps=1, num_stages=1)
        buf20 = empty_strided_cuda((4, 4), (4, 1), torch.float32)
        buf43 = empty_strided_cuda((4, 4), (4, 1), torch.float32)
        buf66 = empty_strided_cuda((4, 4), (4, 1), torch.float32)
        triton_poi_fused_26[grid(16)](arg1_1, buf20, buf43, buf66, 16,
            XBLOCK=16, num_warps=1, num_stages=1)
        buf28 = empty_strided_cuda((4, 4), (1, 4), torch.float32)
        triton_poi_fused_div_5[grid(16)](buf27, buf28, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf29 = buf27
        del buf27
        triton_poi_fused_div_6[grid(16)](buf28, buf29, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf30 = buf28
        del buf28
        triton_poi_fused_mul_7[grid(16)](buf29, buf30, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf31 = reinterpret_tensor(buf29, (4, 4), (4, 1), 0)
        del buf29
        buf54 = empty_strided_cuda((4, 4), (4, 1), torch.float32)
        buf77 = empty_strided_cuda((4, 4), (4, 1), torch.float32)
        triton_poi_fused_23[grid(16)](arg0_1, buf31, buf54, buf77, 16,
            XBLOCK=16, num_warps=1, num_stages=1)
        del arg0_1
        buf51 = empty_strided_cuda((4, 4), (1, 4), torch.float32)
        triton_poi_fused_div_5[grid(16)](buf50, buf51, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf52 = buf50
        del buf50
        triton_poi_fused_div_6[grid(16)](buf51, buf52, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf53 = buf51
        del buf51
        triton_poi_fused_mul_7[grid(16)](buf52, buf53, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf74 = buf52
        del buf52
        triton_poi_fused_div_5[grid(16)](buf73, buf74, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf75 = buf73
        del buf73
        triton_poi_fused_div_6[grid(16)](buf74, buf75, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf76 = buf74
        del buf74
        triton_poi_fused_mul_7[grid(16)](buf75, buf76, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf83 = reinterpret_tensor(buf75, (4, 4), (4, 1), 0)
        del buf75
        triton_poi_fused_27[grid(16)](arg1_1, buf83, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf85 = empty_strided_cuda((4, 4), (4, 1), torch.float32)
        triton_poi_fused_28[grid(16)](arg1_1, buf85, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf87 = empty_strided_cuda((4, 4), (4, 1), torch.float32)
        triton_poi_fused_29[grid(16)](arg1_1, buf87, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf89 = empty_strided_cuda((4, 4), (4, 1), torch.float32)
        triton_poi_fused_30[grid(16)](arg1_1, buf89, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        del arg1_1
        buf11 = buf69
        del buf69
        buf22 = buf11
        del buf11
        buf92 = buf22
        del buf22
        triton_per_fused__log_softmax_add_div_mean_mul_neg_sum_31[grid(1)](
            buf92, buf76, buf77, buf79, buf81, buf83, buf85, buf87, buf89,
            buf53, buf54, buf56, buf58, buf60, buf62, buf64, buf66, buf30,
            buf31, buf33, buf35, buf37, buf39, buf41, buf43, buf7, buf8,
            buf10, buf12, buf14, buf16, buf18, buf20, 1, 4, XBLOCK=1,
            num_warps=2, num_stages=1)
        del buf10
        del buf12
        del buf14
        del buf16
        del buf18
        del buf20
        del buf30
        del buf31
        del buf33
        del buf35
        del buf37
        del buf39
        del buf41
        del buf43
        del buf53
        del buf54
        del buf56
        del buf58
        del buf60
        del buf62
        del buf64
        del buf66
        del buf7
        del buf76
        del buf77
        del buf79
        del buf8
        del buf81
        del buf83
        del buf85
        del buf87
        del buf89
    return buf92,


@torch.no_grad()
def sinkhorn(out: 'torch.Tensor', iterations: 'int'=3, epsilon: 'float'=0.05):
    """Distributed sinkhorn algorithm.

    As outlined in [0] and implemented in [1].
    
    [0]: SwaV, 2020, https://arxiv.org/abs/2006.09882
    [1]: https://github.com/facebookresearch/swav/ 

    Args:
        out:
            Similarity of the features and the SwaV prototypes.
        iterations:
            Number of sinkhorn iterations.
        epsilon:
            Temperature parameter.

    Returns:
        Soft codes Q assigning each feature to a prototype.
    
    """
    Q = torch.exp(out / epsilon).t()
    sum_Q = torch.sum(Q)
    Q /= sum_Q
    B = Q.shape[1]
    K = Q.shape[0]
    for i in range(iterations):
        sum_of_rows = torch.sum(Q, dim=1, keepdim=True)
        Q /= sum_of_rows
        Q /= K
        Q /= torch.sum(Q, dim=0, keepdim=True)
        Q /= B
    Q *= B
    return Q.t()


class SwaVLossNew(nn.Module):
    """Implementation of the SwaV loss.

    Attributes:
        temperature:
            Temperature parameter used for cross entropy calculations.
        sinkhorn_iterations:
            Number of iterations of the sinkhorn algorithm.
        sinkhorn_epsilon:
            Temperature parameter used in the sinkhorn algorithm.
    
    """

    def __init__(self, temperature: 'float'=0.1, sinkhorn_iterations: 'int'
        =3, sinkhorn_epsilon: 'float'=0.05):
        super(SwaVLossNew, self).__init__()
        self.temperature = temperature
        self.sinkhorn_iterations = sinkhorn_iterations
        self.sinkhorn_epsilon = sinkhorn_epsilon

    def subloss(self, z: 'torch.Tensor', q: 'torch.Tensor'):
        """Calculates the cross entropy for the SwaV prediction problem.

        Args:
            z:
                Similarity of the features and the SwaV prototypes.
            q:
                Codes obtained from Sinkhorn iterations.

        Returns:
            Cross entropy between predictions z and codes q.

        """
        return -torch.mean(torch.sum(q * F.log_softmax(z / self.temperature,
            dim=1), dim=1))

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


Measured latency_ms: 0.592601