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 numpy as np
from torch import nn
from sklearn.neighbors import NearestNeighbors
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_red_fused_linalg_vector_norm_0(in_ptr0, out_ptr0, xnumel, rnumel,
    XBLOCK: tl.constexpr, RBLOCK: tl.constexpr):
    rnumel = 128
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:, None]
    tl.full([XBLOCK, RBLOCK], True, tl.int1)
    rbase = tl.arange(0, RBLOCK)[None, :]
    x0 = xindex % 4096
    x1 = xindex // 4096
    _tmp3 = tl.full([XBLOCK, RBLOCK], 0, tl.float32)
    x3 = xindex
    for roffset in range(0, rnumel, RBLOCK):
        rindex = roffset + rbase
        rmask = rindex < rnumel
        r2 = rindex
        tmp0 = tl.load(in_ptr0 + (x0 + 4096 * r2 + 524288 * x1), rmask,
            eviction_policy='evict_last', other=0.0)
        tmp1 = tmp0 * tmp0
        tmp2 = tl.broadcast_to(tmp1, [XBLOCK, RBLOCK])
        tmp4 = _tmp3 + tmp2
        _tmp3 = tl.where(rmask, tmp4, _tmp3)
    tmp3 = tl.sum(_tmp3, 1)[:, None]
    tl.store(out_ptr0 + x3, tmp3, None)


@triton.jit
def triton_poi_fused_div_sub_1(in_ptr0, in_ptr1, in_ptr2, out_ptr0,
    out_ptr1, out_ptr2, out_ptr3, out_ptr4, out_ptr5, out_ptr6, out_ptr7,
    out_ptr8, out_ptr9, out_ptr10, out_ptr11, out_ptr12, out_ptr13,
    out_ptr14, out_ptr15, out_ptr16, out_ptr17, out_ptr18, out_ptr19,
    out_ptr20, out_ptr21, out_ptr22, out_ptr23, out_ptr24, out_ptr25,
    out_ptr26, out_ptr27, out_ptr28, out_ptr29, out_ptr30, out_ptr31,
    out_ptr32, out_ptr33, out_ptr34, out_ptr35, out_ptr36, out_ptr37,
    out_ptr38, out_ptr39, out_ptr40, out_ptr41, out_ptr42, out_ptr43,
    out_ptr44, out_ptr45, out_ptr46, out_ptr47, out_ptr48, out_ptr49,
    out_ptr50, out_ptr51, out_ptr52, out_ptr53, out_ptr54, out_ptr55,
    out_ptr56, out_ptr57, out_ptr58, out_ptr59, out_ptr60, out_ptr61,
    out_ptr62, out_ptr63, xnumel, XBLOCK: tl.constexpr):
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    tl.full([XBLOCK], True, tl.int1)
    x3 = xindex
    x0 = xindex % 4096
    x2 = xindex // 524288
    x1 = xindex // 4096 % 128
    tmp0 = tl.load(in_ptr0 + x3, None)
    tmp1 = tl.load(in_ptr1 + (x0 + 4096 * x2), None, eviction_policy=
        'evict_last')
    tmp6 = tl.load(in_ptr2 + (128 + x1), None, eviction_policy='evict_last')
    tmp8 = tl.load(in_ptr2 + (256 + x1), None, eviction_policy='evict_last')
    tmp10 = tl.load(in_ptr2 + (384 + x1), None, eviction_policy='evict_last')
    tmp12 = tl.load(in_ptr2 + (512 + x1), None, eviction_policy='evict_last')
    tmp14 = tl.load(in_ptr2 + (640 + x1), None, eviction_policy='evict_last')
    tmp16 = tl.load(in_ptr2 + (768 + x1), None, eviction_policy='evict_last')
    tmp18 = tl.load(in_ptr2 + (896 + x1), None, eviction_policy='evict_last')
    tmp20 = tl.load(in_ptr2 + (1024 + x1), None, eviction_policy='evict_last')
    tmp22 = tl.load(in_ptr2 + (1152 + x1), None, eviction_policy='evict_last')
    tmp24 = tl.load(in_ptr2 + (1280 + x1), None, eviction_policy='evict_last')
    tmp26 = tl.load(in_ptr2 + (1408 + x1), None, eviction_policy='evict_last')
    tmp28 = tl.load(in_ptr2 + (1536 + x1), None, eviction_policy='evict_last')
    tmp30 = tl.load(in_ptr2 + (1664 + x1), None, eviction_policy='evict_last')
    tmp32 = tl.load(in_ptr2 + (1792 + x1), None, eviction_policy='evict_last')
    tmp34 = tl.load(in_ptr2 + (1920 + x1), None, eviction_policy='evict_last')
    tmp36 = tl.load(in_ptr2 + (2048 + x1), None, eviction_policy='evict_last')
    tmp38 = tl.load(in_ptr2 + (2176 + x1), None, eviction_policy='evict_last')
    tmp40 = tl.load(in_ptr2 + (2304 + x1), None, eviction_policy='evict_last')
    tmp42 = tl.load(in_ptr2 + (2432 + x1), None, eviction_policy='evict_last')
    tmp44 = tl.load(in_ptr2 + (2560 + x1), None, eviction_policy='evict_last')
    tmp46 = tl.load(in_ptr2 + (2688 + x1), None, eviction_policy='evict_last')
    tmp48 = tl.load(in_ptr2 + (2816 + x1), None, eviction_policy='evict_last')
    tmp50 = tl.load(in_ptr2 + (2944 + x1), None, eviction_policy='evict_last')
    tmp52 = tl.load(in_ptr2 + (3072 + x1), None, eviction_policy='evict_last')
    tmp54 = tl.load(in_ptr2 + (3200 + x1), None, eviction_policy='evict_last')
    tmp56 = tl.load(in_ptr2 + (3328 + x1), None, eviction_policy='evict_last')
    tmp58 = tl.load(in_ptr2 + (3456 + x1), None, eviction_policy='evict_last')
    tmp60 = tl.load(in_ptr2 + (3584 + x1), None, eviction_policy='evict_last')
    tmp62 = tl.load(in_ptr2 + (3712 + x1), None, eviction_policy='evict_last')
    tmp64 = tl.load(in_ptr2 + (3840 + x1), None, eviction_policy='evict_last')
    tmp66 = tl.load(in_ptr2 + (3968 + x1), None, eviction_policy='evict_last')
    tmp68 = tl.load(in_ptr2 + (4096 + x1), None, eviction_policy='evict_last')
    tmp70 = tl.load(in_ptr2 + (4224 + x1), None, eviction_policy='evict_last')
    tmp72 = tl.load(in_ptr2 + (4352 + x1), None, eviction_policy='evict_last')
    tmp74 = tl.load(in_ptr2 + (4480 + x1), None, eviction_policy='evict_last')
    tmp76 = tl.load(in_ptr2 + (4608 + x1), None, eviction_policy='evict_last')
    tmp78 = tl.load(in_ptr2 + (4736 + x1), None, eviction_policy='evict_last')
    tmp80 = tl.load(in_ptr2 + (4864 + x1), None, eviction_policy='evict_last')
    tmp82 = tl.load(in_ptr2 + (4992 + x1), None, eviction_policy='evict_last')
    tmp84 = tl.load(in_ptr2 + (5120 + x1), None, eviction_policy='evict_last')
    tmp86 = tl.load(in_ptr2 + (5248 + x1), None, eviction_policy='evict_last')
    tmp88 = tl.load(in_ptr2 + (5376 + x1), None, eviction_policy='evict_last')
    tmp90 = tl.load(in_ptr2 + (5504 + x1), None, eviction_policy='evict_last')
    tmp92 = tl.load(in_ptr2 + (5632 + x1), None, eviction_policy='evict_last')
    tmp94 = tl.load(in_ptr2 + (5760 + x1), None, eviction_policy='evict_last')
    tmp96 = tl.load(in_ptr2 + (5888 + x1), None, eviction_policy='evict_last')
    tmp98 = tl.load(in_ptr2 + (6016 + x1), None, eviction_policy='evict_last')
    tmp100 = tl.load(in_ptr2 + (6144 + x1), None, eviction_policy='evict_last')
    tmp102 = tl.load(in_ptr2 + (6272 + x1), None, eviction_policy='evict_last')
    tmp104 = tl.load(in_ptr2 + (6400 + x1), None, eviction_policy='evict_last')
    tmp106 = tl.load(in_ptr2 + (6528 + x1), None, eviction_policy='evict_last')
    tmp108 = tl.load(in_ptr2 + (6656 + x1), None, eviction_policy='evict_last')
    tmp110 = tl.load(in_ptr2 + (6784 + x1), None, eviction_policy='evict_last')
    tmp112 = tl.load(in_ptr2 + (6912 + x1), None, eviction_policy='evict_last')
    tmp114 = tl.load(in_ptr2 + (7040 + x1), None, eviction_policy='evict_last')
    tmp116 = tl.load(in_ptr2 + (7168 + x1), None, eviction_policy='evict_last')
    tmp118 = tl.load(in_ptr2 + (7296 + x1), None, eviction_policy='evict_last')
    tmp120 = tl.load(in_ptr2 + (7424 + x1), None, eviction_policy='evict_last')
    tmp122 = tl.load(in_ptr2 + (7552 + x1), None, eviction_policy='evict_last')
    tmp124 = tl.load(in_ptr2 + (7680 + x1), None, eviction_policy='evict_last')
    tmp126 = tl.load(in_ptr2 + (7808 + x1), None, eviction_policy='evict_last')
    tmp128 = tl.load(in_ptr2 + (7936 + x1), None, eviction_policy='evict_last')
    tmp130 = tl.load(in_ptr2 + (8064 + x1), None, eviction_policy='evict_last')
    tmp2 = libdevice.sqrt(tmp1)
    tmp3 = 1e-12
    tmp4 = triton_helpers.maximum(tmp2, tmp3)
    tmp5 = tmp0 / tmp4
    tmp7 = tmp5 - tmp6
    tmp9 = tmp5 - tmp8
    tmp11 = tmp5 - tmp10
    tmp13 = tmp5 - tmp12
    tmp15 = tmp5 - tmp14
    tmp17 = tmp5 - tmp16
    tmp19 = tmp5 - tmp18
    tmp21 = tmp5 - tmp20
    tmp23 = tmp5 - tmp22
    tmp25 = tmp5 - tmp24
    tmp27 = tmp5 - tmp26
    tmp29 = tmp5 - tmp28
    tmp31 = tmp5 - tmp30
    tmp33 = tmp5 - tmp32
    tmp35 = tmp5 - tmp34
    tmp37 = tmp5 - tmp36
    tmp39 = tmp5 - tmp38
    tmp41 = tmp5 - tmp40
    tmp43 = tmp5 - tmp42
    tmp45 = tmp5 - tmp44
    tmp47 = tmp5 - tmp46
    tmp49 = tmp5 - tmp48
    tmp51 = tmp5 - tmp50
    tmp53 = tmp5 - tmp52
    tmp55 = tmp5 - tmp54
    tmp57 = tmp5 - tmp56
    tmp59 = tmp5 - tmp58
    tmp61 = tmp5 - tmp60
    tmp63 = tmp5 - tmp62
    tmp65 = tmp5 - tmp64
    tmp67 = tmp5 - tmp66
    tmp69 = tmp5 - tmp68
    tmp71 = tmp5 - tmp70
    tmp73 = tmp5 - tmp72
    tmp75 = tmp5 - tmp74
    tmp77 = tmp5 - tmp76
    tmp79 = tmp5 - tmp78
    tmp81 = tmp5 - tmp80
    tmp83 = tmp5 - tmp82
    tmp85 = tmp5 - tmp84
    tmp87 = tmp5 - tmp86
    tmp89 = tmp5 - tmp88
    tmp91 = tmp5 - tmp90
    tmp93 = tmp5 - tmp92
    tmp95 = tmp5 - tmp94
    tmp97 = tmp5 - tmp96
    tmp99 = tmp5 - tmp98
    tmp101 = tmp5 - tmp100
    tmp103 = tmp5 - tmp102
    tmp105 = tmp5 - tmp104
    tmp107 = tmp5 - tmp106
    tmp109 = tmp5 - tmp108
    tmp111 = tmp5 - tmp110
    tmp113 = tmp5 - tmp112
    tmp115 = tmp5 - tmp114
    tmp117 = tmp5 - tmp116
    tmp119 = tmp5 - tmp118
    tmp121 = tmp5 - tmp120
    tmp123 = tmp5 - tmp122
    tmp125 = tmp5 - tmp124
    tmp127 = tmp5 - tmp126
    tmp129 = tmp5 - tmp128
    tmp131 = tmp5 - tmp130
    tl.store(out_ptr0 + x3, tmp5, None)
    tl.store(out_ptr1 + x3, tmp7, None)
    tl.store(out_ptr2 + x3, tmp9, None)
    tl.store(out_ptr3 + x3, tmp11, None)
    tl.store(out_ptr4 + x3, tmp13, None)
    tl.store(out_ptr5 + x3, tmp15, None)
    tl.store(out_ptr6 + x3, tmp17, None)
    tl.store(out_ptr7 + x3, tmp19, None)
    tl.store(out_ptr8 + x3, tmp21, None)
    tl.store(out_ptr9 + x3, tmp23, None)
    tl.store(out_ptr10 + x3, tmp25, None)
    tl.store(out_ptr11 + x3, tmp27, None)
    tl.store(out_ptr12 + x3, tmp29, None)
    tl.store(out_ptr13 + x3, tmp31, None)
    tl.store(out_ptr14 + x3, tmp33, None)
    tl.store(out_ptr15 + x3, tmp35, None)
    tl.store(out_ptr16 + x3, tmp37, None)
    tl.store(out_ptr17 + x3, tmp39, None)
    tl.store(out_ptr18 + x3, tmp41, None)
    tl.store(out_ptr19 + x3, tmp43, None)
    tl.store(out_ptr20 + x3, tmp45, None)
    tl.store(out_ptr21 + x3, tmp47, None)
    tl.store(out_ptr22 + x3, tmp49, None)
    tl.store(out_ptr23 + x3, tmp51, None)
    tl.store(out_ptr24 + x3, tmp53, None)
    tl.store(out_ptr25 + x3, tmp55, None)
    tl.store(out_ptr26 + x3, tmp57, None)
    tl.store(out_ptr27 + x3, tmp59, None)
    tl.store(out_ptr28 + x3, tmp61, None)
    tl.store(out_ptr29 + x3, tmp63, None)
    tl.store(out_ptr30 + x3, tmp65, None)
    tl.store(out_ptr31 + x3, tmp67, None)
    tl.store(out_ptr32 + x3, tmp69, None)
    tl.store(out_ptr33 + x3, tmp71, None)
    tl.store(out_ptr34 + x3, tmp73, None)
    tl.store(out_ptr35 + x3, tmp75, None)
    tl.store(out_ptr36 + x3, tmp77, None)
    tl.store(out_ptr37 + x3, tmp79, None)
    tl.store(out_ptr38 + x3, tmp81, None)
    tl.store(out_ptr39 + x3, tmp83, None)
    tl.store(out_ptr40 + x3, tmp85, None)
    tl.store(out_ptr41 + x3, tmp87, None)
    tl.store(out_ptr42 + x3, tmp89, None)
    tl.store(out_ptr43 + x3, tmp91, None)
    tl.store(out_ptr44 + x3, tmp93, None)
    tl.store(out_ptr45 + x3, tmp95, None)
    tl.store(out_ptr46 + x3, tmp97, None)
    tl.store(out_ptr47 + x3, tmp99, None)
    tl.store(out_ptr48 + x3, tmp101, None)
    tl.store(out_ptr49 + x3, tmp103, None)
    tl.store(out_ptr50 + x3, tmp105, None)
    tl.store(out_ptr51 + x3, tmp107, None)
    tl.store(out_ptr52 + x3, tmp109, None)
    tl.store(out_ptr53 + x3, tmp111, None)
    tl.store(out_ptr54 + x3, tmp113, None)
    tl.store(out_ptr55 + x3, tmp115, None)
    tl.store(out_ptr56 + x3, tmp117, None)
    tl.store(out_ptr57 + x3, tmp119, None)
    tl.store(out_ptr58 + x3, tmp121, None)
    tl.store(out_ptr59 + x3, tmp123, None)
    tl.store(out_ptr60 + x3, tmp125, None)
    tl.store(out_ptr61 + x3, tmp127, None)
    tl.store(out_ptr62 + x3, tmp129, None)
    tl.store(out_ptr63 + x3, tmp131, None)


