import torch
from torch._inductor.select_algorithm import extern_kernels
import triton
import triton.language as tl
from torch._inductor.runtime.triton_heuristics import grid
from torch._C import _cuda_getCurrentRawStream as get_raw_stream
from torch._inductor.runtime import triton_helpers
from torch._inductor.runtime.triton_helpers import libdevice, math as tl_math
import torch.nn as nn
import torch.nn.functional as F
from torch.functional import F
from torch.nn import functional as F
import torch.nn.parallel
import torch.optim
import torch.utils.data
import torch.utils.data.distributed
import torch.distributions
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_stack_0(in_ptr0, out_ptr0, out_ptr1, out_ptr2,
    out_ptr3, out_ptr4, out_ptr5, out_ptr6, xnumel, XBLOCK: tl.constexpr):
    xnumel = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex
    tmp0 = x0
    tl.full([1], 0, tl.int64)
    tmp3 = tl.full([1], 4, tl.int64)
    tmp4 = tmp0 < tmp3
    tmp5 = tl.load(in_ptr0 + (4 + x0), tmp4 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp6 = tl.load(in_ptr0 + (8 + x0), tmp4 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp7 = tmp5 + tmp6
    tmp8 = tl.load(in_ptr0 + (12 + x0), tmp4 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp9 = tmp7 + tmp8
    tmp10 = 3.0
    tmp11 = tmp9 / tmp10
    tmp12 = tl.full(tmp11.shape, 0.0, tmp11.dtype)
    tmp13 = tl.where(tmp4, tmp11, tmp12)
    tmp14 = tmp0 >= tmp3
    tmp15 = tl.full([1], 8, tl.int64)
    tmp16 = tmp0 < tmp15
    tmp17 = tmp14 & tmp16
    tmp18 = tl.load(in_ptr0 + (16 + (-4 + x0)), tmp17 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp19 = tl.load(in_ptr0 + (20 + (-4 + x0)), tmp17 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp20 = tmp18 + tmp19
    tmp21 = tl.load(in_ptr0 + (24 + (-4 + x0)), tmp17 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp22 = tmp20 + tmp21
    tmp23 = tl.load(in_ptr0 + (28 + (-4 + x0)), tmp17 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp24 = tmp22 + tmp23
    tmp25 = 4.0
    tmp26 = tmp24 / tmp25
    tmp27 = tl.full(tmp26.shape, 0.0, tmp26.dtype)
    tmp28 = tl.where(tmp17, tmp26, tmp27)
    tmp29 = tmp0 >= tmp15
    tmp30 = tl.full([1], 12, tl.int64)
    tmp31 = tmp0 < tmp30
    tmp32 = tmp29 & tmp31
    tmp33 = tl.load(in_ptr0 + (32 + (-8 + x0)), tmp32 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp34 = tl.load(in_ptr0 + (36 + (-8 + x0)), tmp32 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp35 = tmp33 + tmp34
    tmp36 = tl.load(in_ptr0 + (40 + (-8 + x0)), tmp32 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp37 = tmp35 + tmp36
    tmp38 = tl.load(in_ptr0 + (44 + (-8 + x0)), tmp32 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp39 = tmp37 + tmp38
    tmp40 = tmp39 / tmp25
    tmp41 = tl.full(tmp40.shape, 0.0, tmp40.dtype)
    tmp42 = tl.where(tmp32, tmp40, tmp41)
    tmp43 = tmp0 >= tmp30
    tl.full([1], 16, tl.int64)
    tmp46 = tl.load(in_ptr0 + (48 + (-12 + x0)), tmp43 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp47 = tl.load(in_ptr0 + (52 + (-12 + x0)), tmp43 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp48 = tmp46 + tmp47
    tmp49 = tl.load(in_ptr0 + (56 + (-12 + x0)), tmp43 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp50 = tmp48 + tmp49
    tmp51 = tl.load(in_ptr0 + (60 + (-12 + x0)), tmp43 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp52 = tmp50 + tmp51
    tmp53 = tmp52 / tmp25
    tmp54 = tl.full(tmp53.shape, 0.0, tmp53.dtype)
    tmp55 = tl.where(tmp43, tmp53, tmp54)
    tmp56 = tl.where(tmp32, tmp42, tmp55)
    tmp57 = tl.where(tmp17, tmp28, tmp56)
    tmp58 = tl.where(tmp4, tmp13, tmp57)
    tmp59 = tl.load(in_ptr0 + x0, tmp4 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp60 = tmp59 + tmp5
    tmp61 = tmp60 + tmp6
    tmp62 = tmp61 / tmp10
    tmp63 = tl.full(tmp62.shape, 0.0, tmp62.dtype)
    tmp64 = tl.where(tmp4, tmp62, tmp63)
    tmp65 = tl.where(tmp4, tmp64, tmp57)
    tmp66 = tmp61 + tmp8
    tmp67 = tmp66 / tmp25
    tmp68 = tl.full(tmp67.shape, 0.0, tmp67.dtype)
    tmp69 = tl.where(tmp4, tmp67, tmp68)
    tmp70 = tmp19 + tmp21
    tmp71 = tmp70 + tmp23
    tmp72 = tmp71 / tmp10
    tmp73 = tl.full(tmp72.shape, 0.0, tmp72.dtype)
    tmp74 = tl.where(tmp17, tmp72, tmp73)
    tmp75 = tl.where(tmp17, tmp74, tmp56)
    tmp76 = tl.where(tmp4, tmp69, tmp75)
    tmp77 = tmp22 / tmp10
    tmp78 = tl.full(tmp77.shape, 0.0, tmp77.dtype)
    tmp79 = tl.where(tmp17, tmp77, tmp78)
    tmp80 = tl.where(tmp17, tmp79, tmp56)
    tmp81 = tl.where(tmp4, tmp69, tmp80)
    tmp82 = tmp34 + tmp36
    tmp83 = tmp82 + tmp38
    tmp84 = tmp83 / tmp10
    tmp85 = tl.full(tmp84.shape, 0.0, tmp84.dtype)
    tmp86 = tl.where(tmp32, tmp84, tmp85)
    tmp87 = tl.where(tmp32, tmp86, tmp55)
    tmp88 = tl.where(tmp17, tmp28, tmp87)
    tmp89 = tl.where(tmp4, tmp69, tmp88)
    tmp90 = tmp37 / tmp10
    tmp91 = tl.full(tmp90.shape, 0.0, tmp90.dtype)
    tmp92 = tl.where(tmp32, tmp90, tmp91)
    tmp93 = tl.where(tmp32, tmp92, tmp55)
    tmp94 = tl.where(tmp17, tmp28, tmp93)
    tmp95 = tl.where(tmp4, tmp69, tmp94)
    tmp96 = tmp47 + tmp49
    tmp97 = tmp96 + tmp51
    tmp98 = tmp97 / tmp10
    tmp99 = tl.full(tmp98.shape, 0.0, tmp98.dtype)
    tmp100 = tl.where(tmp43, tmp98, tmp99)
    tmp101 = tl.where(tmp32, tmp42, tmp100)
    tmp102 = tl.where(tmp17, tmp28, tmp101)
    tmp103 = tl.where(tmp4, tmp69, tmp102)
    tl.store(out_ptr0 + x0, tmp58, xmask)
    tl.store(out_ptr1 + x0, tmp65, xmask)
    tl.store(out_ptr2 + x0, tmp76, xmask)
    tl.store(out_ptr3 + x0, tmp81, xmask)
    tl.store(out_ptr4 + x0, tmp89, xmask)
    tl.store(out_ptr5 + x0, tmp95, xmask)
    tl.store(out_ptr6 + x0, tmp103, xmask)


@triton.jit
def triton_per_fused_cat_linalg_vector_norm_mean_1(in_ptr0, out_ptr0,
    out_ptr1, 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 + r0, None)
    tmp1 = tmp0 * tmp0
    tmp2 = tl.broadcast_to(tmp1, [XBLOCK, RBLOCK])
    tmp4 = tl.sum(tmp2, 1)[:, None]
    tmp5 = tl.full([1, 1], 0, tl.int64)
    tmp7 = tl.full([1, 1], 1, tl.int64)
    tmp8 = tmp5 < tmp7
    tmp9 = tl.load(in_ptr0 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp8,
        other=0.0)
    tmp10 = tmp5 >= tmp7
    tl.full([1, 1], 3, tl.int64)
    tmp13 = tl.load(in_ptr0 + tl.broadcast_to(8 + r0 + 4 * -1, [XBLOCK,
        RBLOCK]), tmp10, other=0.0)
    tmp14 = tl.where(tmp8, tmp9, tmp13)
    tmp16 = tmp7 < tmp7
    tmp17 = tl.load(in_ptr0 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp16,
        other=0.0)
    tmp18 = tmp7 >= tmp7
    tmp20 = tl.load(in_ptr0 + tl.broadcast_to(8 + r0 + 4 * 0, [XBLOCK,
        RBLOCK]), tmp18, other=0.0)
    tmp21 = tl.where(tmp16, tmp17, tmp20)
    tmp22 = tmp14 + tmp21
    tmp23 = tl.full([1, 1], 2, tl.int64)
    tmp25 = tmp23 < tmp7
    tmp26 = tl.load(in_ptr0 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp25,
        other=0.0)
    tmp27 = tmp23 >= tmp7
    tmp29 = tl.load(in_ptr0 + tl.broadcast_to(8 + r0 + 4 * 1, [XBLOCK,
        RBLOCK]), tmp27, other=0.0)
    tmp30 = tl.where(tmp25, tmp26, tmp29)
    tmp31 = tmp22 + tmp30
    tmp32 = 3.0
    tmp33 = tmp31 / tmp32
    tl.store(out_ptr1 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp33, None)
    tl.store(out_ptr0 + tl.full([XBLOCK, 1], 0, tl.int32), tmp4, None)


@triton.jit
def triton_per_fused_cat_linalg_vector_norm_mean_stack_2(in_ptr0, out_ptr0,
    out_ptr1, out_ptr2, out_ptr3, 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 + (16 + r0), None)
    tmp5 = tl.load(in_ptr0 + (20 + r0), None)
    tmp7 = tl.load(in_ptr0 + (24 + r0), None)
    tmp9 = tl.load(in_ptr0 + (28 + r0), None)
    tmp1 = tmp0 * tmp0
    tmp2 = tl.broadcast_to(tmp1, [XBLOCK, RBLOCK])
    tmp4 = tl.sum(tmp2, 1)[:, None]
    tmp6 = tmp0 + tmp5
    tmp8 = tmp6 + tmp7
    tmp10 = tmp8 + tmp9
    tmp11 = 4.0
    tmp12 = tmp10 / tmp11
    tmp13 = tl.full([1, 1], 0, tl.int64)
    tmp15 = tl.full([1, 1], 1, tl.int64)
    tmp16 = tmp13 < tmp15
    tmp17 = tl.load(in_ptr0 + tl.broadcast_to(16 + r0, [XBLOCK, RBLOCK]),
        tmp16, other=0.0)
    tmp18 = tmp13 >= tmp15
    tl.full([1, 1], 3, tl.int64)
    tmp21 = tl.load(in_ptr0 + tl.broadcast_to(24 + r0 + 4 * -1, [XBLOCK,
        RBLOCK]), tmp18, other=0.0)
    tmp22 = tl.where(tmp16, tmp17, tmp21)
    tmp24 = tmp15 < tmp15
    tmp25 = tl.load(in_ptr0 + tl.broadcast_to(16 + r0, [XBLOCK, RBLOCK]),
        tmp24, other=0.0)
    tmp26 = tmp15 >= tmp15
    tmp28 = tl.load(in_ptr0 + tl.broadcast_to(24 + r0 + 4 * 0, [XBLOCK,
        RBLOCK]), tmp26, other=0.0)
    tmp29 = tl.where(tmp24, tmp25, tmp28)
    tmp30 = tmp22 + tmp29
    tmp31 = tl.full([1, 1], 2, tl.int64)
    tmp33 = tmp31 < tmp15
    tmp34 = tl.load(in_ptr0 + tl.broadcast_to(16 + r0, [XBLOCK, RBLOCK]),
        tmp33, other=0.0)
    tmp35 = tmp31 >= tmp15
    tmp37 = tl.load(in_ptr0 + tl.broadcast_to(24 + r0 + 4 * 1, [XBLOCK,
        RBLOCK]), tmp35, other=0.0)
    tmp38 = tl.where(tmp33, tmp34, tmp37)
    tmp39 = tmp30 + tmp38
    tmp40 = 3.0
    tmp41 = tmp39 / tmp40
    tl.store(out_ptr1 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp12, None)
    tl.store(out_ptr2 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp12, None)
    tl.store(out_ptr3 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp41, None)
    tl.store(out_ptr0 + tl.full([XBLOCK, 1], 0, tl.int32), tmp4, None)


@triton.jit
def triton_per_fused_cat_linalg_vector_norm_mean_stack_3(in_ptr0, out_ptr0,
    out_ptr1, out_ptr2, out_ptr3, out_ptr4, out_ptr5, out_ptr6, 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 + (32 + r0), None)
    tmp5 = tl.load(in_ptr0 + (36 + r0), None)
    tmp7 = tl.load(in_ptr0 + (40 + r0), None)
    tmp9 = tl.load(in_ptr0 + (44 + r0), None)
    tmp1 = tmp0 * tmp0
    tmp2 = tl.broadcast_to(tmp1, [XBLOCK, RBLOCK])
    tmp4 = tl.sum(tmp2, 1)[:, None]
    tmp6 = tmp0 + tmp5
    tmp8 = tmp6 + tmp7
    tmp10 = tmp8 + tmp9
    tmp11 = 4.0
    tmp12 = tmp10 / tmp11
    tmp13 = tl.full([1, 1], 0, tl.int64)
    tmp15 = tl.full([1, 1], 1, tl.int64)
    tmp16 = tmp13 < tmp15
    tmp17 = tl.load(in_ptr0 + tl.broadcast_to(32 + r0, [XBLOCK, RBLOCK]),
        tmp16, other=0.0)
    tmp18 = tmp13 >= tmp15
    tl.full([1, 1], 3, tl.int64)
    tmp21 = tl.load(in_ptr0 + tl.broadcast_to(40 + r0 + 4 * -1, [XBLOCK,
        RBLOCK]), tmp18, other=0.0)
    tmp22 = tl.where(tmp16, tmp17, tmp21)
    tmp24 = tmp15 < tmp15
    tmp25 = tl.load(in_ptr0 + tl.broadcast_to(32 + r0, [XBLOCK, RBLOCK]),
        tmp24, other=0.0)
    tmp26 = tmp15 >= tmp15
    tmp28 = tl.load(in_ptr0 + tl.broadcast_to(40 + r0 + 4 * 0, [XBLOCK,
        RBLOCK]), tmp26, other=0.0)
    tmp29 = tl.where(tmp24, tmp25, tmp28)
    tmp30 = tmp22 + tmp29
    tmp31 = tl.full([1, 1], 2, tl.int64)
    tmp33 = tmp31 < tmp15
    tmp34 = tl.load(in_ptr0 + tl.broadcast_to(32 + r0, [XBLOCK, RBLOCK]),
        tmp33, other=0.0)
    tmp35 = tmp31 >= tmp15
    tmp37 = tl.load(in_ptr0 + tl.broadcast_to(40 + r0 + 4 * 1, [XBLOCK,
        RBLOCK]), tmp35, other=0.0)
    tmp38 = tl.where(tmp33, tmp34, tmp37)
    tmp39 = tmp30 + tmp38
    tmp40 = 3.0
    tmp41 = tmp39 / tmp40
    tmp42 = tmp5 * tmp5
    tmp43 = tl.broadcast_to(tmp42, [XBLOCK, RBLOCK])
    tmp45 = tl.sum(tmp43, 1)[:, None]
    tl.store(out_ptr1 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp12, None)
    tl.store(out_ptr2 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp12, None)
    tl.store(out_ptr3 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp12, None)
    tl.store(out_ptr4 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp12, None)
    tl.store(out_ptr5 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp41, None)
    tl.store(out_ptr0 + tl.full([XBLOCK, 1], 0, tl.int32), tmp4, None)
    tl.store(out_ptr6 + tl.full([XBLOCK, 1], 0, tl.int32), tmp45, None)


@triton.jit
def triton_per_fused_cat_linalg_vector_norm_mean_stack_4(in_ptr0, out_ptr0,
    out_ptr1, out_ptr2, out_ptr3, out_ptr4, out_ptr5, out_ptr6, out_ptr7,
    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 + (48 + r0), None)
    tmp5 = tl.load(in_ptr0 + (52 + r0), None)
    tmp7 = tl.load(in_ptr0 + (56 + r0), None)
    tmp9 = tl.load(in_ptr0 + (60 + r0), None)
    tmp1 = tmp0 * tmp0
    tmp2 = tl.broadcast_to(tmp1, [XBLOCK, RBLOCK])
    tmp4 = tl.sum(tmp2, 1)[:, None]
    tmp6 = tmp0 + tmp5
    tmp8 = tmp6 + tmp7
    tmp10 = tmp8 + tmp9
    tmp11 = 4.0
    tmp12 = tmp10 / tmp11
    tmp13 = tl.full([1, 1], 0, tl.int64)
    tmp15 = tl.full([1, 1], 1, tl.int64)
    tmp16 = tmp13 < tmp15
    tmp17 = tl.load(in_ptr0 + tl.broadcast_to(48 + r0, [XBLOCK, RBLOCK]),
        tmp16, other=0.0)
    tmp18 = tmp13 >= tmp15
    tl.full([1, 1], 3, tl.int64)
    tmp21 = tl.load(in_ptr0 + tl.broadcast_to(56 + r0 + 4 * -1, [XBLOCK,
        RBLOCK]), tmp18, other=0.0)
    tmp22 = tl.where(tmp16, tmp17, tmp21)
    tmp24 = tmp15 < tmp15
    tmp25 = tl.load(in_ptr0 + tl.broadcast_to(48 + r0, [XBLOCK, RBLOCK]),
        tmp24, other=0.0)
    tmp26 = tmp15 >= tmp15
    tmp28 = tl.load(in_ptr0 + tl.broadcast_to(56 + r0 + 4 * 0, [XBLOCK,
        RBLOCK]), tmp26, other=0.0)
    tmp29 = tl.where(tmp24, tmp25, tmp28)
    tmp30 = tmp22 + tmp29
    tmp31 = tl.full([1, 1], 2, tl.int64)
    tmp33 = tmp31 < tmp15
    tmp34 = tl.load(in_ptr0 + tl.broadcast_to(48 + r0, [XBLOCK, RBLOCK]),
        tmp33, other=0.0)
    tmp35 = tmp31 >= tmp15
    tmp37 = tl.load(in_ptr0 + tl.broadcast_to(56 + r0 + 4 * 1, [XBLOCK,
        RBLOCK]), tmp35, other=0.0)
    tmp38 = tl.where(tmp33, tmp34, tmp37)
    tmp39 = tmp30 + tmp38
    tmp40 = 3.0
    tmp41 = tmp39 / tmp40
    tl.store(out_ptr1 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp12, None)
    tl.store(out_ptr2 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp12, None)
    tl.store(out_ptr3 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp12, None)
    tl.store(out_ptr4 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp12, None)
    tl.store(out_ptr5 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp12, None)
    tl.store(out_ptr6 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp12, None)
    tl.store(out_ptr7 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp41, None)
    tl.store(out_ptr0 + tl.full([XBLOCK, 1], 0, tl.int32), tmp4, None)


@triton.jit
def triton_per_fused_linalg_vector_norm_5(in_ptr0, out_ptr0, 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)
    tmp1 = tmp0 * tmp0
    tmp2 = tl.broadcast_to(tmp1, [XBLOCK, RBLOCK])
    tmp4 = tl.sum(tmp2, 1)[:, None]
    tl.store(out_ptr0 + tl.full([XBLOCK, 1], 0, tl.int32), tmp4, None)


@triton.jit
def triton_per_fused_cat_linalg_vector_norm_mean_stack_6(in_ptr0, out_ptr0,
    out_ptr1, out_ptr2, out_ptr3, out_ptr4, out_ptr5, out_ptr6, out_ptr7,
    out_ptr8, 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 + (8 + r0), None)
    tmp5 = tl.load(in_ptr0 + (12 + r0), None)
    tmp39 = tl.load(in_ptr0 + r0, None)
    tmp40 = tl.load(in_ptr0 + (4 + r0), None)
    tmp1 = tmp0 * tmp0
    tmp2 = tl.broadcast_to(tmp1, [XBLOCK, RBLOCK])
    tmp4 = tl.sum(tmp2, 1)[:, None]
    tmp6 = tmp5 * tmp5
    tmp7 = tl.broadcast_to(tmp6, [XBLOCK, RBLOCK])
    tmp9 = tl.sum(tmp7, 1)[:, None]
    tmp10 = tl.full([1, 1], 0, tl.int64)
    tmp12 = tl.full([1, 1], 2, tl.int64)
    tmp13 = tmp10 < tmp12
    tmp14 = tl.load(in_ptr0 + tl.broadcast_to(r0 + 4 * 0, [XBLOCK, RBLOCK]),
        tmp13, other=0.0)
    tmp15 = tmp10 >= tmp12
    tl.full([1, 1], 3, tl.int64)
    tmp18 = tl.load(in_ptr0 + tl.broadcast_to(12 + r0, [XBLOCK, RBLOCK]),
        tmp15, other=0.0)
    tmp19 = tl.where(tmp13, tmp14, tmp18)
    tmp20 = tl.full([1, 1], 1, tl.int64)
    tmp22 = tmp20 < tmp12
    tmp23 = tl.load(in_ptr0 + tl.broadcast_to(r0 + 4 * 1, [XBLOCK, RBLOCK]),
        tmp22, other=0.0)
    tmp24 = tmp20 >= tmp12
    tmp26 = tl.load(in_ptr0 + tl.broadcast_to(12 + r0, [XBLOCK, RBLOCK]),
        tmp24, other=0.0)
    tmp27 = tl.where(tmp22, tmp23, tmp26)
    tmp28 = tmp19 + tmp27
    tmp30 = tmp12 < tmp12
    tmp31 = tl.load(in_ptr0 + tl.broadcast_to(r0 + 4 * 2, [XBLOCK, RBLOCK]),
        tmp30, other=0.0)
    tmp32 = tmp12 >= tmp12
    tmp34 = tl.load(in_ptr0 + tl.broadcast_to(12 + r0, [XBLOCK, RBLOCK]),
        tmp32, other=0.0)
    tmp35 = tl.where(tmp30, tmp31, tmp34)
    tmp36 = tmp28 + tmp35
    tmp37 = 3.0
    tmp38 = tmp36 / tmp37
    tmp41 = tmp39 + tmp40
    tmp42 = tmp41 + tmp0
    tmp43 = tmp42 + tmp5
    tmp44 = 4.0
    tmp45 = tmp43 / tmp44
    tl.store(out_ptr2 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp38, None)
    tl.store(out_ptr3 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp45, None)
    tl.store(out_ptr4 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp45, None)
    tl.store(out_ptr5 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp45, None)
    tl.store(out_ptr6 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp45, None)
    tl.store(out_ptr7 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp45, None)
    tl.store(out_ptr8 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp45, None)
    tl.store(out_ptr0 + tl.full([XBLOCK, 1], 0, tl.int32), tmp4, None)
    tl.store(out_ptr1 + tl.full([XBLOCK, 1], 0, tl.int32), tmp9, None)


@triton.jit
def triton_poi_fused_cat_7(in_ptr0, in_ptr1, in_ptr2, in_ptr3, in_ptr4,
    in_ptr5, in_ptr6, in_ptr7, in_ptr8, in_ptr9, in_ptr10, in_ptr11,
    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
    tmp6 = tl.load(in_ptr1 + 0)
    tmp7 = tl.broadcast_to(tmp6, [XBLOCK])
    tmp32 = tl.load(in_ptr4 + 0)
    tmp33 = tl.broadcast_to(tmp32, [XBLOCK])
    tmp57 = tl.load(in_ptr7 + 0)
    tmp58 = tl.broadcast_to(tmp57, [XBLOCK])
    tmp81 = tl.load(in_ptr10 + 0)
    tmp82 = tl.broadcast_to(tmp81, [XBLOCK])
    tmp0 = x1
    tl.full([1], 0, tl.int64)
    tmp3 = tl.full([1], 1, tl.int64)
    tmp4 = tmp0 < tmp3
    tmp5 = tl.load(in_ptr0 + x0, tmp4 & xmask, eviction_policy='evict_last',
        other=0.0)
    tmp8 = libdevice.sqrt(tmp7)
    tmp9 = tl.load(in_ptr2 + 4 * x0, tmp4 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp10 = tmp9 * tmp9
    tmp11 = tl.load(in_ptr2 + (1 + 4 * x0), tmp4 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp12 = tmp11 * tmp11
    tmp13 = tmp10 + tmp12
    tmp14 = tl.load(in_ptr2 + (2 + 4 * x0), tmp4 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp15 = tmp14 * tmp14
    tmp16 = tmp13 + tmp15
    tmp17 = tl.load(in_ptr2 + (3 + 4 * x0), tmp4 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp18 = tmp17 * tmp17
    tmp19 = tmp16 + tmp18
    tmp20 = libdevice.sqrt(tmp19)
    tmp21 = tmp8 * tmp20
    tmp22 = tmp5 / tmp21
    tmp23 = 1e-06
    tmp24 = triton_helpers.maximum(tmp22, tmp23)
    tmp25 = tl.full(tmp24.shape, 0.0, tmp24.dtype)
    tmp26 = tl.where(tmp4, tmp24, tmp25)
    tmp27 = tmp0 >= tmp3
    tmp28 = tl.full([1], 2, tl.int64)
    tmp29 = tmp0 < tmp28
    tmp30 = tmp27 & tmp29
    tmp31 = tl.load(in_ptr3 + x0, tmp30 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp34 = libdevice.sqrt(tmp33)
    tmp35 = tl.load(in_ptr5 + 4 * x0, tmp30 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp36 = tmp35 * tmp35
    tmp37 = tl.load(in_ptr5 + (1 + 4 * x0), tmp30 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp38 = tmp37 * tmp37
    tmp39 = tmp36 + tmp38
    tmp40 = tl.load(in_ptr5 + (2 + 4 * x0), tmp30 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp41 = tmp40 * tmp40
    tmp42 = tmp39 + tmp41
    tmp43 = tl.load(in_ptr5 + (3 + 4 * x0), tmp30 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp44 = tmp43 * tmp43
    tmp45 = tmp42 + tmp44
    tmp46 = libdevice.sqrt(tmp45)
    tmp47 = tmp34 * tmp46
    tmp48 = tmp31 / tmp47
    tmp49 = triton_helpers.maximum(tmp48, tmp23)
    tmp50 = tl.full(tmp49.shape, 0.0, tmp49.dtype)
    tmp51 = tl.where(tmp30, tmp49, tmp50)
    tmp52 = tmp0 >= tmp28
    tmp53 = tl.full([1], 3, tl.int64)
    tmp54 = tmp0 < tmp53
    tmp55 = tmp52 & tmp54
    tmp56 = tl.load(in_ptr6 + x0, tmp55 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp59 = libdevice.sqrt(tmp58)
    tmp60 = tl.load(in_ptr8 + 4 * x0, tmp55 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp61 = tmp60 * tmp60
    tmp62 = tl.load(in_ptr8 + (1 + 4 * x0), tmp55 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp63 = tmp62 * tmp62
    tmp64 = tmp61 + tmp63
    tmp65 = tl.load(in_ptr8 + (2 + 4 * x0), tmp55 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp66 = tmp65 * tmp65
    tmp67 = tmp64 + tmp66
    tmp68 = tl.load(in_ptr8 + (3 + 4 * x0), tmp55 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp69 = tmp68 * tmp68
    tmp70 = tmp67 + tmp69
    tmp71 = libdevice.sqrt(tmp70)
    tmp72 = tmp59 * tmp71
    tmp73 = tmp56 / tmp72
    tmp74 = triton_helpers.maximum(tmp73, tmp23)
    tmp75 = tl.full(tmp74.shape, 0.0, tmp74.dtype)
    tmp76 = tl.where(tmp55, tmp74, tmp75)
    tmp77 = tmp0 >= tmp53
    tl.full([1], 4, tl.int64)
    tmp80 = tl.load(in_ptr9 + x0, tmp77 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp83 = libdevice.sqrt(tmp82)
    tmp84 = tl.load(in_ptr11 + 4 * x0, tmp77 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp85 = tmp84 * tmp84
    tmp86 = tl.load(in_ptr11 + (1 + 4 * x0), tmp77 & xmask, eviction_policy
        ='evict_last', other=0.0)
    tmp87 = tmp86 * tmp86
    tmp88 = tmp85 + tmp87
    tmp89 = tl.load(in_ptr11 + (2 + 4 * x0), tmp77 & xmask, eviction_policy
        ='evict_last', other=0.0)
    tmp90 = tmp89 * tmp89
    tmp91 = tmp88 + tmp90
    tmp92 = tl.load(in_ptr11 + (3 + 4 * x0), tmp77 & xmask, eviction_policy
        ='evict_last', other=0.0)
    tmp93 = tmp92 * tmp92
    tmp94 = tmp91 + tmp93
    tmp95 = libdevice.sqrt(tmp94)
    tmp96 = tmp83 * tmp95
    tmp97 = tmp80 / tmp96
    tmp98 = triton_helpers.maximum(tmp97, tmp23)
    tmp99 = tl.full(tmp98.shape, 0.0, tmp98.dtype)
    tmp100 = tl.where(tmp77, tmp98, tmp99)
    tmp101 = tl.where(tmp55, tmp76, tmp100)
    tmp102 = tl.where(tmp30, tmp51, tmp101)
    tmp103 = tl.where(tmp4, tmp26, tmp102)
    tl.store(out_ptr0 + x2, tmp103, xmask)


@triton.jit
def triton_per_fused_cat_linalg_vector_norm_mean_stack_8(in_ptr0, out_ptr0,
    out_ptr1, out_ptr2, out_ptr3, out_ptr4, out_ptr5, out_ptr6, out_ptr7,
    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 + (20 + r0), None)
    tmp5 = tl.load(in_ptr0 + (24 + r0), None)
    tmp10 = tl.load(in_ptr0 + (28 + r0), None)
    tmp44 = tl.load(in_ptr0 + (16 + r0), None)
    tmp1 = tmp0 * tmp0
    tmp2 = tl.broadcast_to(tmp1, [XBLOCK, RBLOCK])
    tmp4 = tl.sum(tmp2, 1)[:, None]
    tmp6 = tmp5 * tmp5
    tmp7 = tl.broadcast_to(tmp6, [XBLOCK, RBLOCK])
    tmp9 = tl.sum(tmp7, 1)[:, None]
    tmp11 = tmp10 * tmp10
    tmp12 = tl.broadcast_to(tmp11, [XBLOCK, RBLOCK])
    tmp14 = tl.sum(tmp12, 1)[:, None]
    tmp15 = tl.full([1, 1], 0, tl.int64)
    tmp17 = tl.full([1, 1], 2, tl.int64)
    tmp18 = tmp15 < tmp17
    tmp19 = tl.load(in_ptr0 + tl.broadcast_to(16 + r0 + 4 * 0, [XBLOCK,
        RBLOCK]), tmp18, other=0.0)
    tmp20 = tmp15 >= tmp17
    tl.full([1, 1], 3, tl.int64)
    tmp23 = tl.load(in_ptr0 + tl.broadcast_to(28 + r0, [XBLOCK, RBLOCK]),
        tmp20, other=0.0)
    tmp24 = tl.where(tmp18, tmp19, tmp23)
    tmp25 = tl.full([1, 1], 1, tl.int64)
    tmp27 = tmp25 < tmp17
    tmp28 = tl.load(in_ptr0 + tl.broadcast_to(16 + r0 + 4 * 1, [XBLOCK,
        RBLOCK]), tmp27, other=0.0)
    tmp29 = tmp25 >= tmp17
    tmp31 = tl.load(in_ptr0 + tl.broadcast_to(28 + r0, [XBLOCK, RBLOCK]),
        tmp29, other=0.0)
    tmp32 = tl.where(tmp27, tmp28, tmp31)
    tmp33 = tmp24 + tmp32
    tmp35 = tmp17 < tmp17
    tmp36 = tl.load(in_ptr0 + tl.broadcast_to(16 + r0 + 4 * 2, [XBLOCK,
        RBLOCK]), tmp35, other=0.0)
    tmp37 = tmp17 >= tmp17
    tmp39 = tl.load(in_ptr0 + tl.broadcast_to(28 + r0, [XBLOCK, RBLOCK]),
        tmp37, other=0.0)
    tmp40 = tl.where(tmp35, tmp36, tmp39)
    tmp41 = tmp33 + tmp40
    tmp42 = 3.0
    tmp43 = tmp41 / tmp42
    tmp45 = tmp44 + tmp0
    tmp46 = tmp45 + tmp5
    tmp47 = tmp46 + tmp10
    tmp48 = 4.0
    tmp49 = tmp47 / tmp48
    tl.store(out_ptr3 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp43, None)
    tl.store(out_ptr4 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp49, None)
    tl.store(out_ptr5 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp49, None)
    tl.store(out_ptr6 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp49, None)
    tl.store(out_ptr7 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp49, None)
    tl.store(out_ptr0 + tl.full([XBLOCK, 1], 0, tl.int32), tmp4, None)
    tl.store(out_ptr1 + tl.full([XBLOCK, 1], 0, tl.int32), tmp9, None)
    tl.store(out_ptr2 + tl.full([XBLOCK, 1], 0, tl.int32), tmp14, None)


@triton.jit
def triton_per_fused_cat_linalg_vector_norm_mean_stack_9(in_ptr0, out_ptr0,
    out_ptr1, out_ptr2, out_ptr3, out_ptr4, 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 + (40 + r0), None)
    tmp5 = tl.load(in_ptr0 + (44 + r0), None)
    tmp39 = tl.load(in_ptr0 + (32 + r0), None)
    tmp40 = tl.load(in_ptr0 + (36 + r0), None)
    tmp1 = tmp0 * tmp0
    tmp2 = tl.broadcast_to(tmp1, [XBLOCK, RBLOCK])
    tmp4 = tl.sum(tmp2, 1)[:, None]
    tmp6 = tmp5 * tmp5
    tmp7 = tl.broadcast_to(tmp6, [XBLOCK, RBLOCK])
    tmp9 = tl.sum(tmp7, 1)[:, None]
    tmp10 = tl.full([1, 1], 0, tl.int64)
    tmp12 = tl.full([1, 1], 2, tl.int64)
    tmp13 = tmp10 < tmp12
    tmp14 = tl.load(in_ptr0 + tl.broadcast_to(32 + r0 + 4 * 0, [XBLOCK,
        RBLOCK]), tmp13, other=0.0)
    tmp15 = tmp10 >= tmp12
    tl.full([1, 1], 3, tl.int64)
    tmp18 = tl.load(in_ptr0 + tl.broadcast_to(44 + r0, [XBLOCK, RBLOCK]),
        tmp15, other=0.0)
    tmp19 = tl.where(tmp13, tmp14, tmp18)
    tmp20 = tl.full([1, 1], 1, tl.int64)
    tmp22 = tmp20 < tmp12
    tmp23 = tl.load(in_ptr0 + tl.broadcast_to(32 + r0 + 4 * 1, [XBLOCK,
        RBLOCK]), tmp22, other=0.0)
    tmp24 = tmp20 >= tmp12
    tmp26 = tl.load(in_ptr0 + tl.broadcast_to(44 + r0, [XBLOCK, RBLOCK]),
        tmp24, other=0.0)
    tmp27 = tl.where(tmp22, tmp23, tmp26)
    tmp28 = tmp19 + tmp27
    tmp30 = tmp12 < tmp12
    tmp31 = tl.load(in_ptr0 + tl.broadcast_to(32 + r0 + 4 * 2, [XBLOCK,
        RBLOCK]), tmp30, other=0.0)
    tmp32 = tmp12 >= tmp12
    tmp34 = tl.load(in_ptr0 + tl.broadcast_to(44 + r0, [XBLOCK, RBLOCK]),
        tmp32, other=0.0)
    tmp35 = tl.where(tmp30, tmp31, tmp34)
    tmp36 = tmp28 + tmp35
    tmp37 = 3.0
    tmp38 = tmp36 / tmp37
    tmp41 = tmp39 + tmp40
    tmp42 = tmp41 + tmp0
    tmp43 = tmp42 + tmp5
    tmp44 = 4.0
    tmp45 = tmp43 / tmp44
    tl.store(out_ptr2 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp38, None)
    tl.store(out_ptr3 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp45, None)
    tl.store(out_ptr4 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp45, None)
    tl.store(out_ptr0 + tl.full([XBLOCK, 1], 0, tl.int32), tmp4, None)
    tl.store(out_ptr1 + tl.full([XBLOCK, 1], 0, tl.int32), tmp9, None)


@triton.jit
def triton_per_fused_linalg_vector_norm_10(in_ptr0, out_ptr0, 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 + (52 + r0), None)
    tmp1 = tmp0 * tmp0
    tmp2 = tl.broadcast_to(tmp1, [XBLOCK, RBLOCK])
    tmp4 = tl.sum(tmp2, 1)[:, None]
    tl.store(out_ptr0 + tl.full([XBLOCK, 1], 0, tl.int32), tmp4, None)


@triton.jit
def triton_per_fused_cat_linalg_vector_norm_mean_11(in_ptr0, out_ptr0,
    out_ptr1, 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
    tmp29 = tl.load(in_ptr0 + (60 + r0), None)
    tmp0 = tl.full([1, 1], 0, tl.int64)
    tmp2 = tl.full([1, 1], 2, tl.int64)
    tmp3 = tmp0 < tmp2
    tmp4 = tl.load(in_ptr0 + tl.broadcast_to(48 + r0 + 4 * 0, [XBLOCK,
        RBLOCK]), tmp3, other=0.0)
    tmp5 = tmp0 >= tmp2
    tl.full([1, 1], 3, tl.int64)
    tmp8 = tl.load(in_ptr0 + tl.broadcast_to(60 + r0, [XBLOCK, RBLOCK]),
        tmp5, other=0.0)
    tmp9 = tl.where(tmp3, tmp4, tmp8)
    tmp10 = tl.full([1, 1], 1, tl.int64)
    tmp12 = tmp10 < tmp2
    tmp13 = tl.load(in_ptr0 + tl.broadcast_to(48 + r0 + 4 * 1, [XBLOCK,
        RBLOCK]), tmp12, other=0.0)
    tmp14 = tmp10 >= tmp2
    tmp16 = tl.load(in_ptr0 + tl.broadcast_to(60 + r0, [XBLOCK, RBLOCK]),
        tmp14, other=0.0)
    tmp17 = tl.where(tmp12, tmp13, tmp16)
    tmp18 = tmp9 + tmp17
    tmp20 = tmp2 < tmp2
    tmp21 = tl.load(in_ptr0 + tl.broadcast_to(48 + r0 + 4 * 2, [XBLOCK,
        RBLOCK]), tmp20, other=0.0)
    tmp22 = tmp2 >= tmp2
    tmp24 = tl.load(in_ptr0 + tl.broadcast_to(60 + r0, [XBLOCK, RBLOCK]),
        tmp22, other=0.0)
    tmp25 = tl.where(tmp20, tmp21, tmp24)
    tmp26 = tmp18 + tmp25
    tmp27 = 3.0
    tmp28 = tmp26 / tmp27
    tmp30 = tmp29 * tmp29
    tmp31 = tl.broadcast_to(tmp30, [XBLOCK, RBLOCK])
    tmp33 = tl.sum(tmp31, 1)[:, None]
    tl.store(out_ptr0 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp28, None)
    tl.store(out_ptr1 + tl.full([XBLOCK, 1], 0, tl.int32), tmp33, None)


@triton.jit
def triton_per_fused_linalg_vector_norm_12(in_ptr0, out_ptr0, 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 + (56 + r0), None)
    tmp1 = tmp0 * tmp0
    tmp2 = tl.broadcast_to(tmp1, [XBLOCK, RBLOCK])
    tmp4 = tl.sum(tmp2, 1)[:, None]
    tl.store(out_ptr0 + tl.full([XBLOCK, 1], 0, tl.int32), tmp4, None)


@triton.jit
def triton_poi_fused_stack_13(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
    x0 = xindex
    tmp0 = x0
    tl.full([1], 0, tl.int64)
    tmp3 = tl.full([1], 4, tl.int64)
    tmp4 = tmp0 < tmp3
    tmp5 = tl.load(in_ptr0 + x0, tmp4 & xmask, eviction_policy='evict_last',
        other=0.0)
    tmp6 = tl.load(in_ptr0 + (4 + x0), tmp4 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp7 = tmp5 + tmp6
    tmp8 = tl.load(in_ptr0 + (8 + x0), tmp4 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp9 = tmp7 + tmp8
    tmp10 = tl.load(in_ptr0 + (12 + x0), tmp4 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp11 = tmp9 + tmp10
    tmp12 = 4.0
    tmp13 = tmp11 / tmp12
    tmp14 = tl.full(tmp13.shape, 0.0, tmp13.dtype)
    tmp15 = tl.where(tmp4, tmp13, tmp14)
    tmp16 = tmp0 >= tmp3
    tmp17 = tl.full([1], 8, tl.int64)
    tmp18 = tmp0 < tmp17
    tmp19 = tmp16 & tmp18
    tmp20 = tl.load(in_ptr0 + (16 + (-4 + x0)), tmp19 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp21 = tl.load(in_ptr0 + (20 + (-4 + x0)), tmp19 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp22 = tmp20 + tmp21
    tmp23 = tl.load(in_ptr0 + (24 + (-4 + x0)), tmp19 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp24 = tmp22 + tmp23
    tmp25 = tl.load(in_ptr0 + (28 + (-4 + x0)), tmp19 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp26 = tmp24 + tmp25
    tmp27 = tmp26 / tmp12
    tmp28 = tl.full(tmp27.shape, 0.0, tmp27.dtype)
    tmp29 = tl.where(tmp19, tmp27, tmp28)
    tmp30 = tmp0 >= tmp17
    tmp31 = tl.full([1], 12, tl.int64)
    tmp32 = tmp0 < tmp31
    tmp33 = tmp30 & tmp32
    tmp34 = tl.load(in_ptr0 + (32 + (-8 + x0)), tmp33 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp35 = tl.load(in_ptr0 + (36 + (-8 + x0)), tmp33 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp36 = tmp34 + tmp35
    tmp37 = tl.load(in_ptr0 + (40 + (-8 + x0)), tmp33 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp38 = tmp36 + tmp37
    tmp39 = tl.load(in_ptr0 + (44 + (-8 + x0)), tmp33 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp40 = tmp38 + tmp39
    tmp41 = tmp40 / tmp12
    tmp42 = tl.full(tmp41.shape, 0.0, tmp41.dtype)
    tmp43 = tl.where(tmp33, tmp41, tmp42)
    tmp44 = tmp0 >= tmp31
    tl.full([1], 16, tl.int64)
    tmp47 = tl.load(in_ptr0 + (48 + (-12 + x0)), tmp44 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp48 = tl.load(in_ptr0 + (52 + (-12 + x0)), tmp44 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp49 = tmp47 + tmp48
    tmp50 = tl.load(in_ptr0 + (56 + (-12 + x0)), tmp44 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp51 = tmp49 + tmp50
    tmp52 = 3.0
    tmp53 = tmp51 / tmp52
    tmp54 = tl.full(tmp53.shape, 0.0, tmp53.dtype)
    tmp55 = tl.where(tmp44, tmp53, tmp54)
    tmp56 = tl.where(tmp33, tmp43, tmp55)
    tmp57 = tl.where(tmp19, tmp29, tmp56)
    tmp58 = tl.where(tmp4, tmp15, tmp57)
    tl.store(out_ptr0 + x0, tmp58, xmask)


@triton.jit
def triton_per_fused__log_softmax__log_softmax_backward_data_stack_14(in_ptr0,
    in_ptr1, in_ptr2, out_ptr32, out_ptr33, out_ptr34, out_ptr35, out_ptr36,
    out_ptr37, out_ptr38, out_ptr39, out_ptr40, out_ptr41, out_ptr42,
    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 + 0)
    tmp1 = tl.broadcast_to(tmp0, [XBLOCK, RBLOCK])
    tmp2 = tl.load(in_ptr1 + r0, None)
    tmp4 = tl.load(in_ptr2 + 0)
    tmp5 = tl.broadcast_to(tmp4, [XBLOCK, RBLOCK])
    tmp15 = tl.load(in_ptr1 + (4 + r0), None)
    tmp26 = tl.load(in_ptr1 + (8 + r0), None)
    tmp37 = tl.load(in_ptr1 + (12 + r0), None)
    tmp48 = tl.load(in_ptr1 + (16 + r0), None)
    tmp59 = tl.load(in_ptr1 + (20 + r0), None)
    tmp70 = tl.load(in_ptr1 + (24 + r0), None)
    tmp81 = tl.load(in_ptr1 + (28 + r0), None)
    tmp92 = tl.load(in_ptr1 + (32 + r0), None)
    tmp103 = tl.load(in_ptr1 + (36 + r0), None)
    tmp114 = tl.load(in_ptr1 + (40 + r0), None)
    tmp125 = tl.load(in_ptr1 + (44 + r0), None)
    tmp136 = tl.load(in_ptr1 + (48 + r0), None)
    tmp147 = tl.load(in_ptr1 + (52 + r0), None)
    tmp158 = tl.load(in_ptr1 + (56 + r0), None)
    tmp169 = tl.load(in_ptr1 + (60 + r0), None)
    tmp185 = tl.load(in_ptr1 + 0)
    tmp186 = tl.broadcast_to(tmp185, [XBLOCK, RBLOCK])
    tmp199 = tl.load(in_ptr1 + 4)
    tmp200 = tl.broadcast_to(tmp199, [XBLOCK, RBLOCK])
    tmp213 = tl.load(in_ptr1 + 8)
    tmp214 = tl.broadcast_to(tmp213, [XBLOCK, RBLOCK])
    tmp226 = tl.load(in_ptr1 + 12)
    tmp227 = tl.broadcast_to(tmp226, [XBLOCK, RBLOCK])
    tmp239 = tl.load(in_ptr1 + 17)
    tmp240 = tl.broadcast_to(tmp239, [XBLOCK, RBLOCK])
    tmp249 = tl.load(in_ptr1 + 21)
    tmp250 = tl.broadcast_to(tmp249, [XBLOCK, RBLOCK])
    tmp259 = tl.load(in_ptr1 + 25)
    tmp260 = tl.broadcast_to(tmp259, [XBLOCK, RBLOCK])
    tmp269 = tl.load(in_ptr1 + 29)
    tmp270 = tl.broadcast_to(tmp269, [XBLOCK, RBLOCK])
    tmp282 = tl.load(in_ptr1 + 34)
    tmp283 = tl.broadcast_to(tmp282, [XBLOCK, RBLOCK])
    tmp292 = tl.load(in_ptr1 + 38)
    tmp293 = tl.broadcast_to(tmp292, [XBLOCK, RBLOCK])
    tmp302 = tl.load(in_ptr1 + 42)
    tmp303 = tl.broadcast_to(tmp302, [XBLOCK, RBLOCK])
    tmp312 = tl.load(in_ptr1 + 46)
    tmp313 = tl.broadcast_to(tmp312, [XBLOCK, RBLOCK])
    tmp3 = tmp1 * tmp2
    tmp6 = tmp3 + tmp5
    tmp7 = tl.broadcast_to(tmp6, [XBLOCK, RBLOCK])
    tmp9 = triton_helpers.max2(tmp7, 1)[:, None]
    tmp10 = tmp6 - tmp9
    tmp11 = tl_math.exp(tmp10)
    tmp12 = tl.broadcast_to(tmp11, [XBLOCK, RBLOCK])
    tmp14 = tl.sum(tmp12, 1)[:, None]
    tmp16 = tmp1 * tmp15
    tmp17 = tmp16 + tmp5
    tmp18 = tl.broadcast_to(tmp17, [XBLOCK, RBLOCK])
    tmp20 = triton_helpers.max2(tmp18, 1)[:, None]
    tmp21 = tmp17 - tmp20
    tmp22 = tl_math.exp(tmp21)
    tmp23 = tl.broadcast_to(tmp22, [XBLOCK, RBLOCK])
    tmp25 = tl.sum(tmp23, 1)[:, None]
    tmp27 = tmp1 * tmp26
    tmp28 = tmp27 + tmp5
    tmp29 = tl.broadcast_to(tmp28, [XBLOCK, RBLOCK])
    tmp31 = triton_helpers.max2(tmp29, 1)[:, None]
    tmp32 = tmp28 - tmp31
    tmp33 = tl_math.exp(tmp32)
    tmp34 = tl.broadcast_to(tmp33, [XBLOCK, RBLOCK])
    tmp36 = tl.sum(tmp34, 1)[:, None]
    tmp38 = tmp1 * tmp37
    tmp39 = tmp38 + tmp5
    tmp40 = tl.broadcast_to(tmp39, [XBLOCK, RBLOCK])
    tmp42 = triton_helpers.max2(tmp40, 1)[:, None]
    tmp43 = tmp39 - tmp42
    tmp44 = tl_math.exp(tmp43)
    tmp45 = tl.broadcast_to(tmp44, [XBLOCK, RBLOCK])
    tmp47 = tl.sum(tmp45, 1)[:, None]
    tmp49 = tmp1 * tmp48
    tmp50 = tmp49 + tmp5
    tmp51 = tl.broadcast_to(tmp50, [XBLOCK, RBLOCK])
    tmp53 = triton_helpers.max2(tmp51, 1)[:, None]
    tmp54 = tmp50 - tmp53
    tmp55 = tl_math.exp(tmp54)
    tmp56 = tl.broadcast_to(tmp55, [XBLOCK, RBLOCK])
    tmp58 = tl.sum(tmp56, 1)[:, None]
    tmp60 = tmp1 * tmp59
    tmp61 = tmp60 + tmp5
    tmp62 = tl.broadcast_to(tmp61, [XBLOCK, RBLOCK])
    tmp64 = triton_helpers.max2(tmp62, 1)[:, None]
    tmp65 = tmp61 - tmp64
    tmp66 = tl_math.exp(tmp65)
    tmp67 = tl.broadcast_to(tmp66, [XBLOCK, RBLOCK])
    tmp69 = tl.sum(tmp67, 1)[:, None]
    tmp71 = tmp1 * tmp70
    tmp72 = tmp71 + tmp5
    tmp73 = tl.broadcast_to(tmp72, [XBLOCK, RBLOCK])
    tmp75 = triton_helpers.max2(tmp73, 1)[:, None]
    tmp76 = tmp72 - tmp75
    tmp77 = tl_math.exp(tmp76)
    tmp78 = tl.broadcast_to(tmp77, [XBLOCK, RBLOCK])
    tmp80 = tl.sum(tmp78, 1)[:, None]
    tmp82 = tmp1 * tmp81
    tmp83 = tmp82 + tmp5
    tmp84 = tl.broadcast_to(tmp83, [XBLOCK, RBLOCK])
    tmp86 = triton_helpers.max2(tmp84, 1)[:, None]
    tmp87 = tmp83 - tmp86
    tmp88 = tl_math.exp(tmp87)
    tmp89 = tl.broadcast_to(tmp88, [XBLOCK, RBLOCK])
    tmp91 = tl.sum(tmp89, 1)[:, None]
    tmp93 = tmp1 * tmp92
    tmp94 = tmp93 + tmp5
    tmp95 = tl.broadcast_to(tmp94, [XBLOCK, RBLOCK])
    tmp97 = triton_helpers.max2(tmp95, 1)[:, None]
    tmp98 = tmp94 - tmp97
    tmp99 = tl_math.exp(tmp98)
    tmp100 = tl.broadcast_to(tmp99, [XBLOCK, RBLOCK])
    tmp102 = tl.sum(tmp100, 1)[:, None]
    tmp104 = tmp1 * tmp103
    tmp105 = tmp104 + tmp5
    tmp106 = tl.broadcast_to(tmp105, [XBLOCK, RBLOCK])
    tmp108 = triton_helpers.max2(tmp106, 1)[:, None]
    tmp109 = tmp105 - tmp108
    tmp110 = tl_math.exp(tmp109)
    tmp111 = tl.broadcast_to(tmp110, [XBLOCK, RBLOCK])
    tmp113 = tl.sum(tmp111, 1)[:, None]
    tmp115 = tmp1 * tmp114
    tmp116 = tmp115 + tmp5
    tmp117 = tl.broadcast_to(tmp116, [XBLOCK, RBLOCK])
    tmp119 = triton_helpers.max2(tmp117, 1)[:, None]
    tmp120 = tmp116 - tmp119
    tmp121 = tl_math.exp(tmp120)
    tmp122 = tl.broadcast_to(tmp121, [XBLOCK, RBLOCK])
    tmp124 = tl.sum(tmp122, 1)[:, None]
    tmp126 = tmp1 * tmp125
    tmp127 = tmp126 + tmp5
    tmp128 = tl.broadcast_to(tmp127, [XBLOCK, RBLOCK])
    tmp130 = triton_helpers.max2(tmp128, 1)[:, None]
    tmp131 = tmp127 - tmp130
    tmp132 = tl_math.exp(tmp131)
    tmp133 = tl.broadcast_to(tmp132, [XBLOCK, RBLOCK])
    tmp135 = tl.sum(tmp133, 1)[:, None]
    tmp137 = tmp1 * tmp136
    tmp138 = tmp137 + tmp5
    tmp139 = tl.broadcast_to(tmp138, [XBLOCK, RBLOCK])
    tmp141 = triton_helpers.max2(tmp139, 1)[:, None]
    tmp142 = tmp138 - tmp141
    tmp143 = tl_math.exp(tmp142)
    tmp144 = tl.broadcast_to(tmp143, [XBLOCK, RBLOCK])
    tmp146 = tl.sum(tmp144, 1)[:, None]
    tmp148 = tmp1 * tmp147
    tmp149 = tmp148 + tmp5
    tmp150 = tl.broadcast_to(tmp149, [XBLOCK, RBLOCK])
    tmp152 = triton_helpers.max2(tmp150, 1)[:, None]
    tmp153 = tmp149 - tmp152
    tmp154 = tl_math.exp(tmp153)
    tmp155 = tl.broadcast_to(tmp154, [XBLOCK, RBLOCK])
    tmp157 = tl.sum(tmp155, 1)[:, None]
    tmp159 = tmp1 * tmp158
    tmp160 = tmp159 + tmp5
    tmp161 = tl.broadcast_to(tmp160, [XBLOCK, RBLOCK])
    tmp163 = triton_helpers.max2(tmp161, 1)[:, None]
    tmp164 = tmp160 - tmp163
    tmp165 = tl_math.exp(tmp164)
    tmp166 = tl.broadcast_to(tmp165, [XBLOCK, RBLOCK])
    tmp168 = tl.sum(tmp166, 1)[:, None]
    tmp170 = tmp1 * tmp169
    tmp171 = tmp170 + tmp5
    tmp172 = tl.broadcast_to(tmp171, [XBLOCK, RBLOCK])
    tmp174 = triton_helpers.max2(tmp172, 1)[:, None]
    tmp175 = tmp171 - tmp174
    tmp176 = tl_math.exp(tmp175)
    tmp177 = tl.broadcast_to(tmp176, [XBLOCK, RBLOCK])
    tmp179 = tl.sum(tmp177, 1)[:, None]
    tmp180 = r0
    tl.full([1, 1], 0, tl.int64)
    tmp183 = tl.full([1, 1], 1, tl.int64)
    tmp184 = tmp180 < tmp183
    tmp187 = tmp1 * tmp186
    tmp188 = tmp187 + tmp5
    tmp189 = tmp188 - tmp9
    tmp190 = tl_math.log(tmp14)
    tmp191 = tmp189 - tmp190
    tmp192 = -tmp191
    tmp193 = tl.full(tmp192.shape, 0.0, tmp192.dtype)
    tmp194 = tl.where(tmp184, tmp192, tmp193)
    tmp195 = tmp180 >= tmp183
    tmp196 = tl.full([1, 1], 2, tl.int64)
    tmp197 = tmp180 < tmp196
    tmp198 = tmp195 & tmp197
    tmp201 = tmp1 * tmp200
    tmp202 = tmp201 + tmp5
    tmp203 = tmp202 - tmp20
    tmp204 = tl_math.log(tmp25)
    tmp205 = tmp203 - tmp204
    tmp206 = -tmp205
    tmp207 = tl.full(tmp206.shape, 0.0, tmp206.dtype)
    tmp208 = tl.where(tmp198, tmp206, tmp207)
    tmp209 = tmp180 >= tmp196
    tmp210 = tl.full([1, 1], 3, tl.int64)
    tmp211 = tmp180 < tmp210
    tmp212 = tmp209 & tmp211
    tmp215 = tmp1 * tmp214
    tmp216 = tmp215 + tmp5
    tmp217 = tmp216 - tmp31
    tmp218 = tl_math.log(tmp36)
    tmp219 = tmp217 - tmp218
    tmp220 = -tmp219
    tmp221 = tl.full(tmp220.shape, 0.0, tmp220.dtype)
    tmp222 = tl.where(tmp212, tmp220, tmp221)
    tmp223 = tmp180 >= tmp210
    tl.full([1, 1], 4, tl.int64)
    tmp228 = tmp1 * tmp227
    tmp229 = tmp228 + tmp5
    tmp230 = tmp229 - tmp42
    tmp231 = tl_math.log(tmp47)
    tmp232 = tmp230 - tmp231
    tmp233 = -tmp232
    tmp234 = tl.full(tmp233.shape, 0.0, tmp233.dtype)
    tmp235 = tl.where(tmp223, tmp233, tmp234)
    tmp236 = tl.where(tmp212, tmp222, tmp235)
    tmp237 = tl.where(tmp198, tmp208, tmp236)
    tmp238 = tl.where(tmp184, tmp194, tmp237)
    tmp241 = tmp1 * tmp240
    tmp242 = tmp241 + tmp5
    tmp243 = tmp242 - tmp53
    tmp244 = tl_math.log(tmp58)
    tmp245 = tmp243 - tmp244
    tmp246 = -tmp245
    tmp247 = tl.full(tmp246.shape, 0.0, tmp246.dtype)
    tmp248 = tl.where(tmp184, tmp246, tmp247)
    tmp251 = tmp1 * tmp250
    tmp252 = tmp251 + tmp5
    tmp253 = tmp252 - tmp64
    tmp254 = tl_math.log(tmp69)
    tmp255 = tmp253 - tmp254
    tmp256 = -tmp255
    tmp257 = tl.full(tmp256.shape, 0.0, tmp256.dtype)
    tmp258 = tl.where(tmp198, tmp256, tmp257)
    tmp261 = tmp1 * tmp260
    tmp262 = tmp261 + tmp5
    tmp263 = tmp262 - tmp75
    tmp264 = tl_math.log(tmp80)
    tmp265 = tmp263 - tmp264
    tmp266 = -tmp265
    tmp267 = tl.full(tmp266.shape, 0.0, tmp266.dtype)
    tmp268 = tl.where(tmp212, tmp266, tmp267)
    tmp271 = tmp1 * tmp270
    tmp272 = tmp271 + tmp5
    tmp273 = tmp272 - tmp86
    tmp274 = tl_math.log(tmp91)
    tmp275 = tmp273 - tmp274
    tmp276 = -tmp275
    tmp277 = tl.full(tmp276.shape, 0.0, tmp276.dtype)
    tmp278 = tl.where(tmp223, tmp276, tmp277)
    tmp279 = tl.where(tmp212, tmp268, tmp278)
    tmp280 = tl.where(tmp198, tmp258, tmp279)
    tmp281 = tl.where(tmp184, tmp248, tmp280)
    tmp284 = tmp1 * tmp283
    tmp285 = tmp284 + tmp5
    tmp286 = tmp285 - tmp97
    tmp287 = tl_math.log(tmp102)
    tmp288 = tmp286 - tmp287
    tmp289 = -tmp288
    tmp290 = tl.full(tmp289.shape, 0.0, tmp289.dtype)
    tmp291 = tl.where(tmp184, tmp289, tmp290)
    tmp294 = tmp1 * tmp293
    tmp295 = tmp294 + tmp5
    tmp296 = tmp295 - tmp108
    tmp297 = tl_math.log(tmp113)
    tmp298 = tmp296 - tmp297
    tmp299 = -tmp298
    tmp300 = tl.full(tmp299.shape, 0.0, tmp299.dtype)
    tmp301 = tl.where(tmp198, tmp299, tmp300)
    tmp304 = tmp1 * tmp303
    tmp305 = tmp304 + tmp5
    tmp306 = tmp305 - tmp119
    tmp307 = tl_math.log(tmp124)
    tmp308 = tmp306 - tmp307
    tmp309 = -tmp308
    tmp310 = tl.full(tmp309.shape, 0.0, tmp309.dtype)
    tmp311 = tl.where(tmp212, tmp309, tmp310)
    tmp314 = tmp1 * tmp313
    tmp315 = tmp314 + tmp5
    tmp316 = tmp315 - tmp130
    tmp317 = tl_math.log(tmp135)
    tmp318 = tmp316 - tmp317
    tmp319 = -tmp318
    tmp320 = tl.full(tmp319.shape, 0.0, tmp319.dtype)
    tmp321 = tl.where(tmp223, tmp319, tmp320)
    tmp322 = tl.where(tmp212, tmp311, tmp321)
    tmp323 = tl.where(tmp198, tmp301, tmp322)
    tmp324 = tl.where(tmp184, tmp291, tmp323)
    tmp325 = tl_math.log(tmp146)
    tmp326 = tmp142 - tmp325
    tmp327 = tl_math.exp(tmp326)
    tmp328 = tl_math.log(tmp157)
    tmp329 = tmp153 - tmp328
    tmp330 = tl_math.exp(tmp329)
    tmp331 = tl_math.log(tmp168)
    tmp332 = tmp164 - tmp331
    tmp333 = tl_math.exp(tmp332)
    tmp334 = tl_math.log(tmp179)
    tmp335 = tmp175 - tmp334
    tmp336 = tl_math.exp(tmp335)
    tl.store(out_ptr32 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp238, None)
    tl.store(out_ptr33 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp281, None)
    tl.store(out_ptr34 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp324, None)
    tl.store(out_ptr35 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp326, None)
    tl.store(out_ptr36 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp327, None)
    tl.store(out_ptr37 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp329, None)
    tl.store(out_ptr38 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp330, None)
    tl.store(out_ptr39 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp332, None)
    tl.store(out_ptr40 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp333, None)
    tl.store(out_ptr41 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp335, None)
    tl.store(out_ptr42 + tl.broadcast_to(r0, [XBLOCK, RBLOCK]), tmp336, None)


@triton.jit
def triton_poi_fused_stack_15(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
    tmp5 = tl.load(in_ptr0 + 3)
    tmp6 = tl.broadcast_to(tmp5, [XBLOCK])
    tmp14 = tl.load(in_ptr1 + 3)
    tmp15 = tl.broadcast_to(tmp14, [XBLOCK])
    tmp23 = tl.load(in_ptr2 + 3)
    tmp24 = tl.broadcast_to(tmp23, [XBLOCK])
    tmp31 = tl.load(in_ptr3 + 3)
    tmp32 = tl.broadcast_to(tmp31, [XBLOCK])
    tmp0 = x0
    tl.full([1], 0, tl.int64)
    tmp3 = tl.full([1], 1, tl.int64)
    tmp4 = tmp0 < tmp3
    tmp7 = -tmp6
    tmp8 = tl.full(tmp7.shape, 0.0, tmp7.dtype)
    tmp9 = tl.where(tmp4, tmp7, tmp8)
    tmp10 = tmp0 >= tmp3
    tmp11 = tl.full([1], 2, tl.int64)
    tmp12 = tmp0 < tmp11
    tmp13 = tmp10 & tmp12
    tmp16 = -tmp15
    tmp17 = tl.full(tmp16.shape, 0.0, tmp16.dtype)
    tmp18 = tl.where(tmp13, tmp16, tmp17)
    tmp19 = tmp0 >= tmp11
    tmp20 = tl.full([1], 3, tl.int64)
    tmp21 = tmp0 < tmp20
    tmp22 = tmp19 & tmp21
    tmp25 = -tmp24
    tmp26 = tl.full(tmp25.shape, 0.0, tmp25.dtype)
    tmp27 = tl.where(tmp22, tmp25, tmp26)
    tmp28 = tmp0 >= tmp20
    tl.full([1], 4, tl.int64)
    tmp33 = -tmp32
    tmp34 = tl.full(tmp33.shape, 0.0, tmp33.dtype)
    tmp35 = tl.where(tmp28, tmp33, tmp34)
    tmp36 = tl.where(tmp22, tmp27, tmp35)
    tmp37 = tl.where(tmp13, tmp18, tmp36)
    tmp38 = tl.where(tmp4, tmp9, tmp37)
    tl.store(out_ptr0 + x0, tmp38, xmask)


@triton.jit
def triton_per_fused_mean_16(in_out_ptr0, in_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 = tl.broadcast_to(tmp0, [XBLOCK, RBLOCK])
    tmp3 = tl.sum(tmp1, 1)[:, None]
    tmp4 = 16.0
    tmp5 = tmp3 / tmp4
    tl.debug_barrier()
    tl.store(in_out_ptr0 + tl.full([XBLOCK, 1], 0, tl.int32), tmp5, None)


def call(args):
    primals_1, primals_2, primals_3 = args
    args.clear()
    assert_size_stride(primals_1, (4, 4, 4), (16, 4, 1))
    assert_size_stride(primals_2, (), ())
    assert_size_stride(primals_3, (), ())
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        buf0 = empty_strided_cuda((16,), (1,), torch.float32)
        buf17 = empty_strided_cuda((16,), (1,), torch.float32)
        buf21 = empty_strided_cuda((16,), (1,), torch.float32)
        buf38 = empty_strided_cuda((16,), (1,), torch.float32)
        buf42 = empty_strided_cuda((16,), (1,), torch.float32)
        buf59 = empty_strided_cuda((16,), (1,), torch.float32)
        buf63 = empty_strided_cuda((16,), (1,), torch.float32)
        get_raw_stream(0)
        triton_poi_fused_stack_0[grid(16)](primals_1, buf0, buf17, buf21,
            buf38, buf42, buf59, buf63, 16, XBLOCK=16, num_warps=1,
            num_stages=1)
        buf1 = empty_strided_cuda((1, 4), (4, 1), torch.float32)
        extern_kernels.mm(reinterpret_tensor(primals_1, (1, 4), (1, 1), 0),
            reinterpret_tensor(buf0, (4, 4), (1, 4), 0), out=buf1)
        buf2 = empty_strided_cuda((), (), torch.float32)
        buf7 = empty_strided_cuda((16,), (1,), torch.float32)
        buf3 = reinterpret_tensor(buf7, (4,), (1,), 0)
        triton_per_fused_cat_linalg_vector_norm_mean_1[grid(1)](primals_1,
            buf2, buf3, 1, 4, XBLOCK=1, num_warps=2, num_stages=1)
        buf23 = empty_strided_cuda((), (), torch.float32)
        buf4 = reinterpret_tensor(buf7, (4,), (1,), 4)
        buf14 = empty_strided_cuda((16,), (1,), torch.float32)
        buf11 = reinterpret_tensor(buf14, (4,), (1,), 4)
        buf28 = empty_strided_cuda((16,), (1,), torch.float32)
        buf24 = reinterpret_tensor(buf28, (4,), (1,), 4)
        triton_per_fused_cat_linalg_vector_norm_mean_stack_2[grid(1)](primals_1
            , buf23, buf4, buf11, buf24, 1, 4, XBLOCK=1, num_warps=2,
            num_stages=1)
        buf44 = empty_strided_cuda((), (), torch.float32)
        buf5 = reinterpret_tensor(buf7, (4,), (1,), 8)
        buf12 = reinterpret_tensor(buf14, (4,), (1,), 8)
        buf26 = reinterpret_tensor(buf28, (4,), (1,), 8)
        buf35 = empty_strided_cuda((16,), (1,), torch.float32)
        buf33 = reinterpret_tensor(buf35, (4,), (1,), 8)
        buf49 = empty_strided_cuda((16,), (1,), torch.float32)
        buf45 = reinterpret_tensor(buf49, (4,), (1,), 8)
        buf51 = empty_strided_cuda((), (), torch.float32)
        triton_per_fused_cat_linalg_vector_norm_mean_stack_3[grid(1)](primals_1
            , buf44, buf5, buf12, buf26, buf33, buf45, buf51, 1, 4, XBLOCK=
            1, num_warps=2, num_stages=1)
        buf65 = empty_strided_cuda((), (), torch.float32)
        buf6 = reinterpret_tensor(buf7, (4,), (1,), 12)
        buf13 = reinterpret_tensor(buf14, (4,), (1,), 12)
        buf27 = reinterpret_tensor(buf28, (4,), (1,), 12)
        buf34 = reinterpret_tensor(buf35, (4,), (1,), 12)
        buf48 = reinterpret_tensor(buf49, (4,), (1,), 12)
        buf56 = empty_strided_cuda((16,), (1,), torch.float32)
        buf55 = reinterpret_tensor(buf56, (4,), (1,), 12)
        buf70 = empty_strided_cuda((16,), (1,), torch.float32)
        buf66 = reinterpret_tensor(buf70, (4,), (1,), 12)
        triton_per_fused_cat_linalg_vector_norm_mean_stack_4[grid(1)](primals_1
            , buf65, buf6, buf13, buf27, buf34, buf48, buf55, buf66, 1, 4,
            XBLOCK=1, num_warps=2, num_stages=1)
        del buf3
        del buf4
        del buf5
        del buf6
        buf8 = empty_strided_cuda((1, 4), (4, 1), torch.float32)
        extern_kernels.mm(reinterpret_tensor(primals_1, (1, 4), (1, 1), 4),
            reinterpret_tensor(buf7, (4, 4), (1, 4), 0), out=buf8)
        buf9 = empty_strided_cuda((), (), torch.float32)
        triton_per_fused_linalg_vector_norm_5[grid(1)](primals_1, buf9, 1, 
            4, XBLOCK=1, num_warps=2, num_stages=1)
        buf16 = empty_strided_cuda((), (), torch.float32)
        buf19 = empty_strided_cuda((), (), torch.float32)
        buf10 = reinterpret_tensor(buf14, (4,), (1,), 0)
        buf25 = reinterpret_tensor(buf28, (4,), (1,), 0)
        buf32 = reinterpret_tensor(buf35, (4,), (1,), 0)
        buf46 = reinterpret_tensor(buf49, (4,), (1,), 0)
        buf53 = reinterpret_tensor(buf56, (4,), (1,), 0)
        buf67 = reinterpret_tensor(buf70, (4,), (1,), 0)
        buf77 = empty_strided_cuda((16,), (1,), torch.float32)
        buf74 = reinterpret_tensor(buf77, (4,), (1,), 0)
        triton_per_fused_cat_linalg_vector_norm_mean_stack_6[grid(1)](primals_1
            , buf16, buf19, buf10, buf25, buf32, buf46, buf53, buf67, buf74,
            1, 4, XBLOCK=1, num_warps=2, num_stages=1)
        del buf10
        del buf11
        del buf12
        del buf13
        buf15 = empty_strided_cuda((1, 4), (4, 1), torch.float32)
        extern_kernels.mm(reinterpret_tensor(primals_1, (1, 4), (1, 1), 8),
            reinterpret_tensor(buf14, (4, 4), (1, 4), 0), out=buf15)
        buf18 = empty_strided_cuda((1, 4), (4, 1), torch.float32)
        extern_kernels.mm(reinterpret_tensor(primals_1, (1, 4), (1, 1), 12),
            reinterpret_tensor(buf17, (4, 4), (1, 4), 0), out=buf18)
        buf84 = empty_strided_cuda((16, 4), (4, 1), torch.float32)
        buf20 = reinterpret_tensor(buf84, (4, 4), (4, 1), 0)
        triton_poi_fused_cat_7[grid(16)](buf1, buf2, buf0, buf8, buf9, buf7,
            buf15, buf16, buf14, buf18, buf19, buf17, buf20, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        del buf0
        del buf14
        del buf16
        del buf17
        del buf7
        buf22 = buf8
        del buf8
        extern_kernels.mm(reinterpret_tensor(primals_1, (1, 4), (1, 1), 16),
            reinterpret_tensor(buf21, (4, 4), (1, 4), 0), out=buf22)
        del buf24
        del buf25
        del buf26
        del buf27
        buf29 = buf18
        del buf18
        extern_kernels.mm(reinterpret_tensor(primals_1, (1, 4), (1, 1), 20),
            reinterpret_tensor(buf28, (4, 4), (1, 4), 0), out=buf29)
        buf30 = buf9
        del buf9
        buf37 = buf2
        del buf2
        buf40 = buf19
        del buf19
        buf31 = reinterpret_tensor(buf35, (4,), (1,), 4)
        buf47 = reinterpret_tensor(buf49, (4,), (1,), 4)
        buf54 = reinterpret_tensor(buf56, (4,), (1,), 4)
        buf68 = reinterpret_tensor(buf70, (4,), (1,), 4)
        buf75 = reinterpret_tensor(buf77, (4,), (1,), 4)
        triton_per_fused_cat_linalg_vector_norm_mean_stack_8[grid(1)](primals_1
            , buf30, buf37, buf40, buf31, buf47, buf54, buf68, buf75, 1, 4,
            XBLOCK=1, num_warps=2, num_stages=1)
        del buf31
        del buf32
        del buf33
        del buf34
        buf36 = buf15
        del buf15
        extern_kernels.mm(reinterpret_tensor(primals_1, (1, 4), (1, 1), 24),
            reinterpret_tensor(buf35, (4, 4), (1, 4), 0), out=buf36)
        buf39 = buf1
        del buf1
        extern_kernels.mm(reinterpret_tensor(primals_1, (1, 4), (1, 1), 28),
            reinterpret_tensor(buf38, (4, 4), (1, 4), 0), out=buf39)
        buf41 = reinterpret_tensor(buf84, (4, 4), (4, 1), 16)
        triton_poi_fused_cat_7[grid(16)](buf22, buf23, buf21, buf29, buf30,
            buf28, buf36, buf37, buf35, buf39, buf40, buf38, buf41, 16,
            XBLOCK=16, num_warps=1, num_stages=1)
        del buf21
        del buf23
        del buf28
        del buf30
        del buf35
        del buf38
        buf43 = buf39
        del buf39
        extern_kernels.mm(reinterpret_tensor(primals_1, (1, 4), (1, 1), 32),
            reinterpret_tensor(buf42, (4, 4), (1, 4), 0), out=buf43)
        del buf45
        del buf46
        del buf47
        del buf48
        buf50 = buf36
        del buf36
        extern_kernels.mm(reinterpret_tensor(primals_1, (1, 4), (1, 1), 36),
            reinterpret_tensor(buf49, (4, 4), (1, 4), 0), out=buf50)
        buf58 = buf40
        del buf40
        buf61 = buf37
        del buf37
        buf52 = reinterpret_tensor(buf56, (4,), (1,), 8)
        buf69 = reinterpret_tensor(buf70, (4,), (1,), 8)
        buf76 = reinterpret_tensor(buf77, (4,), (1,), 8)
        triton_per_fused_cat_linalg_vector_norm_mean_stack_9[grid(1)](primals_1
            , buf58, buf61, buf52, buf69, buf76, 1, 4, XBLOCK=1, num_warps=
            2, num_stages=1)
        del buf52
        del buf53
        del buf54
        del buf55
        buf57 = buf29
        del buf29
        extern_kernels.mm(reinterpret_tensor(primals_1, (1, 4), (1, 1), 40),
            reinterpret_tensor(buf56, (4, 4), (1, 4), 0), out=buf57)
        buf60 = buf22
        del buf22
        extern_kernels.mm(reinterpret_tensor(primals_1, (1, 4), (1, 1), 44),
            reinterpret_tensor(buf59, (4, 4), (1, 4), 0), out=buf60)
        buf62 = reinterpret_tensor(buf84, (4, 4), (4, 1), 32)
        triton_poi_fused_cat_7[grid(16)](buf43, buf44, buf42, buf50, buf51,
            buf49, buf57, buf58, buf56, buf60, buf61, buf59, buf62, 16,
            XBLOCK=16, num_warps=1, num_stages=1)
        del buf42
        del buf44
        del buf49
        del buf56
        buf64 = buf60
        del buf60
        extern_kernels.mm(reinterpret_tensor(primals_1, (1, 4), (1, 1), 48),
            reinterpret_tensor(buf63, (4, 4), (1, 4), 0), out=buf64)
        del buf66
        del buf67
        del buf68
        del buf69
        buf71 = buf57
        del buf57
        extern_kernels.mm(reinterpret_tensor(primals_1, (1, 4), (1, 1), 52),
            reinterpret_tensor(buf70, (4, 4), (1, 4), 0), out=buf71)
        buf72 = buf61
        del buf61
        triton_per_fused_linalg_vector_norm_10[grid(1)](primals_1, buf72, 1,
            4, XBLOCK=1, num_warps=2, num_stages=1)
        buf73 = reinterpret_tensor(buf77, (4,), (1,), 12)
        buf82 = buf58
        del buf58
        triton_per_fused_cat_linalg_vector_norm_mean_11[grid(1)](primals_1,
            buf73, buf82, 1, 4, XBLOCK=1, num_warps=2, num_stages=1)
        del buf73
        del buf74
        del buf75
        del buf76
        buf78 = buf50
        del buf50
        extern_kernels.mm(reinterpret_tensor(primals_1, (1, 4), (1, 1), 56),
            reinterpret_tensor(buf77, (4, 4), (1, 4), 0), out=buf78)
        buf79 = buf51
        del buf51
        triton_per_fused_linalg_vector_norm_12[grid(1)](primals_1, buf79, 1,
            4, XBLOCK=1, num_warps=2, num_stages=1)
        buf80 = buf59
        del buf59
        triton_poi_fused_stack_13[grid(16)](primals_1, buf80, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        buf81 = buf43
        del buf43
        extern_kernels.mm(reinterpret_tensor(primals_1, (1, 4), (1, 1), 60),
            reinterpret_tensor(buf80, (4, 4), (1, 4), 0), out=buf81)
        del primals_1
        buf83 = reinterpret_tensor(buf84, (4, 4), (4, 1), 48)
        triton_poi_fused_cat_7[grid(16)](buf64, buf65, buf63, buf71, buf72,
            buf70, buf78, buf79, buf77, buf81, buf82, buf80, buf83, 16,
            XBLOCK=16, num_warps=1, num_stages=1)
        del buf63
        del buf65
        del buf70
        del buf72
        del buf77
        del buf79
        buf125 = buf80
        del buf80
        buf93 = reinterpret_tensor(buf125, (4,), (1,), 0)
        buf102 = reinterpret_tensor(buf125, (4,), (1,), 4)
        buf111 = reinterpret_tensor(buf125, (4,), (1,), 8)
        buf114 = reinterpret_tensor(buf81, (4,), (1,), 0)
        del buf81
        buf130 = reinterpret_tensor(buf78, (4,), (1,), 0)
        del buf78
        buf117 = reinterpret_tensor(buf71, (4,), (1,), 0)
        del buf71
        buf129 = reinterpret_tensor(buf64, (4,), (1,), 0)
        del buf64
        buf120 = empty_strided_cuda((4,), (1,), torch.float32)
        buf128 = empty_strided_cuda((4,), (1,), torch.float32)
        buf123 = empty_strided_cuda((4,), (1,), torch.float32)
        buf127 = empty_strided_cuda((4,), (1,), torch.float32)
        triton_per_fused__log_softmax__log_softmax_backward_data_stack_14[grid
            (1)](primals_2, buf84, primals_3, buf93, buf102, buf111, buf114,
            buf130, buf117, buf129, buf120, buf128, buf123, buf127, 1, 4,
            XBLOCK=1, num_warps=2, num_stages=1)
        buf124 = reinterpret_tensor(buf125, (4,), (1,), 12)
        triton_poi_fused_stack_15[grid(4)](buf114, buf117, buf120, buf123,
            buf124, 4, XBLOCK=4, num_warps=1, num_stages=1)
        del buf114
        del buf117
        del buf120
        del buf123
        buf126 = buf82
        del buf82
        buf131 = buf126
        del buf126
        triton_per_fused_mean_16[grid(1)](buf131, buf125, 1, 16, XBLOCK=1,
            num_warps=2, num_stages=1)
        del buf102
        del buf111
        del buf124
        del buf125
        del buf93
    return buf131, primals_2, primals_3, buf84, buf127, buf128, buf129, buf130


class GE2ELossNew(nn.Module):

    def __init__(self, init_w=10.0, init_b=-5.0, loss_method='softmax'):
        """
        Implementation of the Generalized End-to-End loss defined in https://arxiv.org/abs/1710.10467 [1]
        Accepts an input of size (N, M, D)
            where N is the number of speakers in the batch,
            M is the number of utterances per speaker,
            and D is the dimensionality of the embedding vector (e.g. d-vector)
        Args:
            - init_w (float): defines the initial value of w in Equation (5) of [1]
            - init_b (float): definies the initial value of b in Equation (5) of [1]
        """
        super(GE2ELossNew, self).__init__()
        self.w = nn.Parameter(torch.tensor(init_w))
        self.b = nn.Parameter(torch.tensor(init_b))
        self.loss_method = loss_method
        assert self.loss_method in ['softmax', 'contrast']
        if self.loss_method == 'softmax':
            self.embed_loss = self.embed_loss_softmax
        if self.loss_method == 'contrast':
            self.embed_loss = self.embed_loss_contrast

    def calc_new_centroids(self, dvecs, centroids, spkr, utt):
        """
        Calculates the new centroids excluding the reference utterance
        """
        excl = torch.cat((dvecs[spkr, :utt], dvecs[spkr, utt + 1:]))
        excl = torch.mean(excl, 0)
        new_centroids = []
        for i, centroid in enumerate(centroids):
            if i == spkr:
                new_centroids.append(excl)
            else:
                new_centroids.append(centroid)
        return torch.stack(new_centroids)

    def calc_cosine_sim(self, dvecs, centroids):
        """
        Make the cosine similarity matrix with dims (N,M,N)
        """
        cos_sim_matrix = []
        for spkr_idx, speaker in enumerate(dvecs):
            cs_row = []
            for utt_idx, utterance in enumerate(speaker):
                new_centroids = self.calc_new_centroids(dvecs, centroids,
                    spkr_idx, utt_idx)
                cs_row.append(torch.clamp(torch.mm(utterance.unsqueeze(1).
                    transpose(0, 1), new_centroids.transpose(0, 1)) / (
                    torch.norm(utterance) * torch.norm(new_centroids, dim=1
                    )), 1e-06))
            cs_row = torch.cat(cs_row, dim=0)
            cos_sim_matrix.append(cs_row)
        return torch.stack(cos_sim_matrix)

    def embed_loss_softmax(self, dvecs, cos_sim_matrix):
        """
        Calculates the loss on each embedding $L(e_{ji})$ by taking softmax
        """
        N, M, _ = dvecs.shape
        L = []
        for j in range(N):
            L_row = []
            for i in range(M):
                L_row.append(-F.log_softmax(cos_sim_matrix[j, i], 0)[j])
            L_row = torch.stack(L_row)
            L.append(L_row)
        return torch.stack(L)

    def embed_loss_contrast(self, dvecs, cos_sim_matrix):
        """
        Calculates the loss on each embedding $L(e_{ji})$ by contrast loss with closest centroid
        """
        N, M, _ = dvecs.shape
        L = []
        for j in range(N):
            L_row = []
            for i in range(M):
                centroids_sigmoids = torch.sigmoid(cos_sim_matrix[j, i])
                excl_centroids_sigmoids = torch.cat((centroids_sigmoids[:j],
                    centroids_sigmoids[j + 1:]))
                L_row.append(1.0 - torch.sigmoid(cos_sim_matrix[j, i, j]) +
                    torch.max(excl_centroids_sigmoids))
            L_row = torch.stack(L_row)
            L.append(L_row)
        return torch.stack(L)

    def forward(self, input_0):
        primals_2 = self.w
        primals_3 = self.b
        primals_1 = input_0
        output = call([primals_1, primals_2, primals_3])
        return output[0]


Measured latency_ms: 0.066956