@triton.jit
def triton_per_fused__softmax_2(in_ptr0, out_ptr0, out_ptr1, xnumel, rnumel,
    XBLOCK: tl.constexpr):
    RBLOCK: tl.constexpr = 64
    xoffset = tl.program_id(0) * XBLOCK
    xindex = 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)
    r2 = rindex
    x0 = xindex % 4096
    x1 = xindex // 4096
    x3 = xindex
    tmp0 = tl.load(in_ptr0 + (x0 + 4096 * r2 + 262144 * x1), None)
    tmp1 = tl.broadcast_to(tmp0, [XBLOCK, RBLOCK])
    tmp3 = triton_helpers.max2(tmp1, 1)[:, None]
    tmp4 = tmp0 - tmp3
    tmp5 = tl_math.exp(tmp4)
    tmp6 = tl.broadcast_to(tmp5, [XBLOCK, RBLOCK])
    tmp8 = tl.sum(tmp6, 1)[:, None]
    tl.store(out_ptr0 + x3, tmp3, None)
    tl.store(out_ptr1 + x3, tmp8, None)


@triton.jit
def triton_red_fused_mul_sub_sum_3(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,
    in_ptr32, out_ptr0, out_ptr1, out_ptr2, out_ptr3, out_ptr4, out_ptr5,
    out_ptr6, out_ptr7, out_ptr8, out_ptr9, out_ptr10, out_ptr11, out_ptr12,
    out_ptr13, out_ptr14, out_ptr15, out_ptr16, out_ptr17, out_ptr18,
    out_ptr19, out_ptr20, out_ptr21, out_ptr22, out_ptr23, out_ptr24,
    out_ptr25, out_ptr26, out_ptr27, out_ptr28, xnumel, rnumel, XBLOCK: tl.
    constexpr, RBLOCK: tl.constexpr):
    xnumel = 512
    rnumel = 4096
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:, None]
    xmask = xindex < xnumel
    rbase = tl.arange(0, RBLOCK)[None, :]
    x3 = xindex
    x0 = xindex % 128
    tmp1 = tl.load(in_ptr1 + x0, xmask, eviction_policy='evict_last')
    x1 = xindex // 128
    _tmp11 = tl.full([XBLOCK, RBLOCK], 0, tl.float32)
    _tmp20 = tl.full([XBLOCK, RBLOCK], 0, tl.float32)
    _tmp29 = tl.full([XBLOCK, RBLOCK], 0, tl.float32)
    _tmp38 = tl.full([XBLOCK, RBLOCK], 0, tl.float32)
    _tmp47 = tl.full([XBLOCK, RBLOCK], 0, tl.float32)
    _tmp56 = tl.full([XBLOCK, RBLOCK], 0, tl.float32)
    _tmp65 = tl.full([XBLOCK, RBLOCK], 0, tl.float32)
    _tmp74 = tl.full([XBLOCK, RBLOCK], 0, tl.float32)
    _tmp83 = tl.full([XBLOCK, RBLOCK], 0, tl.float32)
    _tmp92 = tl.full([XBLOCK, RBLOCK], 0, tl.float32)
    _tmp101 = tl.full([XBLOCK, RBLOCK], 0, tl.float32)
    _tmp110 = tl.full([XBLOCK, RBLOCK], 0, tl.float32)
    _tmp119 = tl.full([XBLOCK, RBLOCK], 0, tl.float32)
    _tmp128 = tl.full([XBLOCK, RBLOCK], 0, tl.float32)
    _tmp137 = tl.full([XBLOCK, RBLOCK], 0, tl.float32)
    _tmp146 = tl.full([XBLOCK, RBLOCK], 0, tl.float32)
    _tmp155 = tl.full([XBLOCK, RBLOCK], 0, tl.float32)
    _tmp164 = tl.full([XBLOCK, RBLOCK], 0, tl.float32)
    _tmp173 = tl.full([XBLOCK, RBLOCK], 0, tl.float32)
    _tmp182 = tl.full([XBLOCK, RBLOCK], 0, tl.float32)
    _tmp191 = tl.full([XBLOCK, RBLOCK], 0, tl.float32)
    _tmp200 = tl.full([XBLOCK, RBLOCK], 0, tl.float32)
    _tmp209 = tl.full([XBLOCK, RBLOCK], 0, tl.float32)
    _tmp218 = tl.full([XBLOCK, RBLOCK], 0, tl.float32)
    _tmp227 = tl.full([XBLOCK, RBLOCK], 0, tl.float32)
    _tmp236 = tl.full([XBLOCK, RBLOCK], 0, tl.float32)
    _tmp245 = tl.full([XBLOCK, RBLOCK], 0, tl.float32)
    _tmp254 = tl.full([XBLOCK, RBLOCK], 0, tl.float32)
    _tmp263 = tl.full([XBLOCK, RBLOCK], 0, tl.float32)
    for roffset in range(0, rnumel, RBLOCK):
        rindex = roffset + rbase
        rmask = rindex < rnumel
        r2 = rindex
        tmp0 = tl.load(in_ptr0 + (r2 + 4096 * x3), rmask & xmask,
            eviction_policy='evict_first', other=0.0)
        tmp3 = tl.load(in_ptr2 + (r2 + 262144 * x1), rmask & xmask,
            eviction_policy='evict_last', other=0.0)
        tmp4 = tl.load(in_ptr3 + (r2 + 4096 * x1), rmask & xmask,
            eviction_policy='evict_last', other=0.0)
        tmp7 = tl.load(in_ptr4 + (r2 + 4096 * x1), rmask & xmask,
            eviction_policy='evict_last', other=0.0)
        tmp13 = tl.load(in_ptr5 + (r2 + 4096 * x3), rmask & xmask,
            eviction_policy='evict_first', other=0.0)
        tmp14 = tl.load(in_ptr2 + (4096 + r2 + 262144 * x1), rmask & xmask,
            eviction_policy='evict_last', other=0.0)
        tmp22 = tl.load(in_ptr6 + (r2 + 4096 * x3), rmask & xmask,
            eviction_policy='evict_first', other=0.0)
        tmp23 = tl.load(in_ptr2 + (8192 + r2 + 262144 * x1), rmask & xmask,
            eviction_policy='evict_last', other=0.0)
        tmp31 = tl.load(in_ptr7 + (r2 + 4096 * x3), rmask & xmask,
            eviction_policy='evict_first', other=0.0)
        tmp32 = tl.load(in_ptr2 + (12288 + r2 + 262144 * x1), rmask & xmask,
            eviction_policy='evict_last', other=0.0)
        tmp40 = tl.load(in_ptr8 + (r2 + 4096 * x3), rmask & xmask,
            eviction_policy='evict_first', other=0.0)
        tmp41 = tl.load(in_ptr2 + (16384 + r2 + 262144 * x1), rmask & xmask,
            eviction_policy='evict_last', other=0.0)
        tmp49 = tl.load(in_ptr9 + (r2 + 4096 * x3), rmask & xmask,
            eviction_policy='evict_first', other=0.0)
        tmp50 = tl.load(in_ptr2 + (20480 + r2 + 262144 * x1), rmask & xmask,
            eviction_policy='evict_last', other=0.0)
        tmp58 = tl.load(in_ptr10 + (r2 + 4096 * x3), rmask & xmask,
            eviction_policy='evict_first', other=0.0)
        tmp59 = tl.load(in_ptr2 + (24576 + r2 + 262144 * x1), rmask & xmask,
            eviction_policy='evict_last', other=0.0)
        tmp67 = tl.load(in_ptr11 + (r2 + 4096 * x3), rmask & xmask,
            eviction_policy='evict_first', other=0.0)
        tmp68 = tl.load(in_ptr2 + (28672 + r2 + 262144 * x1), rmask & xmask,
            eviction_policy='evict_last', other=0.0)
        tmp76 = tl.load(in_ptr12 + (r2 + 4096 * x3), rmask & xmask,
            eviction_policy='evict_first', other=0.0)
        tmp77 = tl.load(in_ptr2 + (32768 + r2 + 262144 * x1), rmask & xmask,
            eviction_policy='evict_last', other=0.0)
        tmp85 = tl.load(in_ptr13 + (r2 + 4096 * x3), rmask & xmask,
            eviction_policy='evict_first', other=0.0)
        tmp86 = tl.load(in_ptr2 + (36864 + r2 + 262144 * x1), rmask & xmask,
            eviction_policy='evict_last', other=0.0)
        tmp94 = tl.load(in_ptr14 + (r2 + 4096 * x3), rmask & xmask,
            eviction_policy='evict_first', other=0.0)
        tmp95 = tl.load(in_ptr2 + (40960 + r2 + 262144 * x1), rmask & xmask,
            eviction_policy='evict_last', other=0.0)
        tmp103 = tl.load(in_ptr15 + (r2 + 4096 * x3), rmask & xmask,
            eviction_policy='evict_first', other=0.0)
        tmp104 = tl.load(in_ptr2 + (45056 + r2 + 262144 * x1), rmask &
            xmask, eviction_policy='evict_last', other=0.0)
        tmp112 = tl.load(in_ptr16 + (r2 + 4096 * x3), rmask & xmask,
            eviction_policy='evict_first', other=0.0)
        tmp113 = tl.load(in_ptr2 + (49152 + r2 + 262144 * x1), rmask &
            xmask, eviction_policy='evict_last', other=0.0)
        tmp121 = tl.load(in_ptr17 + (r2 + 4096 * x3), rmask & xmask,
            eviction_policy='evict_first', other=0.0)
        tmp122 = tl.load(in_ptr2 + (53248 + r2 + 262144 * x1), rmask &
            xmask, eviction_policy='evict_last', other=0.0)
        tmp130 = tl.load(in_ptr18 + (r2 + 4096 * x3), rmask & xmask,
            eviction_policy='evict_first', other=0.0)
        tmp131 = tl.load(in_ptr2 + (57344 + r2 + 262144 * x1), rmask &
            xmask, eviction_policy='evict_last', other=0.0)
        tmp139 = tl.load(in_ptr19 + (r2 + 4096 * x3), rmask & xmask,
            eviction_policy='evict_first', other=0.0)
        tmp140 = tl.load(in_ptr2 + (61440 + r2 + 262144 * x1), rmask &
            xmask, eviction_policy='evict_last', other=0.0)
        tmp148 = tl.load(in_ptr20 + (r2 + 4096 * x3), rmask & xmask,
            eviction_policy='evict_first', other=0.0)
        tmp149 = tl.load(in_ptr2 + (65536 + r2 + 262144 * x1), rmask &
            xmask, eviction_policy='evict_last', other=0.0)
        tmp157 = tl.load(in_ptr21 + (r2 + 4096 * x3), rmask & xmask,
            eviction_policy='evict_first', other=0.0)
        tmp158 = tl.load(in_ptr2 + (69632 + r2 + 262144 * x1), rmask &
            xmask, eviction_policy='evict_last', other=0.0)
        tmp166 = tl.load(in_ptr22 + (r2 + 4096 * x3), rmask & xmask,
            eviction_policy='evict_first', other=0.0)
        tmp167 = tl.load(in_ptr2 + (73728 + r2 + 262144 * x1), rmask &
            xmask, eviction_policy='evict_last', other=0.0)
        tmp175 = tl.load(in_ptr23 + (r2 + 4096 * x3), rmask & xmask,
            eviction_policy='evict_first', other=0.0)
        tmp176 = tl.load(in_ptr2 + (77824 + r2 + 262144 * x1), rmask &
            xmask, eviction_policy='evict_last', other=0.0)
        tmp184 = tl.load(in_ptr24 + (r2 + 4096 * x3), rmask & xmask,
            eviction_policy='evict_first', other=0.0)
        tmp185 = tl.load(in_ptr2 + (81920 + r2 + 262144 * x1), rmask &
            xmask, eviction_policy='evict_last', other=0.0)
        tmp193 = tl.load(in_ptr25 + (r2 + 4096 * x3), rmask & xmask,
            eviction_policy='evict_first', other=0.0)
        tmp194 = tl.load(in_ptr2 + (86016 + r2 + 262144 * x1), rmask &
            xmask, eviction_policy='evict_last', other=0.0)
        tmp202 = tl.load(in_ptr26 + (r2 + 4096 * x3), rmask & xmask,
            eviction_policy='evict_first', other=0.0)
        tmp203 = tl.load(in_ptr2 + (90112 + r2 + 262144 * x1), rmask &
            xmask, eviction_policy='evict_last', other=0.0)
        tmp211 = tl.load(in_ptr27 + (r2 + 4096 * x3), rmask & xmask,
            eviction_policy='evict_first', other=0.0)
        tmp212 = tl.load(in_ptr2 + (94208 + r2 + 262144 * x1), rmask &
            xmask, eviction_policy='evict_last', other=0.0)
        tmp220 = tl.load(in_ptr28 + (r2 + 4096 * x3), rmask & xmask,
            eviction_policy='evict_first', other=0.0)
        tmp221 = tl.load(in_ptr2 + (98304 + r2 + 262144 * x1), rmask &
            xmask, eviction_policy='evict_last', other=0.0)
        tmp229 = tl.load(in_ptr29 + (r2 + 4096 * x3), rmask & xmask,
            eviction_policy='evict_first', other=0.0)
        tmp230 = tl.load(in_ptr2 + (102400 + r2 + 262144 * x1), rmask &
            xmask, eviction_policy='evict_last', other=0.0)
        tmp238 = tl.load(in_ptr30 + (r2 + 4096 * x3), rmask & xmask,
            eviction_policy='evict_first', other=0.0)
        tmp239 = tl.load(in_ptr2 + (106496 + r2 + 262144 * x1), rmask &
            xmask, eviction_policy='evict_last', other=0.0)
        tmp247 = tl.load(in_ptr31 + (r2 + 4096 * x3), rmask & xmask,
            eviction_policy='evict_first', other=0.0)
        tmp248 = tl.load(in_ptr2 + (110592 + r2 + 262144 * x1), rmask &
            xmask, eviction_policy='evict_last', other=0.0)
        tmp256 = tl.load(in_ptr32 + (r2 + 4096 * x3), rmask & xmask,
            eviction_policy='evict_first', other=0.0)
        tmp257 = tl.load(in_ptr2 + (114688 + r2 + 262144 * x1), rmask &
            xmask, eviction_policy='evict_last', other=0.0)
        tmp2 = tmp0 - tmp1
        tmp5 = tmp3 - tmp4
        tmp6 = tl_math.exp(tmp5)
        tmp8 = tmp6 / tmp7
        tmp9 = tmp2 * tmp8
        tmp10 = tl.broadcast_to(tmp9, [XBLOCK, RBLOCK])
        tmp12 = _tmp11 + tmp10
        _tmp11 = tl.where(rmask & xmask, tmp12, _tmp11)
        tmp15 = tmp14 - tmp4
        tmp16 = tl_math.exp(tmp15)
        tmp17 = tmp16 / tmp7
        tmp18 = tmp13 * tmp17
        tmp19 = tl.broadcast_to(tmp18, [XBLOCK, RBLOCK])
        tmp21 = _tmp20 + tmp19
        _tmp20 = tl.where(rmask & xmask, tmp21, _tmp20)
        tmp24 = tmp23 - tmp4
        tmp25 = tl_math.exp(tmp24)
        tmp26 = tmp25 / tmp7
        tmp27 = tmp22 * tmp26
        tmp28 = tl.broadcast_to(tmp27, [XBLOCK, RBLOCK])
        tmp30 = _tmp29 + tmp28
        _tmp29 = tl.where(rmask & xmask, tmp30, _tmp29)
        tmp33 = tmp32 - tmp4
        tmp34 = tl_math.exp(tmp33)
        tmp35 = tmp34 / tmp7
        tmp36 = tmp31 * tmp35
        tmp37 = tl.broadcast_to(tmp36, [XBLOCK, RBLOCK])
        tmp39 = _tmp38 + tmp37
        _tmp38 = tl.where(rmask & xmask, tmp39, _tmp38)
        tmp42 = tmp41 - tmp4
        tmp43 = tl_math.exp(tmp42)
        tmp44 = tmp43 / tmp7
        tmp45 = tmp40 * tmp44
        tmp46 = tl.broadcast_to(tmp45, [XBLOCK, RBLOCK])
        tmp48 = _tmp47 + tmp46
        _tmp47 = tl.where(rmask & xmask, tmp48, _tmp47)
        tmp51 = tmp50 - tmp4
        tmp52 = tl_math.exp(tmp51)
        tmp53 = tmp52 / tmp7
        tmp54 = tmp49 * tmp53
        tmp55 = tl.broadcast_to(tmp54, [XBLOCK, RBLOCK])
        tmp57 = _tmp56 + tmp55
        _tmp56 = tl.where(rmask & xmask, tmp57, _tmp56)
        tmp60 = tmp59 - tmp4
        tmp61 = tl_math.exp(tmp60)
        tmp62 = tmp61 / tmp7
        tmp63 = tmp58 * tmp62
        tmp64 = tl.broadcast_to(tmp63, [XBLOCK, RBLOCK])
        tmp66 = _tmp65 + tmp64
        _tmp65 = tl.where(rmask & xmask, tmp66, _tmp65)
        tmp69 = tmp68 - tmp4
        tmp70 = tl_math.exp(tmp69)
        tmp71 = tmp70 / tmp7
        tmp72 = tmp67 * tmp71
        tmp73 = tl.broadcast_to(tmp72, [XBLOCK, RBLOCK])
        tmp75 = _tmp74 + tmp73
        _tmp74 = tl.where(rmask & xmask, tmp75, _tmp74)
        tmp78 = tmp77 - tmp4
        tmp79 = tl_math.exp(tmp78)
        tmp80 = tmp79 / tmp7
        tmp81 = tmp76 * tmp80
        tmp82 = tl.broadcast_to(tmp81, [XBLOCK, RBLOCK])
        tmp84 = _tmp83 + tmp82
        _tmp83 = tl.where(rmask & xmask, tmp84, _tmp83)
        tmp87 = tmp86 - tmp4
        tmp88 = tl_math.exp(tmp87)
        tmp89 = tmp88 / tmp7
        tmp90 = tmp85 * tmp89
        tmp91 = tl.broadcast_to(tmp90, [XBLOCK, RBLOCK])
        tmp93 = _tmp92 + tmp91
        _tmp92 = tl.where(rmask & xmask, tmp93, _tmp92)
        tmp96 = tmp95 - tmp4
        tmp97 = tl_math.exp(tmp96)
        tmp98 = tmp97 / tmp7
        tmp99 = tmp94 * tmp98
        tmp100 = tl.broadcast_to(tmp99, [XBLOCK, RBLOCK])
        tmp102 = _tmp101 + tmp100
        _tmp101 = tl.where(rmask & xmask, tmp102, _tmp101)
        tmp105 = tmp104 - tmp4
        tmp106 = tl_math.exp(tmp105)
        tmp107 = tmp106 / tmp7
        tmp108 = tmp103 * tmp107
        tmp109 = tl.broadcast_to(tmp108, [XBLOCK, RBLOCK])
        tmp111 = _tmp110 + tmp109
        _tmp110 = tl.where(rmask & xmask, tmp111, _tmp110)
        tmp114 = tmp113 - tmp4
        tmp115 = tl_math.exp(tmp114)
        tmp116 = tmp115 / tmp7
        tmp117 = tmp112 * tmp116
        tmp118 = tl.broadcast_to(tmp117, [XBLOCK, RBLOCK])
        tmp120 = _tmp119 + tmp118
        _tmp119 = tl.where(rmask & xmask, tmp120, _tmp119)
        tmp123 = tmp122 - tmp4
        tmp124 = tl_math.exp(tmp123)
        tmp125 = tmp124 / tmp7
        tmp126 = tmp121 * tmp125
        tmp127 = tl.broadcast_to(tmp126, [XBLOCK, RBLOCK])
        tmp129 = _tmp128 + tmp127
        _tmp128 = tl.where(rmask & xmask, tmp129, _tmp128)
        tmp132 = tmp131 - tmp4
        tmp133 = tl_math.exp(tmp132)
        tmp134 = tmp133 / tmp7
        tmp135 = tmp130 * tmp134
        tmp136 = tl.broadcast_to(tmp135, [XBLOCK, RBLOCK])
        tmp138 = _tmp137 + tmp136
        _tmp137 = tl.where(rmask & xmask, tmp138, _tmp137)
        tmp141 = tmp140 - tmp4
        tmp142 = tl_math.exp(tmp141)
        tmp143 = tmp142 / tmp7
        tmp144 = tmp139 * tmp143
        tmp145 = tl.broadcast_to(tmp144, [XBLOCK, RBLOCK])
        tmp147 = _tmp146 + tmp145
        _tmp146 = tl.where(rmask & xmask, tmp147, _tmp146)
        tmp150 = tmp149 - tmp4
        tmp151 = tl_math.exp(tmp150)
        tmp152 = tmp151 / tmp7
        tmp153 = tmp148 * tmp152
        tmp154 = tl.broadcast_to(tmp153, [XBLOCK, RBLOCK])
        tmp156 = _tmp155 + tmp154
        _tmp155 = tl.where(rmask & xmask, tmp156, _tmp155)
        tmp159 = tmp158 - tmp4
        tmp160 = tl_math.exp(tmp159)
        tmp161 = tmp160 / tmp7
        tmp162 = tmp157 * tmp161
        tmp163 = tl.broadcast_to(tmp162, [XBLOCK, RBLOCK])
        tmp165 = _tmp164 + tmp163
        _tmp164 = tl.where(rmask & xmask, tmp165, _tmp164)
        tmp168 = tmp167 - tmp4
        tmp169 = tl_math.exp(tmp168)
        tmp170 = tmp169 / tmp7
        tmp171 = tmp166 * tmp170
        tmp172 = tl.broadcast_to(tmp171, [XBLOCK, RBLOCK])
        tmp174 = _tmp173 + tmp172
        _tmp173 = tl.where(rmask & xmask, tmp174, _tmp173)
        tmp177 = tmp176 - tmp4
        tmp178 = tl_math.exp(tmp177)
        tmp179 = tmp178 / tmp7
        tmp180 = tmp175 * tmp179
        tmp181 = tl.broadcast_to(tmp180, [XBLOCK, RBLOCK])
        tmp183 = _tmp182 + tmp181
        _tmp182 = tl.where(rmask & xmask, tmp183, _tmp182)
        tmp186 = tmp185 - tmp4
        tmp187 = tl_math.exp(tmp186)
        tmp188 = tmp187 / tmp7
        tmp189 = tmp184 * tmp188
        tmp190 = tl.broadcast_to(tmp189, [XBLOCK, RBLOCK])
        tmp192 = _tmp191 + tmp190
        _tmp191 = tl.where(rmask & xmask, tmp192, _tmp191)
        tmp195 = tmp194 - tmp4
        tmp196 = tl_math.exp(tmp195)
        tmp197 = tmp196 / tmp7
        tmp198 = tmp193 * tmp197
        tmp199 = tl.broadcast_to(tmp198, [XBLOCK, RBLOCK])
        tmp201 = _tmp200 + tmp199
        _tmp200 = tl.where(rmask & xmask, tmp201, _tmp200)
        tmp204 = tmp203 - tmp4
        tmp205 = tl_math.exp(tmp204)
        tmp206 = tmp205 / tmp7
        tmp207 = tmp202 * tmp206
        tmp208 = tl.broadcast_to(tmp207, [XBLOCK, RBLOCK])
        tmp210 = _tmp209 + tmp208
        _tmp209 = tl.where(rmask & xmask, tmp210, _tmp209)
        tmp213 = tmp212 - tmp4
        tmp214 = tl_math.exp(tmp213)
        tmp215 = tmp214 / tmp7
        tmp216 = tmp211 * tmp215
        tmp217 = tl.broadcast_to(tmp216, [XBLOCK, RBLOCK])
        tmp219 = _tmp218 + tmp217
        _tmp218 = tl.where(rmask & xmask, tmp219, _tmp218)
        tmp222 = tmp221 - tmp4
        tmp223 = tl_math.exp(tmp222)
        tmp224 = tmp223 / tmp7
        tmp225 = tmp220 * tmp224
        tmp226 = tl.broadcast_to(tmp225, [XBLOCK, RBLOCK])
        tmp228 = _tmp227 + tmp226
        _tmp227 = tl.where(rmask & xmask, tmp228, _tmp227)
        tmp231 = tmp230 - tmp4
        tmp232 = tl_math.exp(tmp231)
        tmp233 = tmp232 / tmp7
        tmp234 = tmp229 * tmp233
        tmp235 = tl.broadcast_to(tmp234, [XBLOCK, RBLOCK])
        tmp237 = _tmp236 + tmp235
        _tmp236 = tl.where(rmask & xmask, tmp237, _tmp236)
        tmp240 = tmp239 - tmp4
        tmp241 = tl_math.exp(tmp240)
        tmp242 = tmp241 / tmp7
        tmp243 = tmp238 * tmp242
        tmp244 = tl.broadcast_to(tmp243, [XBLOCK, RBLOCK])
        tmp246 = _tmp245 + tmp244
        _tmp245 = tl.where(rmask & xmask, tmp246, _tmp245)
        tmp249 = tmp248 - tmp4
        tmp250 = tl_math.exp(tmp249)
        tmp251 = tmp250 / tmp7
        tmp252 = tmp247 * tmp251
        tmp253 = tl.broadcast_to(tmp252, [XBLOCK, RBLOCK])
        tmp255 = _tmp254 + tmp253
        _tmp254 = tl.where(rmask & xmask, tmp255, _tmp254)
        tmp258 = tmp257 - tmp4
        tmp259 = tl_math.exp(tmp258)
        tmp260 = tmp259 / tmp7
        tmp261 = tmp256 * tmp260
        tmp262 = tl.broadcast_to(tmp261, [XBLOCK, RBLOCK])
        tmp264 = _tmp263 + tmp262
        _tmp263 = tl.where(rmask & xmask, tmp264, _tmp263)
    tmp11 = tl.sum(_tmp11, 1)[:, None]
    tl.store(out_ptr0 + x3, tmp11, xmask)
    tmp20 = tl.sum(_tmp20, 1)[:, None]
    tl.store(out_ptr1 + x3, tmp20, xmask)
    tmp29 = tl.sum(_tmp29, 1)[:, None]
    tl.store(out_ptr2 + x3, tmp29, xmask)
    tmp38 = tl.sum(_tmp38, 1)[:, None]
    tl.store(out_ptr3 + x3, tmp38, xmask)
    tmp47 = tl.sum(_tmp47, 1)[:, None]
    tl.store(out_ptr4 + x3, tmp47, xmask)
    tmp56 = tl.sum(_tmp56, 1)[:, None]
    tl.store(out_ptr5 + x3, tmp56, xmask)
    tmp65 = tl.sum(_tmp65, 1)[:, None]
    tl.store(out_ptr6 + x3, tmp65, xmask)
    tmp74 = tl.sum(_tmp74, 1)[:, None]
    tl.store(out_ptr7 + x3, tmp74, xmask)
    tmp83 = tl.sum(_tmp83, 1)[:, None]
    tl.store(out_ptr8 + x3, tmp83, xmask)
    tmp92 = tl.sum(_tmp92, 1)[:, None]
    tl.store(out_ptr9 + x3, tmp92, xmask)
    tmp101 = tl.sum(_tmp101, 1)[:, None]
    tl.store(out_ptr10 + x3, tmp101, xmask)
    tmp110 = tl.sum(_tmp110, 1)[:, None]
    tl.store(out_ptr11 + x3, tmp110, xmask)
    tmp119 = tl.sum(_tmp119, 1)[:, None]
    tl.store(out_ptr12 + x3, tmp119, xmask)
    tmp128 = tl.sum(_tmp128, 1)[:, None]
    tl.store(out_ptr13 + x3, tmp128, xmask)
    tmp137 = tl.sum(_tmp137, 1)[:, None]
    tl.store(out_ptr14 + x3, tmp137, xmask)
    tmp146 = tl.sum(_tmp146, 1)[:, None]
    tl.store(out_ptr15 + x3, tmp146, xmask)
    tmp155 = tl.sum(_tmp155, 1)[:, None]
    tl.store(out_ptr16 + x3, tmp155, xmask)
    tmp164 = tl.sum(_tmp164, 1)[:, None]
    tl.store(out_ptr17 + x3, tmp164, xmask)
    tmp173 = tl.sum(_tmp173, 1)[:, None]
    tl.store(out_ptr18 + x3, tmp173, xmask)
    tmp182 = tl.sum(_tmp182, 1)[:, None]
    tl.store(out_ptr19 + x3, tmp182, xmask)
    tmp191 = tl.sum(_tmp191, 1)[:, None]
    tl.store(out_ptr20 + x3, tmp191, xmask)
    tmp200 = tl.sum(_tmp200, 1)[:, None]
    tl.store(out_ptr21 + x3, tmp200, xmask)
    tmp209 = tl.sum(_tmp209, 1)[:, None]
    tl.store(out_ptr22 + x3, tmp209, xmask)
    tmp218 = tl.sum(_tmp218, 1)[:, None]
    tl.store(out_ptr23 + x3, tmp218, xmask)
    tmp227 = tl.sum(_tmp227, 1)[:, None]
    tl.store(out_ptr24 + x3, tmp227, xmask)
    tmp236 = tl.sum(_tmp236, 1)[:, None]
    tl.store(out_ptr25 + x3, tmp236, xmask)
    tmp245 = tl.sum(_tmp245, 1)[:, None]
    tl.store(out_ptr26 + x3, tmp245, xmask)
    tmp254 = tl.sum(_tmp254, 1)[:, None]
    tl.store(out_ptr27 + x3, tmp254, xmask)
    tmp263 = tl.sum(_tmp263, 1)[:, None]
    tl.store(out_ptr28 + x3, tmp263, xmask)


@triton.jit
def triton_red_fused_mul_sum_4(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, out_ptr0, out_ptr1,
    out_ptr2, out_ptr3, out_ptr4, out_ptr5, out_ptr6, out_ptr7, out_ptr8,
    out_ptr9, out_ptr10, out_ptr11, out_ptr12, out_ptr13, out_ptr14,
    out_ptr15, out_ptr16, out_ptr17, out_ptr18, out_ptr19, out_ptr20,
    out_ptr21, out_ptr22, out_ptr23, out_ptr24, out_ptr25, out_ptr26,
    out_ptr27, xnumel, rnumel, XBLOCK: tl.constexpr, RBLOCK: tl.constexpr):
    xnumel = 512
    rnumel = 4096
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:, None]
    xmask = xindex < xnumel
    rbase = tl.arange(0, RBLOCK)[None, :]
    x3 = xindex
    x1 = xindex // 128
    _tmp9 = tl.full([XBLOCK, RBLOCK], 0, tl.float32)
    _tmp18 = tl.full([XBLOCK, RBLOCK], 0, tl.float32)
    _tmp27 = tl.full([XBLOCK, RBLOCK], 0, tl.float32)
    _tmp36 = tl.full([XBLOCK, RBLOCK], 0, tl.float32)
    _tmp45 = tl.full([XBLOCK, RBLOCK], 0, tl.float32)
    _tmp54 = tl.full([XBLOCK, RBLOCK], 0, tl.float32)
    _tmp63 = tl.full([XBLOCK, RBLOCK], 0, tl.float32)
    _tmp72 = tl.full([XBLOCK, RBLOCK], 0, tl.float32)
    _tmp81 = tl.full([XBLOCK, RBLOCK], 0, tl.float32)
    _tmp90 = tl.full([XBLOCK, RBLOCK], 0, tl.float32)
    _tmp99 = tl.full([XBLOCK, RBLOCK], 0, tl.float32)
    _tmp108 = tl.full([XBLOCK, RBLOCK], 0, tl.float32)
    _tmp117 = tl.full([XBLOCK, RBLOCK], 0, tl.float32)
    _tmp126 = tl.full([XBLOCK, RBLOCK], 0, tl.float32)
    _tmp135 = tl.full([XBLOCK, RBLOCK], 0, tl.float32)
    _tmp144 = tl.full([XBLOCK, RBLOCK], 0, tl.float32)
    _tmp153 = tl.full([XBLOCK, RBLOCK], 0, tl.float32)
    _tmp162 = tl.full([XBLOCK, RBLOCK], 0, tl.float32)
    _tmp171 = tl.full([XBLOCK, RBLOCK], 0, tl.float32)
    _tmp180 = tl.full([XBLOCK, RBLOCK], 0, tl.float32)
    _tmp189 = tl.full([XBLOCK, RBLOCK], 0, tl.float32)
    _tmp198 = tl.full([XBLOCK, RBLOCK], 0, tl.float32)
    _tmp207 = tl.full([XBLOCK, RBLOCK], 0, tl.float32)
    _tmp216 = tl.full([XBLOCK, RBLOCK], 0, tl.float32)
    _tmp225 = tl.full([XBLOCK, RBLOCK], 0, tl.float32)
    _tmp234 = tl.full([XBLOCK, RBLOCK], 0, tl.float32)
    _tmp243 = tl.full([XBLOCK, RBLOCK], 0, tl.float32)
    _tmp252 = tl.full([XBLOCK, RBLOCK], 0, tl.float32)
    for roffset in range(0, rnumel, RBLOCK):
        rindex = roffset + rbase
        rmask = rindex < rnumel
        r2 = rindex
        tmp0 = tl.load(in_ptr0 + (r2 + 4096 * x3), rmask & xmask,
            eviction_policy='evict_first', other=0.0)
        tmp1 = tl.load(in_ptr1 + (118784 + r2 + 262144 * x1), rmask & xmask,
            eviction_policy='evict_last', other=0.0)
        tmp2 = tl.load(in_ptr2 + (r2 + 4096 * x1), rmask & xmask,
            eviction_policy='evict_last', other=0.0)
        tmp5 = tl.load(in_ptr3 + (r2 + 4096 * x1), rmask & xmask,
            eviction_policy='evict_last', other=0.0)
        tmp11 = tl.load(in_ptr4 + (r2 + 4096 * x3), rmask & xmask,
            eviction_policy='evict_first', other=0.0)
        tmp12 = tl.load(in_ptr1 + (122880 + r2 + 262144 * x1), rmask &
            xmask, eviction_policy='evict_last', other=0.0)
        tmp20 = tl.load(in_ptr5 + (r2 + 4096 * x3), rmask & xmask,
            eviction_policy='evict_first', other=0.0)
        tmp21 = tl.load(in_ptr1 + (126976 + r2 + 262144 * x1), rmask &
            xmask, eviction_policy='evict_last', other=0.0)
        tmp29 = tl.load(in_ptr6 + (r2 + 4096 * x3), rmask & xmask,
            eviction_policy='evict_first', other=0.0)
        tmp30 = tl.load(in_ptr1 + (131072 + r2 + 262144 * x1), rmask &
            xmask, eviction_policy='evict_last', other=0.0)
        tmp38 = tl.load(in_ptr7 + (r2 + 4096 * x3), rmask & xmask,
            eviction_policy='evict_first', other=0.0)
        tmp39 = tl.load(in_ptr1 + (135168 + r2 + 262144 * x1), rmask &
            xmask, eviction_policy='evict_last', other=0.0)
        tmp47 = tl.load(in_ptr8 + (r2 + 4096 * x3), rmask & xmask,
            eviction_policy='evict_first', other=0.0)
        tmp48 = tl.load(in_ptr1 + (139264 + r2 + 262144 * x1), rmask &
            xmask, eviction_policy='evict_last', other=0.0)
        tmp56 = tl.load(in_ptr9 + (r2 + 4096 * x3), rmask & xmask,
            eviction_policy='evict_first', other=0.0)
        tmp57 = tl.load(in_ptr1 + (143360 + r2 + 262144 * x1), rmask &
            xmask, eviction_policy='evict_last', other=0.0)
        tmp65 = tl.load(in_ptr10 + (r2 + 4096 * x3), rmask & xmask,
            eviction_policy='evict_first', other=0.0)
        tmp66 = tl.load(in_ptr1 + (147456 + r2 + 262144 * x1), rmask &
            xmask, eviction_policy='evict_last', other=0.0)
        tmp74 = tl.load(in_ptr11 + (r2 + 4096 * x3), rmask & xmask,
            eviction_policy='evict_first', other=0.0)
        tmp75 = tl.load(in_ptr1 + (151552 + r2 + 262144 * x1), rmask &
            xmask, eviction_policy='evict_last', other=0.0)
        tmp83 = tl.load(in_ptr12 + (r2 + 4096 * x3), rmask & xmask,
            eviction_policy='evict_first', other=0.0)
        tmp84 = tl.load(in_ptr1 + (155648 + r2 + 262144 * x1), rmask &
            xmask, eviction_policy='evict_last', other=0.0)
        tmp92 = tl.load(in_ptr13 + (r2 + 4096 * x3), rmask & xmask,
            eviction_policy='evict_first', other=0.0)
        tmp93 = tl.load(in_ptr1 + (159744 + r2 + 262144 * x1), rmask &
            xmask, eviction_policy='evict_last', other=0.0)
        tmp101 = tl.load(in_ptr14 + (r2 + 4096 * x3), rmask & xmask,
            eviction_policy='evict_first', other=0.0)
        tmp102 = tl.load(in_ptr1 + (163840 + r2 + 262144 * x1), rmask &
            xmask, eviction_policy='evict_last', other=0.0)
        tmp110 = tl.load(in_ptr15 + (r2 + 4096 * x3), rmask & xmask,
            eviction_policy='evict_first', other=0.0)
        tmp111 = tl.load(in_ptr1 + (167936 + r2 + 262144 * x1), rmask &
            xmask, eviction_policy='evict_last', other=0.0)
        tmp119 = tl.load(in_ptr16 + (r2 + 4096 * x3), rmask & xmask,
            eviction_policy='evict_first', other=0.0)
        tmp120 = tl.load(in_ptr1 + (172032 + r2 + 262144 * x1), rmask &
            xmask, eviction_policy='evict_last', other=0.0)
        tmp128 = tl.load(in_ptr17 + (r2 + 4096 * x3), rmask & xmask,
            eviction_policy='evict_first', other=0.0)
        tmp129 = tl.load(in_ptr1 + (176128 + r2 + 262144 * x1), rmask &
            xmask, eviction_policy='evict_last', other=0.0)
        tmp137 = tl.load(in_ptr18 + (r2 + 4096 * x3), rmask & xmask,
            eviction_policy='evict_first', other=0.0)
        tmp138 = tl.load(in_ptr1 + (180224 + r2 + 262144 * x1), rmask &
            xmask, eviction_policy='evict_last', other=0.0)
        tmp146 = tl.load(in_ptr19 + (r2 + 4096 * x3), rmask & xmask,
            eviction_policy='evict_first', other=0.0)
        tmp147 = tl.load(in_ptr1 + (184320 + r2 + 262144 * x1), rmask &
            xmask, eviction_policy='evict_last', other=0.0)
        tmp155 = tl.load(in_ptr20 + (r2 + 4096 * x3), rmask & xmask,
            eviction_policy='evict_first', other=0.0)
        tmp156 = tl.load(in_ptr1 + (188416 + r2 + 262144 * x1), rmask &
            xmask, eviction_policy='evict_last', other=0.0)
        tmp164 = tl.load(in_ptr21 + (r2 + 4096 * x3), rmask & xmask,
            eviction_policy='evict_first', other=0.0)
        tmp165 = tl.load(in_ptr1 + (192512 + r2 + 262144 * x1), rmask &
            xmask, eviction_policy='evict_last', other=0.0)
        tmp173 = tl.load(in_ptr22 + (r2 + 4096 * x3), rmask & xmask,
            eviction_policy='evict_first', other=0.0)
        tmp174 = tl.load(in_ptr1 + (196608 + r2 + 262144 * x1), rmask &
            xmask, eviction_policy='evict_last', other=0.0)
        tmp182 = tl.load(in_ptr23 + (r2 + 4096 * x3), rmask & xmask,
            eviction_policy='evict_first', other=0.0)
        tmp183 = tl.load(in_ptr1 + (200704 + r2 + 262144 * x1), rmask &
            xmask, eviction_policy='evict_last', other=0.0)
        tmp191 = tl.load(in_ptr24 + (r2 + 4096 * x3), rmask & xmask,
            eviction_policy='evict_first', other=0.0)
        tmp192 = tl.load(in_ptr1 + (204800 + r2 + 262144 * x1), rmask &
            xmask, eviction_policy='evict_last', other=0.0)
        tmp200 = tl.load(in_ptr25 + (r2 + 4096 * x3), rmask & xmask,
            eviction_policy='evict_first', other=0.0)
        tmp201 = tl.load(in_ptr1 + (208896 + r2 + 262144 * x1), rmask &
            xmask, eviction_policy='evict_last', other=0.0)
        tmp209 = tl.load(in_ptr26 + (r2 + 4096 * x3), rmask & xmask,
            eviction_policy='evict_first', other=0.0)
        tmp210 = tl.load(in_ptr1 + (212992 + r2 + 262144 * x1), rmask &
            xmask, eviction_policy='evict_last', other=0.0)
        tmp218 = tl.load(in_ptr27 + (r2 + 4096 * x3), rmask & xmask,
            eviction_policy='evict_first', other=0.0)
        tmp219 = tl.load(in_ptr1 + (217088 + r2 + 262144 * x1), rmask &
            xmask, eviction_policy='evict_last', other=0.0)
        tmp227 = tl.load(in_ptr28 + (r2 + 4096 * x3), rmask & xmask,
            eviction_policy='evict_first', other=0.0)
        tmp228 = tl.load(in_ptr1 + (221184 + r2 + 262144 * x1), rmask &
            xmask, eviction_policy='evict_last', other=0.0)
        tmp236 = tl.load(in_ptr29 + (r2 + 4096 * x3), rmask & xmask,
            eviction_policy='evict_first', other=0.0)
        tmp237 = tl.load(in_ptr1 + (225280 + r2 + 262144 * x1), rmask &
            xmask, eviction_policy='evict_last', other=0.0)
        tmp245 = tl.load(in_ptr30 + (r2 + 4096 * x3), rmask & xmask,
            eviction_policy='evict_first', other=0.0)
        tmp246 = tl.load(in_ptr1 + (229376 + r2 + 262144 * x1), rmask &
            xmask, eviction_policy='evict_last', other=0.0)
        tmp3 = tmp1 - tmp2
        tmp4 = tl_math.exp(tmp3)
        tmp6 = tmp4 / tmp5
        tmp7 = tmp0 * tmp6
        tmp8 = tl.broadcast_to(tmp7, [XBLOCK, RBLOCK])
        tmp10 = _tmp9 + tmp8
        _tmp9 = tl.where(rmask & xmask, tmp10, _tmp9)
        tmp13 = tmp12 - tmp2
        tmp14 = tl_math.exp(tmp13)
        tmp15 = tmp14 / tmp5
        tmp16 = tmp11 * tmp15
        tmp17 = tl.broadcast_to(tmp16, [XBLOCK, RBLOCK])
        tmp19 = _tmp18 + tmp17
        _tmp18 = tl.where(rmask & xmask, tmp19, _tmp18)
        tmp22 = tmp21 - tmp2
        tmp23 = tl_math.exp(tmp22)
        tmp24 = tmp23 / tmp5
        tmp25 = tmp20 * tmp24
        tmp26 = tl.broadcast_to(tmp25, [XBLOCK, RBLOCK])
        tmp28 = _tmp27 + tmp26
        _tmp27 = tl.where(rmask & xmask, tmp28, _tmp27)
        tmp31 = tmp30 - tmp2
        tmp32 = tl_math.exp(tmp31)
        tmp33 = tmp32 / tmp5
        tmp34 = tmp29 * tmp33
        tmp35 = tl.broadcast_to(tmp34, [XBLOCK, RBLOCK])
        tmp37 = _tmp36 + tmp35
        _tmp36 = tl.where(rmask & xmask, tmp37, _tmp36)
        tmp40 = tmp39 - tmp2
        tmp41 = tl_math.exp(tmp40)
        tmp42 = tmp41 / tmp5
        tmp43 = tmp38 * tmp42
        tmp44 = tl.broadcast_to(tmp43, [XBLOCK, RBLOCK])
        tmp46 = _tmp45 + tmp44
        _tmp45 = tl.where(rmask & xmask, tmp46, _tmp45)
        tmp49 = tmp48 - tmp2
        tmp50 = tl_math.exp(tmp49)
        tmp51 = tmp50 / tmp5
        tmp52 = tmp47 * tmp51
        tmp53 = tl.broadcast_to(tmp52, [XBLOCK, RBLOCK])
        tmp55 = _tmp54 + tmp53
        _tmp54 = tl.where(rmask & xmask, tmp55, _tmp54)
        tmp58 = tmp57 - tmp2
        tmp59 = tl_math.exp(tmp58)
        tmp60 = tmp59 / tmp5
        tmp61 = tmp56 * tmp60
        tmp62 = tl.broadcast_to(tmp61, [XBLOCK, RBLOCK])
        tmp64 = _tmp63 + tmp62
        _tmp63 = tl.where(rmask & xmask, tmp64, _tmp63)
        tmp67 = tmp66 - tmp2
        tmp68 = tl_math.exp(tmp67)
        tmp69 = tmp68 / tmp5
        tmp70 = tmp65 * tmp69
        tmp71 = tl.broadcast_to(tmp70, [XBLOCK, RBLOCK])
        tmp73 = _tmp72 + tmp71
        _tmp72 = tl.where(rmask & xmask, tmp73, _tmp72)
        tmp76 = tmp75 - tmp2
        tmp77 = tl_math.exp(tmp76)
        tmp78 = tmp77 / tmp5
        tmp79 = tmp74 * tmp78
        tmp80 = tl.broadcast_to(tmp79, [XBLOCK, RBLOCK])
        tmp82 = _tmp81 + tmp80
        _tmp81 = tl.where(rmask & xmask, tmp82, _tmp81)
        tmp85 = tmp84 - tmp2
        tmp86 = tl_math.exp(tmp85)
        tmp87 = tmp86 / tmp5
        tmp88 = tmp83 * tmp87
        tmp89 = tl.broadcast_to(tmp88, [XBLOCK, RBLOCK])
        tmp91 = _tmp90 + tmp89
        _tmp90 = tl.where(rmask & xmask, tmp91, _tmp90)
        tmp94 = tmp93 - tmp2
        tmp95 = tl_math.exp(tmp94)
        tmp96 = tmp95 / tmp5
        tmp97 = tmp92 * tmp96
        tmp98 = tl.broadcast_to(tmp97, [XBLOCK, RBLOCK])
        tmp100 = _tmp99 + tmp98
        _tmp99 = tl.where(rmask & xmask, tmp100, _tmp99)
        tmp103 = tmp102 - tmp2
        tmp104 = tl_math.exp(tmp103)
        tmp105 = tmp104 / tmp5
        tmp106 = tmp101 * tmp105
        tmp107 = tl.broadcast_to(tmp106, [XBLOCK, RBLOCK])
        tmp109 = _tmp108 + tmp107
        _tmp108 = tl.where(rmask & xmask, tmp109, _tmp108)
        tmp112 = tmp111 - tmp2
        tmp113 = tl_math.exp(tmp112)
        tmp114 = tmp113 / tmp5
        tmp115 = tmp110 * tmp114
        tmp116 = tl.broadcast_to(tmp115, [XBLOCK, RBLOCK])
        tmp118 = _tmp117 + tmp116
        _tmp117 = tl.where(rmask & xmask, tmp118, _tmp117)
        tmp121 = tmp120 - tmp2
        tmp122 = tl_math.exp(tmp121)
        tmp123 = tmp122 / tmp5
        tmp124 = tmp119 * tmp123
        tmp125 = tl.broadcast_to(tmp124, [XBLOCK, RBLOCK])
        tmp127 = _tmp126 + tmp125
        _tmp126 = tl.where(rmask & xmask, tmp127, _tmp126)
        tmp130 = tmp129 - tmp2
        tmp131 = tl_math.exp(tmp130)
        tmp132 = tmp131 / tmp5
        tmp133 = tmp128 * tmp132
        tmp134 = tl.broadcast_to(tmp133, [XBLOCK, RBLOCK])
        tmp136 = _tmp135 + tmp134
        _tmp135 = tl.where(rmask & xmask, tmp136, _tmp135)
        tmp139 = tmp138 - tmp2
        tmp140 = tl_math.exp(tmp139)
        tmp141 = tmp140 / tmp5
        tmp142 = tmp137 * tmp141
        tmp143 = tl.broadcast_to(tmp142, [XBLOCK, RBLOCK])
        tmp145 = _tmp144 + tmp143
        _tmp144 = tl.where(rmask & xmask, tmp145, _tmp144)
        tmp148 = tmp147 - tmp2
        tmp149 = tl_math.exp(tmp148)
        tmp150 = tmp149 / tmp5
        tmp151 = tmp146 * tmp150
        tmp152 = tl.broadcast_to(tmp151, [XBLOCK, RBLOCK])
        tmp154 = _tmp153 + tmp152
        _tmp153 = tl.where(rmask & xmask, tmp154, _tmp153)
        tmp157 = tmp156 - tmp2
        tmp158 = tl_math.exp(tmp157)
        tmp159 = tmp158 / tmp5
        tmp160 = tmp155 * tmp159
        tmp161 = tl.broadcast_to(tmp160, [XBLOCK, RBLOCK])
        tmp163 = _tmp162 + tmp161
        _tmp162 = tl.where(rmask & xmask, tmp163, _tmp162)
        tmp166 = tmp165 - tmp2
        tmp167 = tl_math.exp(tmp166)
        tmp168 = tmp167 / tmp5
        tmp169 = tmp164 * tmp168
        tmp170 = tl.broadcast_to(tmp169, [XBLOCK, RBLOCK])
        tmp172 = _tmp171 + tmp170
        _tmp171 = tl.where(rmask & xmask, tmp172, _tmp171)
        tmp175 = tmp174 - tmp2
        tmp176 = tl_math.exp(tmp175)
        tmp177 = tmp176 / tmp5
        tmp178 = tmp173 * tmp177
        tmp179 = tl.broadcast_to(tmp178, [XBLOCK, RBLOCK])
        tmp181 = _tmp180 + tmp179
        _tmp180 = tl.where(rmask & xmask, tmp181, _tmp180)
        tmp184 = tmp183 - tmp2
        tmp185 = tl_math.exp(tmp184)
        tmp186 = tmp185 / tmp5
        tmp187 = tmp182 * tmp186
        tmp188 = tl.broadcast_to(tmp187, [XBLOCK, RBLOCK])
        tmp190 = _tmp189 + tmp188
        _tmp189 = tl.where(rmask & xmask, tmp190, _tmp189)
        tmp193 = tmp192 - tmp2
        tmp194 = tl_math.exp(tmp193)
        tmp195 = tmp194 / tmp5
        tmp196 = tmp191 * tmp195
        tmp197 = tl.broadcast_to(tmp196, [XBLOCK, RBLOCK])
        tmp199 = _tmp198 + tmp197
        _tmp198 = tl.where(rmask & xmask, tmp199, _tmp198)
        tmp202 = tmp201 - tmp2
        tmp203 = tl_math.exp(tmp202)
        tmp204 = tmp203 / tmp5
        tmp205 = tmp200 * tmp204
        tmp206 = tl.broadcast_to(tmp205, [XBLOCK, RBLOCK])
        tmp208 = _tmp207 + tmp206
        _tmp207 = tl.where(rmask & xmask, tmp208, _tmp207)
        tmp211 = tmp210 - tmp2
        tmp212 = tl_math.exp(tmp211)
        tmp213 = tmp212 / tmp5
        tmp214 = tmp209 * tmp213
        tmp215 = tl.broadcast_to(tmp214, [XBLOCK, RBLOCK])
        tmp217 = _tmp216 + tmp215
        _tmp216 = tl.where(rmask & xmask, tmp217, _tmp216)
        tmp220 = tmp219 - tmp2
        tmp221 = tl_math.exp(tmp220)
        tmp222 = tmp221 / tmp5
        tmp223 = tmp218 * tmp222
        tmp224 = tl.broadcast_to(tmp223, [XBLOCK, RBLOCK])
        tmp226 = _tmp225 + tmp224
        _tmp225 = tl.where(rmask & xmask, tmp226, _tmp225)
        tmp229 = tmp228 - tmp2
        tmp230 = tl_math.exp(tmp229)
        tmp231 = tmp230 / tmp5
        tmp232 = tmp227 * tmp231
        tmp233 = tl.broadcast_to(tmp232, [XBLOCK, RBLOCK])
        tmp235 = _tmp234 + tmp233
        _tmp234 = tl.where(rmask & xmask, tmp235, _tmp234)
        tmp238 = tmp237 - tmp2
        tmp239 = tl_math.exp(tmp238)
        tmp240 = tmp239 / tmp5
        tmp241 = tmp236 * tmp240
        tmp242 = tl.broadcast_to(tmp241, [XBLOCK, RBLOCK])
        tmp244 = _tmp243 + tmp242
        _tmp243 = tl.where(rmask & xmask, tmp244, _tmp243)
        tmp247 = tmp246 - tmp2
        tmp248 = tl_math.exp(tmp247)
        tmp249 = tmp248 / tmp5
        tmp250 = tmp245 * tmp249
        tmp251 = tl.broadcast_to(tmp250, [XBLOCK, RBLOCK])
        tmp253 = _tmp252 + tmp251
        _tmp252 = tl.where(rmask & xmask, tmp253, _tmp252)
    tmp9 = tl.sum(_tmp9, 1)[:, None]
    tl.store(out_ptr0 + x3, tmp9, xmask)
    tmp18 = tl.sum(_tmp18, 1)[:, None]
    tl.store(out_ptr1 + x3, tmp18, xmask)
    tmp27 = tl.sum(_tmp27, 1)[:, None]
    tl.store(out_ptr2 + x3, tmp27, xmask)
    tmp36 = tl.sum(_tmp36, 1)[:, None]
    tl.store(out_ptr3 + x3, tmp36, xmask)
    tmp45 = tl.sum(_tmp45, 1)[:, None]
    tl.store(out_ptr4 + x3, tmp45, xmask)
    tmp54 = tl.sum(_tmp54, 1)[:, None]
    tl.store(out_ptr5 + x3, tmp54, xmask)
    tmp63 = tl.sum(_tmp63, 1)[:, None]
    tl.store(out_ptr6 + x3, tmp63, xmask)
    tmp72 = tl.sum(_tmp72, 1)[:, None]
    tl.store(out_ptr7 + x3, tmp72, xmask)
    tmp81 = tl.sum(_tmp81, 1)[:, None]
    tl.store(out_ptr8 + x3, tmp81, xmask)
    tmp90 = tl.sum(_tmp90, 1)[:, None]
    tl.store(out_ptr9 + x3, tmp90, xmask)
    tmp99 = tl.sum(_tmp99, 1)[:, None]
    tl.store(out_ptr10 + x3, tmp99, xmask)
    tmp108 = tl.sum(_tmp108, 1)[:, None]
    tl.store(out_ptr11 + x3, tmp108, xmask)
    tmp117 = tl.sum(_tmp117, 1)[:, None]
    tl.store(out_ptr12 + x3, tmp117, xmask)
    tmp126 = tl.sum(_tmp126, 1)[:, None]
    tl.store(out_ptr13 + x3, tmp126, xmask)
    tmp135 = tl.sum(_tmp135, 1)[:, None]
    tl.store(out_ptr14 + x3, tmp135, xmask)
    tmp144 = tl.sum(_tmp144, 1)[:, None]
    tl.store(out_ptr15 + x3, tmp144, xmask)
    tmp153 = tl.sum(_tmp153, 1)[:, None]
    tl.store(out_ptr16 + x3, tmp153, xmask)
    tmp162 = tl.sum(_tmp162, 1)[:, None]
    tl.store(out_ptr17 + x3, tmp162, xmask)
    tmp171 = tl.sum(_tmp171, 1)[:, None]
    tl.store(out_ptr18 + x3, tmp171, xmask)
    tmp180 = tl.sum(_tmp180, 1)[:, None]
    tl.store(out_ptr19 + x3, tmp180, xmask)
    tmp189 = tl.sum(_tmp189, 1)[:, None]
    tl.store(out_ptr20 + x3, tmp189, xmask)
    tmp198 = tl.sum(_tmp198, 1)[:, None]
    tl.store(out_ptr21 + x3, tmp198, xmask)
    tmp207 = tl.sum(_tmp207, 1)[:, None]
    tl.store(out_ptr22 + x3, tmp207, xmask)
    tmp216 = tl.sum(_tmp216, 1)[:, None]
    tl.store(out_ptr23 + x3, tmp216, xmask)
    tmp225 = tl.sum(_tmp225, 1)[:, None]
    tl.store(out_ptr24 + x3, tmp225, xmask)
    tmp234 = tl.sum(_tmp234, 1)[:, None]
    tl.store(out_ptr25 + x3, tmp234, xmask)
    tmp243 = tl.sum(_tmp243, 1)[:, None]
    tl.store(out_ptr26 + x3, tmp243, xmask)
    tmp252 = tl.sum(_tmp252, 1)[:, None]
    tl.store(out_ptr27 + x3, tmp252, xmask)


@triton.jit
def triton_red_fused_mul_sum_5(in_ptr0, in_ptr1, in_ptr2, in_ptr3, in_ptr4,
    in_ptr5, in_ptr6, in_ptr7, in_ptr8, in_ptr9, out_ptr0, out_ptr1,
    out_ptr2, out_ptr3, out_ptr4, out_ptr5, out_ptr6, xnumel, rnumel,
    XBLOCK: tl.constexpr, RBLOCK: tl.constexpr):
    xnumel = 512
    rnumel = 4096
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:, None]
    xmask = xindex < xnumel
    rbase = tl.arange(0, RBLOCK)[None, :]
    x3 = xindex
    x1 = xindex // 128
    _tmp9 = tl.full([XBLOCK, RBLOCK], 0, tl.float32)
    _tmp18 = tl.full([XBLOCK, RBLOCK], 0, tl.float32)
    _tmp27 = tl.full([XBLOCK, RBLOCK], 0, tl.float32)
    _tmp36 = tl.full([XBLOCK, RBLOCK], 0, tl.float32)
    _tmp45 = tl.full([XBLOCK, RBLOCK], 0, tl.float32)
    _tmp54 = tl.full([XBLOCK, RBLOCK], 0, tl.float32)
    _tmp63 = tl.full([XBLOCK, RBLOCK], 0, tl.float32)
    for roffset in range(0, rnumel, RBLOCK):
        rindex = roffset + rbase
        rmask = rindex < rnumel
        r2 = rindex
        tmp0 = tl.load(in_ptr0 + (r2 + 4096 * x3), rmask & xmask,
            eviction_policy='evict_first', other=0.0)
        tmp1 = tl.load(in_ptr1 + (233472 + r2 + 262144 * x1), rmask & xmask,
            eviction_policy='evict_last', other=0.0)
        tmp2 = tl.load(in_ptr2 + (r2 + 4096 * x1), rmask & xmask,
            eviction_policy='evict_last', other=0.0)
        tmp5 = tl.load(in_ptr3 + (r2 + 4096 * x1), rmask & xmask,
            eviction_policy='evict_last', other=0.0)
        tmp11 = tl.load(in_ptr4 + (r2 + 4096 * x3), rmask & xmask,
            eviction_policy='evict_first', other=0.0)
        tmp12 = tl.load(in_ptr1 + (237568 + r2 + 262144 * x1), rmask &
            xmask, eviction_policy='evict_last', other=0.0)
        tmp20 = tl.load(in_ptr5 + (r2 + 4096 * x3), rmask & xmask,
            eviction_policy='evict_first', other=0.0)
        tmp21 = tl.load(in_ptr1 + (241664 + r2 + 262144 * x1), rmask &
            xmask, eviction_policy='evict_last', other=0.0)
        tmp29 = tl.load(in_ptr6 + (r2 + 4096 * x3), rmask & xmask,
            eviction_policy='evict_first', other=0.0)
        tmp30 = tl.load(in_ptr1 + (245760 + r2 + 262144 * x1), rmask &
            xmask, eviction_policy='evict_last', other=0.0)
        tmp38 = tl.load(in_ptr7 + (r2 + 4096 * x3), rmask & xmask,
            eviction_policy='evict_first', other=0.0)
        tmp39 = tl.load(in_ptr1 + (249856 + r2 + 262144 * x1), rmask &
            xmask, eviction_policy='evict_last', other=0.0)
        tmp47 = tl.load(in_ptr8 + (r2 + 4096 * x3), rmask & xmask,
            eviction_policy='evict_first', other=0.0)
        tmp48 = tl.load(in_ptr1 + (253952 + r2 + 262144 * x1), rmask &
            xmask, eviction_policy='evict_last', other=0.0)
        tmp56 = tl.load(in_ptr9 + (r2 + 4096 * x3), rmask & xmask,
            eviction_policy='evict_first', other=0.0)
        tmp57 = tl.load(in_ptr1 + (258048 + r2 + 262144 * x1), rmask &
            xmask, eviction_policy='evict_last', other=0.0)
        tmp3 = tmp1 - tmp2
        tmp4 = tl_math.exp(tmp3)
        tmp6 = tmp4 / tmp5
        tmp7 = tmp0 * tmp6
        tmp8 = tl.broadcast_to(tmp7, [XBLOCK, RBLOCK])
        tmp10 = _tmp9 + tmp8
        _tmp9 = tl.where(rmask & xmask, tmp10, _tmp9)
        tmp13 = tmp12 - tmp2
        tmp14 = tl_math.exp(tmp13)
        tmp15 = tmp14 / tmp5
        tmp16 = tmp11 * tmp15
        tmp17 = tl.broadcast_to(tmp16, [XBLOCK, RBLOCK])
        tmp19 = _tmp18 + tmp17
        _tmp18 = tl.where(rmask & xmask, tmp19, _tmp18)
        tmp22 = tmp21 - tmp2
        tmp23 = tl_math.exp(tmp22)
        tmp24 = tmp23 / tmp5
        tmp25 = tmp20 * tmp24
        tmp26 = tl.broadcast_to(tmp25, [XBLOCK, RBLOCK])
        tmp28 = _tmp27 + tmp26
        _tmp27 = tl.where(rmask & xmask, tmp28, _tmp27)
        tmp31 = tmp30 - tmp2
        tmp32 = tl_math.exp(tmp31)
        tmp33 = tmp32 / tmp5
        tmp34 = tmp29 * tmp33
        tmp35 = tl.broadcast_to(tmp34, [XBLOCK, RBLOCK])
        tmp37 = _tmp36 + tmp35
        _tmp36 = tl.where(rmask & xmask, tmp37, _tmp36)
        tmp40 = tmp39 - tmp2
        tmp41 = tl_math.exp(tmp40)
        tmp42 = tmp41 / tmp5
        tmp43 = tmp38 * tmp42
        tmp44 = tl.broadcast_to(tmp43, [XBLOCK, RBLOCK])
        tmp46 = _tmp45 + tmp44
        _tmp45 = tl.where(rmask & xmask, tmp46, _tmp45)
        tmp49 = tmp48 - tmp2
        tmp50 = tl_math.exp(tmp49)
        tmp51 = tmp50 / tmp5
        tmp52 = tmp47 * tmp51
        tmp53 = tl.broadcast_to(tmp52, [XBLOCK, RBLOCK])
        tmp55 = _tmp54 + tmp53
        _tmp54 = tl.where(rmask & xmask, tmp55, _tmp54)
        tmp58 = tmp57 - tmp2
        tmp59 = tl_math.exp(tmp58)
        tmp60 = tmp59 / tmp5
        tmp61 = tmp56 * tmp60
        tmp62 = tl.broadcast_to(tmp61, [XBLOCK, RBLOCK])
        tmp64 = _tmp63 + tmp62
        _tmp63 = tl.where(rmask & xmask, tmp64, _tmp63)
    tmp9 = tl.sum(_tmp9, 1)[:, None]
    tl.store(out_ptr0 + x3, tmp9, xmask)
    tmp18 = tl.sum(_tmp18, 1)[:, None]
    tl.store(out_ptr1 + x3, tmp18, xmask)
    tmp27 = tl.sum(_tmp27, 1)[:, None]
    tl.store(out_ptr2 + x3, tmp27, xmask)
    tmp36 = tl.sum(_tmp36, 1)[:, None]
    tl.store(out_ptr3 + x3, tmp36, xmask)
    tmp45 = tl.sum(_tmp45, 1)[:, None]
    tl.store(out_ptr4 + x3, tmp45, xmask)
    tmp54 = tl.sum(_tmp54, 1)[:, None]
    tl.store(out_ptr5 + x3, tmp54, xmask)
    tmp63 = tl.sum(_tmp63, 1)[:, None]
    tl.store(out_ptr6 + x3, tmp63, xmask)


@triton.jit
def triton_per_fused_copy_linalg_vector_norm_zeros_6(in_out_ptr0,
    in_out_ptr1, 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, in_ptr32, in_ptr33,
    in_ptr34, in_ptr35, in_ptr36, in_ptr37, in_ptr38, in_ptr39, in_ptr40,
    in_ptr41, in_ptr42, in_ptr43, in_ptr44, in_ptr45, in_ptr46, in_ptr47,
    in_ptr48, in_ptr49, in_ptr50, in_ptr51, in_ptr52, in_ptr53, in_ptr54,
    in_ptr55, in_ptr56, in_ptr57, in_ptr58, in_ptr59, in_ptr60, in_ptr61,
    in_ptr62, in_ptr63, xnumel, rnumel, XBLOCK: tl.constexpr):
    xnumel = 256
    RBLOCK: tl.constexpr = 128
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:, None]
    xmask = xindex < xnumel
    rindex = tl.arange(0, RBLOCK)[None, :]
    tl.full([XBLOCK, RBLOCK], True, tl.int1)
    x0 = xindex % 64
    r2 = rindex
    x1 = xindex // 64
    x3 = xindex
    tmp0 = x0
    tmp1 = tl.full([1, 1], 4, tl.int64)
    tmp2 = tmp0 >= tmp1
    tmp3 = tl.full([1, 1], 5, tl.int64)
    tmp4 = tmp0 < tmp3
    tmp5 = tmp2 & tmp4
    tmp6 = tl.load(in_ptr0 + (r2 + 128 * x1), tmp5 & xmask, eviction_policy
        ='evict_last', other=0.0)
    tmp7 = tl.full([1, 1], 3, tl.int64)
    tmp8 = tmp0 >= tmp7
    tmp9 = tmp0 < tmp1
    tmp10 = tmp8 & tmp9
    tmp11 = tl.load(in_ptr1 + (r2 + 128 * x1), tmp10 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp12 = tl.full([1, 1], 2, tl.int64)
    tmp13 = tmp0 >= tmp12
    tmp14 = tmp0 < tmp7
    tmp15 = tmp13 & tmp14
    tmp16 = tl.load(in_ptr2 + (r2 + 128 * x1), tmp15 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp17 = tl.full([1, 1], 1, tl.int64)
    tmp18 = tmp0 >= tmp17
    tmp19 = tmp0 < tmp12
    tmp20 = tmp18 & tmp19
    tmp21 = tl.load(in_ptr3 + (r2 + 128 * x1), tmp20 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp22 = tmp0 < tmp17
    tmp23 = tl.load(in_ptr4 + (r2 + 128 * x1), tmp22 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp24 = 0.0
    tmp25 = tl.where(tmp22, tmp23, tmp24)
    tmp26 = tl.where(tmp20, tmp21, tmp25)
    tmp27 = tl.where(tmp15, tmp16, tmp26)
    tmp28 = tl.where(tmp10, tmp11, tmp27)
    tmp29 = tl.where(tmp5, tmp6, tmp28)
    tmp30 = tl.full([1, 1], 8, tl.int64)
    tmp31 = tmp0 >= tmp30
    tmp32 = tl.full([1, 1], 9, tl.int64)
    tmp33 = tmp0 < tmp32
    tmp34 = tmp31 & tmp33
    tmp35 = tl.load(in_ptr5 + (r2 + 128 * x1), tmp34 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp36 = tl.full([1, 1], 7, tl.int64)
    tmp37 = tmp0 >= tmp36
    tmp38 = tmp0 < tmp30
    tmp39 = tmp37 & tmp38
    tmp40 = tl.load(in_ptr6 + (r2 + 128 * x1), tmp39 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp41 = tl.full([1, 1], 6, tl.int64)
    tmp42 = tmp0 >= tmp41
    tmp43 = tmp0 < tmp36
    tmp44 = tmp42 & tmp43
    tmp45 = tl.load(in_ptr7 + (r2 + 128 * x1), tmp44 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp46 = tmp0 >= tmp3
    tmp47 = tmp0 < tmp41
    tmp48 = tmp46 & tmp47
    tmp49 = tl.load(in_ptr8 + (r2 + 128 * x1), tmp48 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp50 = tl.where(tmp48, tmp49, tmp29)
    tmp51 = tl.where(tmp44, tmp45, tmp50)
    tmp52 = tl.where(tmp39, tmp40, tmp51)
    tmp53 = tl.where(tmp34, tmp35, tmp52)
    tmp54 = tl.full([1, 1], 12, tl.int64)
    tmp55 = tmp0 >= tmp54
    tmp56 = tl.full([1, 1], 13, tl.int64)
    tmp57 = tmp0 < tmp56
    tmp58 = tmp55 & tmp57
    tmp59 = tl.load(in_ptr9 + (r2 + 128 * x1), tmp58 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp60 = tl.full([1, 1], 11, tl.int64)
    tmp61 = tmp0 >= tmp60
    tmp62 = tmp0 < tmp54
    tmp63 = tmp61 & tmp62
    tmp64 = tl.load(in_ptr10 + (r2 + 128 * x1), tmp63 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp65 = tl.full([1, 1], 10, tl.int64)
    tmp66 = tmp0 >= tmp65
    tmp67 = tmp0 < tmp60
    tmp68 = tmp66 & tmp67
    tmp69 = tl.load(in_ptr11 + (r2 + 128 * x1), tmp68 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp70 = tmp0 >= tmp32
    tmp71 = tmp0 < tmp65
    tmp72 = tmp70 & tmp71
    tmp73 = tl.load(in_ptr12 + (r2 + 128 * x1), tmp72 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp74 = tl.where(tmp72, tmp73, tmp53)
    tmp75 = tl.where(tmp68, tmp69, tmp74)
    tmp76 = tl.where(tmp63, tmp64, tmp75)
    tmp77 = tl.where(tmp58, tmp59, tmp76)
    tmp78 = tl.full([1, 1], 16, tl.int64)
    tmp79 = tmp0 >= tmp78
    tmp80 = tl.full([1, 1], 17, tl.int64)
    tmp81 = tmp0 < tmp80
    tmp82 = tmp79 & tmp81
    tmp83 = tl.load(in_ptr13 + (r2 + 128 * x1), tmp82 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp84 = tl.full([1, 1], 15, tl.int64)
    tmp85 = tmp0 >= tmp84
    tmp86 = tmp0 < tmp78
    tmp87 = tmp85 & tmp86
    tmp88 = tl.load(in_ptr14 + (r2 + 128 * x1), tmp87 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp89 = tl.full([1, 1], 14, tl.int64)
    tmp90 = tmp0 >= tmp89
    tmp91 = tmp0 < tmp84
    tmp92 = tmp90 & tmp91
    tmp93 = tl.load(in_ptr15 + (r2 + 128 * x1), tmp92 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp94 = tmp0 >= tmp56
    tmp95 = tmp0 < tmp89
    tmp96 = tmp94 & tmp95
    tmp97 = tl.load(in_ptr16 + (r2 + 128 * x1), tmp96 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp98 = tl.where(tmp96, tmp97, tmp77)
    tmp99 = tl.where(tmp92, tmp93, tmp98)
    tmp100 = tl.where(tmp87, tmp88, tmp99)
    tmp101 = tl.where(tmp82, tmp83, tmp100)
    tmp102 = tl.full([1, 1], 20, tl.int64)
    tmp103 = tmp0 >= tmp102
    tmp104 = tl.full([1, 1], 21, tl.int64)
    tmp105 = tmp0 < tmp104
    tmp106 = tmp103 & tmp105
    tmp107 = tl.load(in_ptr17 + (r2 + 128 * x1), tmp106 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp108 = tl.full([1, 1], 19, tl.int64)
    tmp109 = tmp0 >= tmp108
    tmp110 = tmp0 < tmp102
    tmp111 = tmp109 & tmp110
    tmp112 = tl.load(in_ptr18 + (r2 + 128 * x1), tmp111 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp113 = tl.full([1, 1], 18, tl.int64)
    tmp114 = tmp0 >= tmp113
    tmp115 = tmp0 < tmp108
    tmp116 = tmp114 & tmp115
    tmp117 = tl.load(in_ptr19 + (r2 + 128 * x1), tmp116 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp118 = tmp0 >= tmp80
    tmp119 = tmp0 < tmp113
    tmp120 = tmp118 & tmp119
    tmp121 = tl.load(in_ptr20 + (r2 + 128 * x1), tmp120 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp122 = tl.where(tmp120, tmp121, tmp101)
    tmp123 = tl.where(tmp116, tmp117, tmp122)
    tmp124 = tl.where(tmp111, tmp112, tmp123)
    tmp125 = tl.where(tmp106, tmp107, tmp124)
    tmp126 = tl.full([1, 1], 24, tl.int64)
    tmp127 = tmp0 >= tmp126
    tmp128 = tl.full([1, 1], 25, tl.int64)
    tmp129 = tmp0 < tmp128
    tmp130 = tmp127 & tmp129
    tmp131 = tl.load(in_ptr21 + (r2 + 128 * x1), tmp130 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp132 = tl.full([1, 1], 23, tl.int64)
    tmp133 = tmp0 >= tmp132
    tmp134 = tmp0 < tmp126
    tmp135 = tmp133 & tmp134
    tmp136 = tl.load(in_ptr22 + (r2 + 128 * x1), tmp135 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp137 = tl.full([1, 1], 22, tl.int64)
    tmp138 = tmp0 >= tmp137
    tmp139 = tmp0 < tmp132
    tmp140 = tmp138 & tmp139
    tmp141 = tl.load(in_ptr23 + (r2 + 128 * x1), tmp140 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp142 = tmp0 >= tmp104
    tmp143 = tmp0 < tmp137
    tmp144 = tmp142 & tmp143
    tmp145 = tl.load(in_ptr24 + (r2 + 128 * x1), tmp144 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp146 = tl.where(tmp144, tmp145, tmp125)
    tmp147 = tl.where(tmp140, tmp141, tmp146)
    tmp148 = tl.where(tmp135, tmp136, tmp147)
    tmp149 = tl.where(tmp130, tmp131, tmp148)
    tmp150 = tl.full([1, 1], 28, tl.int64)
    tmp151 = tmp0 >= tmp150
    tmp152 = tl.full([1, 1], 29, tl.int64)
    tmp153 = tmp0 < tmp152
    tmp154 = tmp151 & tmp153
    tmp155 = tl.load(in_ptr25 + (r2 + 128 * x1), tmp154 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp156 = tl.full([1, 1], 27, tl.int64)
    tmp157 = tmp0 >= tmp156
    tmp158 = tmp0 < tmp150
    tmp159 = tmp157 & tmp158
    tmp160 = tl.load(in_ptr26 + (r2 + 128 * x1), tmp159 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp161 = tl.full([1, 1], 26, tl.int64)
    tmp162 = tmp0 >= tmp161
    tmp163 = tmp0 < tmp156
    tmp164 = tmp162 & tmp163
    tmp165 = tl.load(in_ptr27 + (r2 + 128 * x1), tmp164 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp166 = tmp0 >= tmp128
    tmp167 = tmp0 < tmp161
    tmp168 = tmp166 & tmp167
    tmp169 = tl.load(in_ptr28 + (r2 + 128 * x1), tmp168 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp170 = tl.where(tmp168, tmp169, tmp149)
    tmp171 = tl.where(tmp164, tmp165, tmp170)
    tmp172 = tl.where(tmp159, tmp160, tmp171)
    tmp173 = tl.where(tmp154, tmp155, tmp172)
    tmp174 = tl.full([1, 1], 32, tl.int64)
    tmp175 = tmp0 >= tmp174
    tmp176 = tl.full([1, 1], 33, tl.int64)
    tmp177 = tmp0 < tmp176
    tmp178 = tmp175 & tmp177
    tmp179 = tl.load(in_ptr29 + (r2 + 128 * x1), tmp178 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp180 = tl.full([1, 1], 31, tl.int64)
    tmp181 = tmp0 >= tmp180
    tmp182 = tmp0 < tmp174
    tmp183 = tmp181 & tmp182
    tmp184 = tl.load(in_ptr30 + (r2 + 128 * x1), tmp183 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp185 = tl.full([1, 1], 30, tl.int64)
    tmp186 = tmp0 >= tmp185
    tmp187 = tmp0 < tmp180
    tmp188 = tmp186 & tmp187
    tmp189 = tl.load(in_ptr31 + (r2 + 128 * x1), tmp188 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp190 = tmp0 >= tmp152
    tmp191 = tmp0 < tmp185
    tmp192 = tmp190 & tmp191
    tmp193 = tl.load(in_ptr32 + (r2 + 128 * x1), tmp192 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp194 = tl.where(tmp192, tmp193, tmp173)
    tmp195 = tl.where(tmp188, tmp189, tmp194)
    tmp196 = tl.where(tmp183, tmp184, tmp195)
    tmp197 = tl.where(tmp178, tmp179, tmp196)
    tmp198 = tl.full([1, 1], 36, tl.int64)
    tmp199 = tmp0 >= tmp198
    tmp200 = tl.full([1, 1], 37, tl.int64)
    tmp201 = tmp0 < tmp200
    tmp202 = tmp199 & tmp201
    tmp203 = tl.load(in_ptr33 + (r2 + 128 * x1), tmp202 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp204 = tl.full([1, 1], 35, tl.int64)
    tmp205 = tmp0 >= tmp204
    tmp206 = tmp0 < tmp198
    tmp207 = tmp205 & tmp206
    tmp208 = tl.load(in_ptr34 + (r2 + 128 * x1), tmp207 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp209 = tl.full([1, 1], 34, tl.int64)
    tmp210 = tmp0 >= tmp209
    tmp211 = tmp0 < tmp204
    tmp212 = tmp210 & tmp211
    tmp213 = tl.load(in_ptr35 + (r2 + 128 * x1), tmp212 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp214 = tmp0 >= tmp176
    tmp215 = tmp0 < tmp209
    tmp216 = tmp214 & tmp215
    tmp217 = tl.load(in_ptr36 + (r2 + 128 * x1), tmp216 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp218 = tl.where(tmp216, tmp217, tmp197)
    tmp219 = tl.where(tmp212, tmp213, tmp218)
    tmp220 = tl.where(tmp207, tmp208, tmp219)
    tmp221 = tl.where(tmp202, tmp203, tmp220)
    tmp222 = tl.full([1, 1], 40, tl.int64)
    tmp223 = tmp0 >= tmp222
    tmp224 = tl.full([1, 1], 41, tl.int64)
    tmp225 = tmp0 < tmp224
    tmp226 = tmp223 & tmp225
    tmp227 = tl.load(in_ptr37 + (r2 + 128 * x1), tmp226 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp228 = tl.full([1, 1], 39, tl.int64)
    tmp229 = tmp0 >= tmp228
    tmp230 = tmp0 < tmp222
    tmp231 = tmp229 & tmp230
    tmp232 = tl.load(in_ptr38 + (r2 + 128 * x1), tmp231 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp233 = tl.full([1, 1], 38, tl.int64)
    tmp234 = tmp0 >= tmp233
    tmp235 = tmp0 < tmp228
    tmp236 = tmp234 & tmp235
    tmp237 = tl.load(in_ptr39 + (r2 + 128 * x1), tmp236 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp238 = tmp0 >= tmp200
    tmp239 = tmp0 < tmp233
    tmp240 = tmp238 & tmp239
    tmp241 = tl.load(in_ptr40 + (r2 + 128 * x1), tmp240 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp242 = tl.where(tmp240, tmp241, tmp221)
    tmp243 = tl.where(tmp236, tmp237, tmp242)
    tmp244 = tl.where(tmp231, tmp232, tmp243)
    tmp245 = tl.where(tmp226, tmp227, tmp244)
    tmp246 = tl.full([1, 1], 44, tl.int64)
    tmp247 = tmp0 >= tmp246
    tmp248 = tl.full([1, 1], 45, tl.int64)
    tmp249 = tmp0 < tmp248
    tmp250 = tmp247 & tmp249
    tmp251 = tl.load(in_ptr41 + (r2 + 128 * x1), tmp250 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp252 = tl.full([1, 1], 43, tl.int64)
    tmp253 = tmp0 >= tmp252
    tmp254 = tmp0 < tmp246
    tmp255 = tmp253 & tmp254
    tmp256 = tl.load(in_ptr42 + (r2 + 128 * x1), tmp255 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp257 = tl.full([1, 1], 42, tl.int64)
    tmp258 = tmp0 >= tmp257
    tmp259 = tmp0 < tmp252
    tmp260 = tmp258 & tmp259
    tmp261 = tl.load(in_ptr43 + (r2 + 128 * x1), tmp260 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp262 = tmp0 >= tmp224
    tmp263 = tmp0 < tmp257
    tmp264 = tmp262 & tmp263
    tmp265 = tl.load(in_ptr44 + (r2 + 128 * x1), tmp264 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp266 = tl.where(tmp264, tmp265, tmp245)
    tmp267 = tl.where(tmp260, tmp261, tmp266)
    tmp268 = tl.where(tmp255, tmp256, tmp267)
    tmp269 = tl.where(tmp250, tmp251, tmp268)
    tmp270 = tl.full([1, 1], 48, tl.int64)
    tmp271 = tmp0 >= tmp270
    tmp272 = tl.full([1, 1], 49, tl.int64)
    tmp273 = tmp0 < tmp272
    tmp274 = tmp271 & tmp273
    tmp275 = tl.load(in_ptr45 + (r2 + 128 * x1), tmp274 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp276 = tl.full([1, 1], 47, tl.int64)
    tmp277 = tmp0 >= tmp276
    tmp278 = tmp0 < tmp270
    tmp279 = tmp277 & tmp278
    tmp280 = tl.load(in_ptr46 + (r2 + 128 * x1), tmp279 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp281 = tl.full([1, 1], 46, tl.int64)
    tmp282 = tmp0 >= tmp281
    tmp283 = tmp0 < tmp276
    tmp284 = tmp282 & tmp283
    tmp285 = tl.load(in_ptr47 + (r2 + 128 * x1), tmp284 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp286 = tmp0 >= tmp248
    tmp287 = tmp0 < tmp281
    tmp288 = tmp286 & tmp287
    tmp289 = tl.load(in_ptr48 + (r2 + 128 * x1), tmp288 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp290 = tl.where(tmp288, tmp289, tmp269)
    tmp291 = tl.where(tmp284, tmp285, tmp290)
    tmp292 = tl.where(tmp279, tmp280, tmp291)
    tmp293 = tl.where(tmp274, tmp275, tmp292)
    tmp294 = tl.full([1, 1], 52, tl.int64)
    tmp295 = tmp0 >= tmp294
    tmp296 = tl.full([1, 1], 53, tl.int64)
    tmp297 = tmp0 < tmp296
    tmp298 = tmp295 & tmp297
    tmp299 = tl.load(in_ptr49 + (r2 + 128 * x1), tmp298 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp300 = tl.full([1, 1], 51, tl.int64)
    tmp301 = tmp0 >= tmp300
    tmp302 = tmp0 < tmp294
    tmp303 = tmp301 & tmp302
    tmp304 = tl.load(in_ptr50 + (r2 + 128 * x1), tmp303 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp305 = tl.full([1, 1], 50, tl.int64)
    tmp306 = tmp0 >= tmp305
    tmp307 = tmp0 < tmp300
    tmp308 = tmp306 & tmp307
    tmp309 = tl.load(in_ptr51 + (r2 + 128 * x1), tmp308 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp310 = tmp0 >= tmp272
    tmp311 = tmp0 < tmp305
    tmp312 = tmp310 & tmp311
    tmp313 = tl.load(in_ptr52 + (r2 + 128 * x1), tmp312 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp314 = tl.where(tmp312, tmp313, tmp293)
    tmp315 = tl.where(tmp308, tmp309, tmp314)
    tmp316 = tl.where(tmp303, tmp304, tmp315)
    tmp317 = tl.where(tmp298, tmp299, tmp316)
    tmp318 = tl.full([1, 1], 56, tl.int64)
    tmp319 = tmp0 >= tmp318
    tmp320 = tl.full([1, 1], 57, tl.int64)
    tmp321 = tmp0 < tmp320
    tmp322 = tmp319 & tmp321
    tmp323 = tl.load(in_ptr53 + (r2 + 128 * x1), tmp322 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp324 = tl.full([1, 1], 55, tl.int64)
    tmp325 = tmp0 >= tmp324
    tmp326 = tmp0 < tmp318
    tmp327 = tmp325 & tmp326
    tmp328 = tl.load(in_ptr54 + (r2 + 128 * x1), tmp327 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp329 = tl.full([1, 1], 54, tl.int64)
    tmp330 = tmp0 >= tmp329
    tmp331 = tmp0 < tmp324
    tmp332 = tmp330 & tmp331
    tmp333 = tl.load(in_ptr55 + (r2 + 128 * x1), tmp332 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp334 = tmp0 >= tmp296
    tmp335 = tmp0 < tmp329
    tmp336 = tmp334 & tmp335
    tmp337 = tl.load(in_ptr56 + (r2 + 128 * x1), tmp336 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp338 = tl.where(tmp336, tmp337, tmp317)
    tmp339 = tl.where(tmp332, tmp333, tmp338)
    tmp340 = tl.where(tmp327, tmp328, tmp339)
    tmp341 = tl.where(tmp322, tmp323, tmp340)
    tmp342 = tl.full([1, 1], 60, tl.int64)
    tmp343 = tmp0 >= tmp342
    tmp344 = tl.full([1, 1], 61, tl.int64)
    tmp345 = tmp0 < tmp344
    tmp346 = tmp343 & tmp345
    tmp347 = tl.load(in_ptr57 + (r2 + 128 * x1), tmp346 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp348 = tl.full([1, 1], 59, tl.int64)
    tmp349 = tmp0 >= tmp348
    tmp350 = tmp0 < tmp342
    tmp351 = tmp349 & tmp350
    tmp352 = tl.load(in_ptr58 + (r2 + 128 * x1), tmp351 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp353 = tl.full([1, 1], 58, tl.int64)
    tmp354 = tmp0 >= tmp353
    tmp355 = tmp0 < tmp348
    tmp356 = tmp354 & tmp355
    tmp357 = tl.load(in_ptr59 + (r2 + 128 * x1), tmp356 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp358 = tmp0 >= tmp320
    tmp359 = tmp0 < tmp353
    tmp360 = tmp358 & tmp359
    tmp361 = tl.load(in_ptr60 + (r2 + 128 * x1), tmp360 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp362 = tl.where(tmp360, tmp361, tmp341)
    tmp363 = tl.where(tmp356, tmp357, tmp362)
    tmp364 = tl.where(tmp351, tmp352, tmp363)
    tmp365 = tl.where(tmp346, tmp347, tmp364)
    tmp366 = tl.full([1, 1], 63, tl.int64)
    tmp367 = tmp0 >= tmp366
    tmp368 = tl.load(in_ptr61 + (r2 + 128 * x1), tmp367 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp369 = tl.full([1, 1], 62, tl.int64)
    tmp370 = tmp0 >= tmp369
    tmp371 = tmp0 < tmp366
    tmp372 = tmp370 & tmp371
    tmp373 = tl.load(in_ptr62 + (r2 + 128 * x1), tmp372 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp374 = tmp0 >= tmp344
    tmp375 = tmp0 < tmp369
    tmp376 = tmp374 & tmp375
    tmp377 = tl.load(in_ptr63 + (r2 + 128 * x1), tmp376 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp378 = tl.where(tmp376, tmp377, tmp365)
    tmp379 = tl.where(tmp372, tmp373, tmp378)
    tmp380 = tl.where(tmp367, tmp368, tmp379)
    tmp381 = tmp380 * tmp380
    tmp382 = tl.broadcast_to(tmp381, [XBLOCK, RBLOCK])
    tmp384 = tl.where(xmask, tmp382, 0)
    tmp385 = tl.sum(tmp384, 1)[:, None]
    tmp386 = libdevice.sqrt(tmp385)
    tl.store(in_out_ptr0 + (r2 + 128 * x3), tmp380, xmask)
    tl.debug_barrier()
    tl.store(in_out_ptr1 + x3, tmp386, xmask)


@triton.jit
def triton_red_fused_div_linalg_vector_norm_7(in_out_ptr0, in_ptr0, in_ptr1,
    out_ptr0, xnumel, rnumel, XBLOCK: tl.constexpr, RBLOCK: tl.constexpr):
    xnumel = 4
    rnumel = 8192
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:, None]
    xmask = xindex < xnumel
    rbase = tl.arange(0, RBLOCK)[None, :]
    x0 = xindex
    _tmp7 = tl.full([XBLOCK, RBLOCK], 0, tl.float32)
    for roffset in range(0, rnumel, RBLOCK):
        rindex = roffset + rbase
        rmask = rindex < rnumel
        r1 = rindex
        tmp0 = tl.load(in_ptr0 + (r1 + 8192 * x0), rmask & xmask,
            eviction_policy='evict_last', other=0.0)
        tmp1 = tl.load(in_ptr1 + (64 * x0 + r1 // 128), rmask & xmask,
            eviction_policy='evict_last', other=0.0)
        tmp2 = 1e-12
        tmp3 = triton_helpers.maximum(tmp1, tmp2)
        tmp4 = tmp0 / tmp3
        tmp5 = tmp4 * tmp4
        tmp6 = tl.broadcast_to(tmp5, [XBLOCK, RBLOCK])
        tmp8 = _tmp7 + tmp6
        _tmp7 = tl.where(rmask & xmask, tmp8, _tmp7)
    tmp7 = tl.sum(_tmp7, 1)[:, None]
    tmp9 = libdevice.sqrt(tmp7)
    tl.debug_barrier()
    tl.store(in_out_ptr0 + x0, tmp9, xmask)
    for roffset in range(0, rnumel, RBLOCK):
        rindex = roffset + rbase
        rmask = rindex < rnumel
        r1 = rindex
        tmp10 = tl.load(in_ptr0 + (r1 + 8192 * x0), rmask & xmask,
            eviction_policy='evict_first', other=0.0)
        tmp11 = tl.load(in_ptr1 + (64 * x0 + r1 // 128), rmask & xmask,
            eviction_policy='evict_last', other=0.0)
        tmp12 = 1e-12
        tmp13 = triton_helpers.maximum(tmp11, tmp12)
        tmp14 = tmp10 / tmp13
        tmp15 = triton_helpers.maximum(tmp9, tmp12)
        tmp16 = tmp14 / tmp15
        tl.store(out_ptr0 + (r1 + 8192 * x0), tmp16, rmask & xmask)


def call(args):
    primals_1, primals_2, primals_3 = args
    args.clear()
    assert_size_stride(primals_1, (4, 128, 64, 64), (524288, 4096, 64, 1))
    assert_size_stride(primals_2, (64, 128, 1, 1), (128, 1, 1, 1))
    assert_size_stride(primals_3, (64, 128), (128, 1))
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        buf0 = empty_strided_cuda((4, 1, 64, 64), (4096, 16384, 64, 1),
            torch.float32)
        get_raw_stream(0)
        triton_red_fused_linalg_vector_norm_0[grid(16384)](primals_1, buf0,
            16384, 128, XBLOCK=64, RBLOCK=4, num_warps=8, num_stages=1)
        buf1 = empty_strided_cuda((4, 128, 64, 64), (524288, 4096, 64, 1),
            torch.float32)
        buf6 = empty_strided_cuda((4, 1, 128, 4096), (524288, 2097152, 4096,
            1), torch.float32)
        buf8 = empty_strided_cuda((4, 1, 128, 4096), (524288, 2097152, 4096,
            1), torch.float32)
        buf10 = empty_strided_cuda((4, 1, 128, 4096), (524288, 2097152, 
            4096, 1), torch.float32)
        buf12 = empty_strided_cuda((4, 1, 128, 4096), (524288, 2097152, 
            4096, 1), torch.float32)
        buf15 = empty_strided_cuda((4, 1, 128, 4096), (524288, 2097152, 
            4096, 1), torch.float32)
        buf17 = empty_strided_cuda((4, 1, 128, 4096), (524288, 2097152, 
            4096, 1), torch.float32)
        buf19 = empty_strided_cuda((4, 1, 128, 4096), (524288, 2097152, 
            4096, 1), torch.float32)
        buf21 = empty_strided_cuda((4, 1, 128, 4096), (524288, 2097152, 
            4096, 1), torch.float32)
        buf24 = empty_strided_cuda((4, 1, 128, 4096), (524288, 2097152, 
            4096, 1), torch.float32)
        buf26 = empty_strided_cuda((4, 1, 128, 4096), (524288, 2097152, 
            4096, 1), torch.float32)
        buf28 = empty_strided_cuda((4, 1, 128, 4096), (524288, 2097152, 
            4096, 1), torch.float32)
        buf30 = empty_strided_cuda((4, 1, 128, 4096), (524288, 2097152, 
            4096, 1), torch.float32)
        buf33 = empty_strided_cuda((4, 1, 128, 4096), (524288, 2097152, 
            4096, 1), torch.float32)
        buf35 = empty_strided_cuda((4, 1, 128, 4096), (524288, 2097152, 
            4096, 1), torch.float32)
        buf37 = empty_strided_cuda((4, 1, 128, 4096), (524288, 2097152, 
            4096, 1), torch.float32)
        buf39 = empty_strided_cuda((4, 1, 128, 4096), (524288, 2097152, 
            4096, 1), torch.float32)
        buf42 = empty_strided_cuda((4, 1, 128, 4096), (524288, 2097152, 
            4096, 1), torch.float32)
        buf44 = empty_strided_cuda((4, 1, 128, 4096), (524288, 2097152, 
            4096, 1), torch.float32)
        buf46 = empty_strided_cuda((4, 1, 128, 4096), (524288, 2097152, 
            4096, 1), torch.float32)
        buf48 = empty_strided_cuda((4, 1, 128, 4096), (524288, 2097152, 
            4096, 1), torch.float32)
        buf51 = empty_strided_cuda((4, 1, 128, 4096), (524288, 2097152, 
            4096, 1), torch.float32)
        buf53 = empty_strided_cuda((4, 1, 128, 4096), (524288, 2097152, 
            4096, 1), torch.float32)
        buf55 = empty_strided_cuda((4, 1, 128, 4096), (524288, 2097152, 
            4096, 1), torch.float32)
        buf57 = empty_strided_cuda((4, 1, 128, 4096), (524288, 2097152, 
            4096, 1), torch.float32)
        buf60 = empty_strided_cuda((4, 1, 128, 4096), (524288, 2097152, 
            4096, 1), torch.float32)
        buf62 = empty_strided_cuda((4, 1, 128, 4096), (524288, 2097152, 
            4096, 1), torch.float32)
        buf64 = empty_strided_cuda((4, 1, 128, 4096), (524288, 2097152, 
            4096, 1), torch.float32)
        buf66 = empty_strided_cuda((4, 1, 128, 4096), (524288, 2097152, 
            4096, 1), torch.float32)
        buf69 = empty_strided_cuda((4, 1, 128, 4096), (524288, 2097152, 
            4096, 1), torch.float32)
        buf71 = empty_strided_cuda((4, 1, 128, 4096), (524288, 2097152, 
            4096, 1), torch.float32)
        buf73 = empty_strided_cuda((4, 1, 128, 4096), (524288, 2097152, 
            4096, 1), torch.float32)
        buf75 = empty_strided_cuda((4, 1, 128, 4096), (524288, 2097152, 
            4096, 1), torch.float32)
        buf78 = empty_strided_cuda((4, 1, 128, 4096), (524288, 2097152, 
            4096, 1), torch.float32)
        buf80 = empty_strided_cuda((4, 1, 128, 4096), (524288, 2097152, 
            4096, 1), torch.float32)
        buf82 = empty_strided_cuda((4, 1, 128, 4096), (524288, 2097152, 
            4096, 1), torch.float32)
        buf84 = empty_strided_cuda((4, 1, 128, 4096), (524288, 2097152, 
            4096, 1), torch.float32)
        buf87 = empty_strided_cuda((4, 1, 128, 4096), (524288, 2097152, 
            4096, 1), torch.float32)
        buf89 = empty_strided_cuda((4, 1, 128, 4096), (524288, 2097152, 
            4096, 1), torch.float32)
        buf91 = empty_strided_cuda((4, 1, 128, 4096), (524288, 2097152, 
            4096, 1), torch.float32)
        buf93 = empty_strided_cuda((4, 1, 128, 4096), (524288, 2097152, 
            4096, 1), torch.float32)
        buf96 = empty_strided_cuda((4, 1, 128, 4096), (524288, 2097152, 
            4096, 1), torch.float32)
        buf98 = empty_strided_cuda((4, 1, 128, 4096), (524288, 2097152, 
            4096, 1), torch.float32)
        buf100 = empty_strided_cuda((4, 1, 128, 4096), (524288, 2097152, 
            4096, 1), torch.float32)
        buf102 = empty_strided_cuda((4, 1, 128, 4096), (524288, 2097152, 
            4096, 1), torch.float32)
        buf105 = empty_strided_cuda((4, 1, 128, 4096), (524288, 2097152, 
            4096, 1), torch.float32)
        buf107 = empty_strided_cuda((4, 1, 128, 4096), (524288, 2097152, 
            4096, 1), torch.float32)
        buf109 = empty_strided_cuda((4, 1, 128, 4096), (524288, 2097152, 
            4096, 1), torch.float32)
        buf111 = empty_strided_cuda((4, 1, 128, 4096), (524288, 2097152, 
            4096, 1), torch.float32)
        buf114 = empty_strided_cuda((4, 1, 128, 4096), (524288, 2097152, 
            4096, 1), torch.float32)
        buf116 = empty_strided_cuda((4, 1, 128, 4096), (524288, 2097152, 
            4096, 1), torch.float32)
        buf118 = empty_strided_cuda((4, 1, 128, 4096), (524288, 2097152, 
            4096, 1), torch.float32)
        buf120 = empty_strided_cuda((4, 1, 128, 4096), (524288, 2097152, 
            4096, 1), torch.float32)
        buf123 = empty_strided_cuda((4, 1, 128, 4096), (524288, 2097152, 
            4096, 1), torch.float32)
        buf125 = empty_strided_cuda((4, 1, 128, 4096), (524288, 2097152, 
            4096, 1), torch.float32)
        buf127 = empty_strided_cuda((4, 1, 128, 4096), (524288, 2097152, 
            4096, 1), torch.float32)
        buf129 = empty_strided_cuda((4, 1, 128, 4096), (524288, 2097152, 
            4096, 1), torch.float32)
        buf132 = empty_strided_cuda((4, 1, 128, 4096), (524288, 2097152, 
            4096, 1), torch.float32)
        buf134 = empty_strided_cuda((4, 1, 128, 4096), (524288, 2097152, 
            4096, 1), torch.float32)
        buf136 = empty_strided_cuda((4, 1, 128, 4096), (524288, 2097152, 
            4096, 1), torch.float32)
        buf138 = empty_strided_cuda((4, 1, 128, 4096), (524288, 2097152, 
            4096, 1), torch.float32)
        buf141 = empty_strided_cuda((4, 1, 128, 4096), (524288, 2097152, 
            4096, 1), torch.float32)
        buf143 = empty_strided_cuda((4, 1, 128, 4096), (524288, 2097152, 
            4096, 1), torch.float32)
        buf145 = empty_strided_cuda((4, 1, 128, 4096), (524288, 2097152, 
            4096, 1), torch.float32)
        triton_poi_fused_div_sub_1[grid(2097152)](primals_1, buf0,
            primals_3, buf1, buf6, buf8, buf10, buf12, buf15, buf17, buf19,
            buf21, buf24, buf26, buf28, buf30, buf33, buf35, buf37, buf39,
            buf42, buf44, buf46, buf48, buf51, buf53, buf55, buf57, buf60,
            buf62, buf64, buf66, buf69, buf71, buf73, buf75, buf78, buf80,
            buf82, buf84, buf87, buf89, buf91, buf93, buf96, buf98, buf100,
            buf102, buf105, buf107, buf109, buf111, buf114, buf116, buf118,
            buf120, buf123, buf125, buf127, buf129, buf132, buf134, buf136,
            buf138, buf141, buf143, buf145, 2097152, XBLOCK=512, num_warps=
            8, num_stages=1)
        del primals_1
        buf2 = extern_kernels.convolution(buf1, primals_2, stride=(1, 1),
            padding=(0, 0), dilation=(1, 1), transposed=False,
            output_padding=(0, 0), groups=1, bias=None)
        assert_size_stride(buf2, (4, 64, 64, 64), (262144, 4096, 64, 1))
        buf3 = reinterpret_tensor(buf0, (4, 1, 4096), (4096, 4096, 1), 0)
        del buf0
        buf4 = empty_strided_cuda((4, 1, 4096), (4096, 4096, 1), torch.float32)
        triton_per_fused__softmax_2[grid(16384)](buf2, buf3, buf4, 16384, 
            64, XBLOCK=8, num_warps=4, num_stages=1)
        buf5 = empty_strided_cuda((4, 1, 128), (128, 512, 1), torch.float32)
        buf7 = empty_strided_cuda((4, 1, 128), (128, 512, 1), torch.float32)
        buf9 = empty_strided_cuda((4, 1, 128), (128, 512, 1), torch.float32)
        buf11 = empty_strided_cuda((4, 1, 128), (128, 512, 1), torch.float32)
        buf13 = empty_strided_cuda((4, 1, 128), (128, 512, 1), torch.float32)
        buf16 = empty_strided_cuda((4, 1, 128), (128, 512, 1), torch.float32)
        buf18 = empty_strided_cuda((4, 1, 128), (128, 512, 1), torch.float32)
        buf20 = empty_strided_cuda((4, 1, 128), (128, 512, 1), torch.float32)
        buf22 = empty_strided_cuda((4, 1, 128), (128, 512, 1), torch.float32)
        buf25 = empty_strided_cuda((4, 1, 128), (128, 512, 1), torch.float32)
        buf27 = empty_strided_cuda((4, 1, 128), (128, 512, 1), torch.float32)
        buf29 = empty_strided_cuda((4, 1, 128), (128, 512, 1), torch.float32)
        buf31 = empty_strided_cuda((4, 1, 128), (128, 512, 1), torch.float32)
        buf34 = empty_strided_cuda((4, 1, 128), (128, 512, 1), torch.float32)
        buf36 = empty_strided_cuda((4, 1, 128), (128, 512, 1), torch.float32)
        buf38 = empty_strided_cuda((4, 1, 128), (128, 512, 1), torch.float32)
        buf40 = empty_strided_cuda((4, 1, 128), (128, 512, 1), torch.float32)
        buf43 = empty_strided_cuda((4, 1, 128), (128, 512, 1), torch.float32)
        buf45 = empty_strided_cuda((4, 1, 128), (128, 512, 1), torch.float32)
        buf47 = empty_strided_cuda((4, 1, 128), (128, 512, 1), torch.float32)
        buf49 = empty_strided_cuda((4, 1, 128), (128, 512, 1), torch.float32)
        buf52 = empty_strided_cuda((4, 1, 128), (128, 512, 1), torch.float32)
        buf54 = empty_strided_cuda((4, 1, 128), (128, 512, 1), torch.float32)
        buf56 = empty_strided_cuda((4, 1, 128), (128, 512, 1), torch.float32)
        buf58 = empty_strided_cuda((4, 1, 128), (128, 512, 1), torch.float32)
        buf61 = empty_strided_cuda((4, 1, 128), (128, 512, 1), torch.float32)
        buf63 = empty_strided_cuda((4, 1, 128), (128, 512, 1), torch.float32)
        buf65 = empty_strided_cuda((4, 1, 128), (128, 512, 1), torch.float32)
        buf67 = empty_strided_cuda((4, 1, 128), (128, 512, 1), torch.float32)
        triton_red_fused_mul_sub_sum_3[grid(512)](buf1, primals_3, buf2,
            buf3, buf4, buf6, buf8, buf10, buf12, buf15, buf17, buf19,
            buf21, buf24, buf26, buf28, buf30, buf33, buf35, buf37, buf39,
            buf42, buf44, buf46, buf48, buf51, buf53, buf55, buf57, buf60,
            buf62, buf64, buf66, buf5, buf7, buf9, buf11, buf13, buf16,
            buf18, buf20, buf22, buf25, buf27, buf29, buf31, buf34, buf36,
            buf38, buf40, buf43, buf45, buf47, buf49, buf52, buf54, buf56,
            buf58, buf61, buf63, buf65, buf67, 512, 4096, XBLOCK=1, RBLOCK=
            2048, num_warps=16, num_stages=1)
        buf70 = empty_strided_cuda((4, 1, 128), (128, 512, 1), torch.float32)
        buf72 = empty_strided_cuda((4, 1, 128), (128, 512, 1), torch.float32)
        buf74 = empty_strided_cuda((4, 1, 128), (128, 512, 1), torch.float32)
        buf76 = empty_strided_cuda((4, 1, 128), (128, 512, 1), torch.float32)
        buf79 = empty_strided_cuda((4, 1, 128), (128, 512, 1), torch.float32)
        buf81 = empty_strided_cuda((4, 1, 128), (128, 512, 1), torch.float32)
        buf83 = empty_strided_cuda((4, 1, 128), (128, 512, 1), torch.float32)
        buf85 = empty_strided_cuda((4, 1, 128), (128, 512, 1), torch.float32)
        buf88 = empty_strided_cuda((4, 1, 128), (128, 512, 1), torch.float32)
        buf90 = empty_strided_cuda((4, 1, 128), (128, 512, 1), torch.float32)
        buf92 = empty_strided_cuda((4, 1, 128), (128, 512, 1), torch.float32)
        buf94 = empty_strided_cuda((4, 1, 128), (128, 512, 1), torch.float32)
        buf97 = empty_strided_cuda((4, 1, 128), (128, 512, 1), torch.float32)
        buf99 = empty_strided_cuda((4, 1, 128), (128, 512, 1), torch.float32)
        buf101 = empty_strided_cuda((4, 1, 128), (128, 512, 1), torch.float32)
        buf103 = empty_strided_cuda((4, 1, 128), (128, 512, 1), torch.float32)
        buf106 = empty_strided_cuda((4, 1, 128), (128, 512, 1), torch.float32)
        buf108 = empty_strided_cuda((4, 1, 128), (128, 512, 1), torch.float32)
        buf110 = empty_strided_cuda((4, 1, 128), (128, 512, 1), torch.float32)
        buf112 = empty_strided_cuda((4, 1, 128), (128, 512, 1), torch.float32)
        buf115 = empty_strided_cuda((4, 1, 128), (128, 512, 1), torch.float32)
        buf117 = empty_strided_cuda((4, 1, 128), (128, 512, 1), torch.float32)
        buf119 = empty_strided_cuda((4, 1, 128), (128, 512, 1), torch.float32)
        buf121 = empty_strided_cuda((4, 1, 128), (128, 512, 1), torch.float32)
        buf124 = empty_strided_cuda((4, 1, 128), (128, 512, 1), torch.float32)
        buf126 = empty_strided_cuda((4, 1, 128), (128, 512, 1), torch.float32)
        buf128 = empty_strided_cuda((4, 1, 128), (128, 512, 1), torch.float32)
        buf130 = empty_strided_cuda((4, 1, 128), (128, 512, 1), torch.float32)
        triton_red_fused_mul_sum_4[grid(512)](buf69, buf2, buf3, buf4,
            buf71, buf73, buf75, buf78, buf80, buf82, buf84, buf87, buf89,
            buf91, buf93, buf96, buf98, buf100, buf102, buf105, buf107,
            buf109, buf111, buf114, buf116, buf118, buf120, buf123, buf125,
            buf127, buf129, buf70, buf72, buf74, buf76, buf79, buf81, buf83,
            buf85, buf88, buf90, buf92, buf94, buf97, buf99, buf101, buf103,
            buf106, buf108, buf110, buf112, buf115, buf117, buf119, buf121,
            buf124, buf126, buf128, buf130, 512, 4096, XBLOCK=1, RBLOCK=
            2048, num_warps=16, num_stages=1)
        buf133 = empty_strided_cuda((4, 1, 128), (128, 512, 1), torch.float32)
        buf135 = empty_strided_cuda((4, 1, 128), (128, 512, 1), torch.float32)
        buf137 = empty_strided_cuda((4, 1, 128), (128, 512, 1), torch.float32)
        buf139 = empty_strided_cuda((4, 1, 128), (128, 512, 1), torch.float32)
        buf142 = empty_strided_cuda((4, 1, 128), (128, 512, 1), torch.float32)
        buf144 = empty_strided_cuda((4, 1, 128), (128, 512, 1), torch.float32)
        buf146 = empty_strided_cuda((4, 1, 128), (128, 512, 1), torch.float32)
        triton_red_fused_mul_sum_5[grid(512)](buf132, buf2, buf3, buf4,
            buf134, buf136, buf138, buf141, buf143, buf145, buf133, buf135,
            buf137, buf139, buf142, buf144, buf146, 512, 4096, XBLOCK=1,
            RBLOCK=1024, num_warps=16, num_stages=1)
        buf14 = empty_strided_cuda((4, 64, 128), (8192, 128, 1), torch.float32)
        buf23 = buf14
        del buf14
        buf32 = buf23
        del buf23
        buf41 = buf32
        del buf32
        buf50 = buf41
        del buf41
        buf59 = buf50
        del buf50
        buf68 = buf59
        del buf59
        buf77 = buf68
        del buf68
        buf86 = buf77
        del buf77
        buf95 = buf86
        del buf86
        buf104 = buf95
        del buf95
        buf113 = buf104
        del buf104
        buf122 = buf113
        del buf113
        buf131 = buf122
        del buf122
        buf140 = buf131
        del buf131
        buf147 = buf140
        del buf140
        buf148 = empty_strided_cuda((4, 64, 1), (64, 1, 256), torch.float32)
        buf149 = reinterpret_tensor(buf148, (4, 64, 1), (64, 1, 1), 0)
        del buf148
        triton_per_fused_copy_linalg_vector_norm_zeros_6[grid(256)](buf147,
            buf149, buf13, buf11, buf9, buf7, buf5, buf22, buf20, buf18,
            buf16, buf31, buf29, buf27, buf25, buf40, buf38, buf36, buf34,
            buf49, buf47, buf45, buf43, buf58, buf56, buf54, buf52, buf67,
            buf65, buf63, buf61, buf76, buf74, buf72, buf70, buf85, buf83,
            buf81, buf79, buf94, buf92, buf90, buf88, buf103, buf101, buf99,
            buf97, buf112, buf110, buf108, buf106, buf121, buf119, buf117,
            buf115, buf130, buf128, buf126, buf124, buf139, buf137, buf135,
            buf133, buf146, buf144, buf142, 256, 128, XBLOCK=1, num_warps=2,
            num_stages=1)
        del buf101
        del buf103
        del buf106
        del buf108
        del buf11
        del buf110
        del buf112
        del buf115
        del buf117
        del buf119
        del buf121
        del buf124
        del buf126
        del buf128
        del buf13
        del buf130
        del buf133
        del buf135
        del buf137
        del buf139
        del buf142
        del buf144
        del buf146
        del buf16
        del buf18
        del buf20
        del buf22
        del buf25
        del buf27
        del buf29
        del buf31
        del buf34
        del buf36
        del buf38
        del buf40
        del buf43
        del buf45
        del buf47
        del buf49
        del buf5
        del buf52
        del buf54
        del buf56
        del buf58
        del buf61
        del buf63
        del buf65
        del buf67
        del buf7
        del buf70
        del buf72
        del buf74
        del buf76
        del buf79
        del buf81
        del buf83
        del buf85
        del buf88
        del buf9
        del buf90
        del buf92
        del buf94
        del buf97
        del buf99
        buf150 = empty_strided_cuda((4, 1), (1, 4), torch.float32)
        buf151 = reinterpret_tensor(buf150, (4, 1), (1, 1), 0)
        del buf150
        buf152 = empty_strided_cuda((4, 8192), (8192, 1), torch.float32)
        triton_red_fused_div_linalg_vector_norm_7[grid(4)](buf151, buf147,
            buf149, buf152, 4, 8192, XBLOCK=1, RBLOCK=2048, num_warps=16,
            num_stages=1)
    return (buf152, primals_2, buf1, buf2, buf3, buf4, reinterpret_tensor(
        primals_3, (1, 128), (128, 1), 0), buf6, buf8, buf10, buf12, buf15,
        buf17, buf19, buf21, buf24, buf26, buf28, buf30, buf33, buf35,
        buf37, buf39, buf42, buf44, buf46, buf48, buf51, buf53, buf55,
        buf57, buf60, buf62, buf64, buf66, buf69, buf71, buf73, buf75,
        buf78, buf80, buf82, buf84, buf87, buf89, buf91, buf93, buf96,
        buf98, buf100, buf102, buf105, buf107, buf109, buf111, buf114,
        buf116, buf118, buf120, buf123, buf125, buf127, buf129, buf132,
        buf134, buf136, buf138, buf141, buf143, buf145, buf147, buf149, buf151)


class NetVLADNew(nn.Module):
    """NetVLAD layer implementation"""

    def __init__(self, num_clusters=64, dim=128, normalize_input=True,
        vladv2=False):
        """
        Args:
            num_clusters : int
                The number of clusters
            dim : int
                Dimension of descriptors
            alpha : float
                Parameter of initialization. Larger value is harder assignment.
            normalize_input : bool
                If true, descriptor-wise L2 normalization is applied to input.
            vladv2 : bool
                If true, use vladv2 otherwise use vladv1
        """
        super(NetVLADNew, self).__init__()
        self.num_clusters = num_clusters
        self.dim = dim
        self.alpha = 0
        self.vladv2 = vladv2
        self.normalize_input = normalize_input
        self.conv = nn.Conv2d(dim, num_clusters, kernel_size=(1, 1), bias=
            vladv2)
        self.centroids = nn.Parameter(torch.rand(num_clusters, dim))

    def init_params(self, clsts, traindescs):
        if self.vladv2 is False:
            clstsAssign = clsts / np.linalg.norm(clsts, axis=1, keepdims=True)
            dots = np.dot(clstsAssign, traindescs.T)
            dots.sort(0)
            dots = dots[::-1, :]
            self.alpha = (-np.log(0.01) / np.mean(dots[0, :] - dots[1, :])
                ).item()
            self.centroids = nn.Parameter(torch.from_numpy(clsts))
            self.conv.weight = nn.Parameter(torch.from_numpy(self.alpha *
                clstsAssign).unsqueeze(2).unsqueeze(3))
            self.conv.bias = None
        else:
            knn = NearestNeighbors(n_jobs=-1)
            knn.fit(traindescs)
            del traindescs
            dsSq = np.square(knn.kneighbors(clsts, 2)[1])
            del knn
            self.alpha = (-np.log(0.01) / np.mean(dsSq[:, 1] - dsSq[:, 0])
                ).item()
            self.centroids = nn.Parameter(torch.from_numpy(clsts))
            del clsts, dsSq
            self.conv.weight = nn.Parameter((2.0 * self.alpha * self.
                centroids).unsqueeze(-1).unsqueeze(-1))
            self.conv.bias = nn.Parameter(-self.alpha * self.centroids.norm
                (dim=1))

    def forward(self, input_0):
        primals_3 = self.centroids
        primals_2 = self.conv.weight
        primals_1 = input_0
        output = call([primals_1, primals_2, primals_3])
        return output[0]


Measured latency_ms: 0.015863