import torch
import triton
import triton.language as tl
from torch._inductor.runtime.triton_heuristics import grid
from torch._C import _cuda_getCurrentRawStream as get_raw_stream
import torch.nn as nn
assert_size_stride = torch._C._dynamo.guards.assert_size_stride
empty_strided_cuda = torch._C._dynamo.guards._empty_strided_cuda
reinterpret_tensor = torch._C._dynamo.guards._reinterpret_tensor


@triton.jit
def triton_poi_fused_stack_0(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 256
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x1 = xindex // 4 % 4
    x0 = xindex % 4
    x4 = xindex
    x2 = xindex // 16
    x5 = xindex % 16
    tmp0 = -2 + x1
    tmp1 = tl.full([1], 0, tl.int64)
    tmp2 = tmp0 >= tmp1
    tmp3 = tl.full([1], 4, tl.int64)
    tmp4 = tmp0 < tmp3
    tmp5 = -2 + x0
    tmp6 = tmp5 >= tmp1
    tmp7 = tmp5 < tmp3
    tmp8 = tmp2 & tmp4
    tmp9 = tmp8 & tmp6
    tmp10 = tmp9 & tmp7
    tmp11 = tl.load(in_ptr0 + (-10 + x4), tmp10 & xmask, other=0.0)
    tl.store(out_ptr0 + (x5 + 400 * x2), tmp11, xmask)


@triton.jit
def triton_poi_fused_stack_1(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 256
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x1 = xindex // 4 % 4
    x0 = xindex % 4
    x4 = xindex
    x2 = xindex // 16
    x5 = xindex % 16
    tmp0 = -2 + x1
    tmp1 = tl.full([1], 0, tl.int64)
    tmp2 = tmp0 >= tmp1
    tmp3 = tl.full([1], 4, tl.int64)
    tmp4 = tmp0 < tmp3
    tmp5 = -1 + x0
    tmp6 = tmp5 >= tmp1
    tmp7 = tmp5 < tmp3
    tmp8 = tmp2 & tmp4
    tmp9 = tmp8 & tmp6
    tmp10 = tmp9 & tmp7
    tmp11 = tl.load(in_ptr0 + (-9 + x4), tmp10 & xmask, other=0.0)
    tl.store(out_ptr0 + (x5 + 400 * x2), tmp11, xmask)


@triton.jit
def triton_poi_fused_stack_2(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 256
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x1 = xindex // 4 % 4
    x0 = xindex % 4
    x3 = xindex
    x2 = xindex // 16
    x4 = xindex % 16
    tmp0 = -2 + x1
    tmp1 = tl.full([1], 0, tl.int64)
    tmp2 = tmp0 >= tmp1
    tmp3 = tl.full([1], 4, tl.int64)
    tmp4 = tmp0 < tmp3
    tmp5 = x0
    tmp6 = tmp5 >= tmp1
    tmp7 = tmp5 < tmp3
    tmp8 = tmp2 & tmp4
    tmp9 = tmp8 & tmp6
    tmp10 = tmp9 & tmp7
    tmp11 = tl.load(in_ptr0 + (-8 + x3), tmp10 & xmask, other=0.0)
    tl.store(out_ptr0 + (x4 + 400 * x2), tmp11, xmask)


@triton.jit
def triton_poi_fused_stack_3(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 256
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x1 = xindex // 4 % 4
    x0 = xindex % 4
    x4 = xindex
    x2 = xindex // 16
    x5 = xindex % 16
    tmp0 = -2 + x1
    tmp1 = tl.full([1], 0, tl.int64)
    tmp2 = tmp0 >= tmp1
    tmp3 = tl.full([1], 4, tl.int64)
    tmp4 = tmp0 < tmp3
    tmp5 = 1 + x0
    tmp6 = tmp5 >= tmp1
    tmp7 = tmp5 < tmp3
    tmp8 = tmp2 & tmp4
    tmp9 = tmp8 & tmp6
    tmp10 = tmp9 & tmp7
    tmp11 = tl.load(in_ptr0 + (-7 + x4), tmp10 & xmask, other=0.0)
    tl.store(out_ptr0 + (x5 + 400 * x2), tmp11, xmask)


@triton.jit
def triton_poi_fused_stack_4(in_ptr0, out_ptr0, out_ptr1, xnumel, XBLOCK:
    tl.constexpr):
    xnumel = 256
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x1 = xindex // 4 % 4
    x0 = xindex % 4
    x4 = xindex
    x2 = xindex // 16
    x5 = xindex % 16
    tmp0 = -2 + x1
    tmp1 = tl.full([1], 0, tl.int64)
    tmp2 = tmp0 >= tmp1
    tmp3 = tl.full([1], 4, tl.int64)
    tmp4 = tmp0 < tmp3
    tmp5 = 2 + x0
    tmp6 = tmp5 >= tmp1
    tmp7 = tmp5 < tmp3
    tmp8 = tmp2 & tmp4
    tmp9 = tmp8 & tmp6
    tmp10 = tmp9 & tmp7
    tmp11 = tl.load(in_ptr0 + (-6 + x4), tmp10 & xmask, other=0.0)
    tmp12 = -1 + x1
    tmp13 = tmp12 >= tmp1
    tmp14 = tmp12 < tmp3
    tmp15 = -2 + x0
    tmp16 = tmp15 >= tmp1
    tmp17 = tmp15 < tmp3
    tmp18 = tmp13 & tmp14
    tmp19 = tmp18 & tmp16
    tmp20 = tmp19 & tmp17
    tmp21 = tl.load(in_ptr0 + (-6 + x4), tmp20 & xmask, other=0.0)
    tl.store(out_ptr0 + (x5 + 400 * x2), tmp11, xmask)
    tl.store(out_ptr1 + (x5 + 400 * x2), tmp21, xmask)


@triton.jit
def triton_poi_fused_stack_5(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 256
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x1 = xindex // 4 % 4
    x0 = xindex % 4
    x4 = xindex
    x2 = xindex // 16
    x5 = xindex % 16
    tmp0 = -1 + x1
    tmp1 = tl.full([1], 0, tl.int64)
    tmp2 = tmp0 >= tmp1
    tmp3 = tl.full([1], 4, tl.int64)
    tmp4 = tmp0 < tmp3
    tmp5 = -1 + x0
    tmp6 = tmp5 >= tmp1
    tmp7 = tmp5 < tmp3
    tmp8 = tmp2 & tmp4
    tmp9 = tmp8 & tmp6
    tmp10 = tmp9 & tmp7
    tmp11 = tl.load(in_ptr0 + (-5 + x4), tmp10 & xmask, other=0.0)
    tl.store(out_ptr0 + (x5 + 400 * x2), tmp11, xmask)


@triton.jit
def triton_poi_fused_stack_6(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 256
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x1 = xindex // 4 % 4
    x0 = xindex % 4
    x3 = xindex
    x2 = xindex // 16
    x4 = xindex % 16
    tmp0 = -1 + x1
    tmp1 = tl.full([1], 0, tl.int64)
    tmp2 = tmp0 >= tmp1
    tmp3 = tl.full([1], 4, tl.int64)
    tmp4 = tmp0 < tmp3
    tmp5 = x0
    tmp6 = tmp5 >= tmp1
    tmp7 = tmp5 < tmp3
    tmp8 = tmp2 & tmp4
    tmp9 = tmp8 & tmp6
    tmp10 = tmp9 & tmp7
    tmp11 = tl.load(in_ptr0 + (-4 + x3), tmp10 & xmask, other=0.0)
    tl.store(out_ptr0 + (x4 + 400 * x2), tmp11, xmask)


@triton.jit
def triton_poi_fused_stack_7(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 256
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x1 = xindex // 4 % 4
    x0 = xindex % 4
    x4 = xindex
    x2 = xindex // 16
    x5 = xindex % 16
    tmp0 = -1 + x1
    tmp1 = tl.full([1], 0, tl.int64)
    tmp2 = tmp0 >= tmp1
    tmp3 = tl.full([1], 4, tl.int64)
    tmp4 = tmp0 < tmp3
    tmp5 = 1 + x0
    tmp6 = tmp5 >= tmp1
    tmp7 = tmp5 < tmp3
    tmp8 = tmp2 & tmp4
    tmp9 = tmp8 & tmp6
    tmp10 = tmp9 & tmp7
    tmp11 = tl.load(in_ptr0 + (-3 + x4), tmp10 & xmask, other=0.0)
    tl.store(out_ptr0 + (x5 + 400 * x2), tmp11, xmask)


@triton.jit
def triton_poi_fused_stack_8(in_ptr0, out_ptr0, out_ptr1, xnumel, XBLOCK:
    tl.constexpr):
    xnumel = 256
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x1 = xindex // 4 % 4
    x0 = xindex % 4
    x4 = xindex
    x2 = xindex // 16
    x5 = xindex % 16
    tmp0 = -1 + x1
    tmp1 = tl.full([1], 0, tl.int64)
    tmp2 = tmp0 >= tmp1
    tmp3 = tl.full([1], 4, tl.int64)
    tmp4 = tmp0 < tmp3
    tmp5 = 2 + x0
    tmp6 = tmp5 >= tmp1
    tmp7 = tmp5 < tmp3
    tmp8 = tmp2 & tmp4
    tmp9 = tmp8 & tmp6
    tmp10 = tmp9 & tmp7
    tmp11 = tl.load(in_ptr0 + (-2 + x4), tmp10 & xmask, other=0.0)
    tmp12 = x1
    tmp13 = tmp12 >= tmp1
    tmp14 = tmp12 < tmp3
    tmp15 = -2 + x0
    tmp16 = tmp15 >= tmp1
    tmp17 = tmp15 < tmp3
    tmp18 = tmp13 & tmp14
    tmp19 = tmp18 & tmp16
    tmp20 = tmp19 & tmp17
    tmp21 = tl.load(in_ptr0 + (-2 + x4), tmp20 & xmask, other=0.0)
    tl.store(out_ptr0 + (x5 + 400 * x2), tmp11, xmask)
    tl.store(out_ptr1 + (x5 + 400 * x2), tmp21, xmask)


@triton.jit
def triton_poi_fused_stack_9(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 256
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x1 = xindex // 4 % 4
    x0 = xindex % 4
    x4 = xindex
    x2 = xindex // 16
    x5 = xindex % 16
    tmp0 = x1
    tmp1 = tl.full([1], 0, tl.int64)
    tmp2 = tmp0 >= tmp1
    tmp3 = tl.full([1], 4, tl.int64)
    tmp4 = tmp0 < tmp3
    tmp5 = -1 + x0
    tmp6 = tmp5 >= tmp1
    tmp7 = tmp5 < tmp3
    tmp8 = tmp2 & tmp4
    tmp9 = tmp8 & tmp6
    tmp10 = tmp9 & tmp7
    tmp11 = tl.load(in_ptr0 + (-1 + x4), tmp10 & xmask, other=0.0)
    tl.store(out_ptr0 + (x5 + 400 * x2), tmp11, xmask)


@triton.jit
def triton_poi_fused_stack_10(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 256
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x1 = xindex // 4 % 4
    x0 = xindex % 4
    x3 = xindex
    x2 = xindex // 16
    x4 = xindex % 16
    tmp0 = x1
    tmp1 = tl.full([1], 0, tl.int64)
    tmp2 = tmp0 >= tmp1
    tmp3 = tl.full([1], 4, tl.int64)
    tmp4 = tmp0 < tmp3
    tmp5 = x0
    tmp6 = tmp5 >= tmp1
    tmp7 = tmp5 < tmp3
    tmp8 = tmp2 & tmp4
    tmp9 = tmp8 & tmp6
    tmp10 = tmp9 & tmp7
    tmp11 = tl.load(in_ptr0 + x3, tmp10 & xmask, other=0.0)
    tl.store(out_ptr0 + (x4 + 400 * x2), tmp11, xmask)


@triton.jit
def triton_poi_fused_stack_11(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 256
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x1 = xindex // 4 % 4
    x0 = xindex % 4
    x4 = xindex
    x2 = xindex // 16
    x5 = xindex % 16
    tmp0 = x1
    tmp1 = tl.full([1], 0, tl.int64)
    tmp2 = tmp0 >= tmp1
    tmp3 = tl.full([1], 4, tl.int64)
    tmp4 = tmp0 < tmp3
    tmp5 = 1 + x0
    tmp6 = tmp5 >= tmp1
    tmp7 = tmp5 < tmp3
    tmp8 = tmp2 & tmp4
    tmp9 = tmp8 & tmp6
    tmp10 = tmp9 & tmp7
    tmp11 = tl.load(in_ptr0 + (1 + x4), tmp10 & xmask, other=0.0)
    tl.store(out_ptr0 + (x5 + 400 * x2), tmp11, xmask)


@triton.jit
def triton_poi_fused_stack_12(in_ptr0, out_ptr0, out_ptr1, xnumel, XBLOCK:
    tl.constexpr):
    xnumel = 256
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x1 = xindex // 4 % 4
    x0 = xindex % 4
    x4 = xindex
    x2 = xindex // 16
    x5 = xindex % 16
    tmp0 = x1
    tmp1 = tl.full([1], 0, tl.int64)
    tmp2 = tmp0 >= tmp1
    tmp3 = tl.full([1], 4, tl.int64)
    tmp4 = tmp0 < tmp3
    tmp5 = 2 + x0
    tmp6 = tmp5 >= tmp1
    tmp7 = tmp5 < tmp3
    tmp8 = tmp2 & tmp4
    tmp9 = tmp8 & tmp6
    tmp10 = tmp9 & tmp7
    tmp11 = tl.load(in_ptr0 + (2 + x4), tmp10 & xmask, other=0.0)
    tmp12 = 1 + x1
    tmp13 = tmp12 >= tmp1
    tmp14 = tmp12 < tmp3
    tmp15 = -2 + x0
    tmp16 = tmp15 >= tmp1
    tmp17 = tmp15 < tmp3
    tmp18 = tmp13 & tmp14
    tmp19 = tmp18 & tmp16
    tmp20 = tmp19 & tmp17
    tmp21 = tl.load(in_ptr0 + (2 + x4), tmp20 & xmask, other=0.0)
    tl.store(out_ptr0 + (x5 + 400 * x2), tmp11, xmask)
    tl.store(out_ptr1 + (x5 + 400 * x2), tmp21, xmask)


@triton.jit
def triton_poi_fused_stack_13(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 256
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x1 = xindex // 4 % 4
    x0 = xindex % 4
    x4 = xindex
    x2 = xindex // 16
    x5 = xindex % 16
    tmp0 = 1 + x1
    tmp1 = tl.full([1], 0, tl.int64)
    tmp2 = tmp0 >= tmp1
    tmp3 = tl.full([1], 4, tl.int64)
    tmp4 = tmp0 < tmp3
    tmp5 = -1 + x0
    tmp6 = tmp5 >= tmp1
    tmp7 = tmp5 < tmp3
    tmp8 = tmp2 & tmp4
    tmp9 = tmp8 & tmp6
    tmp10 = tmp9 & tmp7
    tmp11 = tl.load(in_ptr0 + (3 + x4), tmp10 & xmask, other=0.0)
    tl.store(out_ptr0 + (x5 + 400 * x2), tmp11, xmask)


@triton.jit
def triton_poi_fused_stack_14(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 256
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x1 = xindex // 4 % 4
    x0 = xindex % 4
    x3 = xindex
    x2 = xindex // 16
    x4 = xindex % 16
    tmp0 = 1 + x1
    tmp1 = tl.full([1], 0, tl.int64)
    tmp2 = tmp0 >= tmp1
    tmp3 = tl.full([1], 4, tl.int64)
    tmp4 = tmp0 < tmp3
    tmp5 = x0
    tmp6 = tmp5 >= tmp1
    tmp7 = tmp5 < tmp3
    tmp8 = tmp2 & tmp4
    tmp9 = tmp8 & tmp6
    tmp10 = tmp9 & tmp7
    tmp11 = tl.load(in_ptr0 + (4 + x3), tmp10 & xmask, other=0.0)
    tl.store(out_ptr0 + (x4 + 400 * x2), tmp11, xmask)


@triton.jit
def triton_poi_fused_stack_15(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 256
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x1 = xindex // 4 % 4
    x0 = xindex % 4
    x4 = xindex
    x2 = xindex // 16
    x5 = xindex % 16
    tmp0 = 1 + x1
    tmp1 = tl.full([1], 0, tl.int64)
    tmp2 = tmp0 >= tmp1
    tmp3 = tl.full([1], 4, tl.int64)
    tmp4 = tmp0 < tmp3
    tmp5 = 1 + x0
    tmp6 = tmp5 >= tmp1
    tmp7 = tmp5 < tmp3
    tmp8 = tmp2 & tmp4
    tmp9 = tmp8 & tmp6
    tmp10 = tmp9 & tmp7
    tmp11 = tl.load(in_ptr0 + (5 + x4), tmp10 & xmask, other=0.0)
    tl.store(out_ptr0 + (x5 + 400 * x2), tmp11, xmask)


@triton.jit
def triton_poi_fused_stack_16(in_ptr0, out_ptr0, out_ptr1, xnumel, XBLOCK:
    tl.constexpr):
    xnumel = 256
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x1 = xindex // 4 % 4
    x0 = xindex % 4
    x4 = xindex
    x2 = xindex // 16
    x5 = xindex % 16
    tmp0 = 1 + x1
    tmp1 = tl.full([1], 0, tl.int64)
    tmp2 = tmp0 >= tmp1
    tmp3 = tl.full([1], 4, tl.int64)
    tmp4 = tmp0 < tmp3
    tmp5 = 2 + x0
    tmp6 = tmp5 >= tmp1
    tmp7 = tmp5 < tmp3
    tmp8 = tmp2 & tmp4
    tmp9 = tmp8 & tmp6
    tmp10 = tmp9 & tmp7
    tmp11 = tl.load(in_ptr0 + (6 + x4), tmp10 & xmask, other=0.0)
    tmp12 = 2 + x1
    tmp13 = tmp12 >= tmp1
    tmp14 = tmp12 < tmp3
    tmp15 = -2 + x0
    tmp16 = tmp15 >= tmp1
    tmp17 = tmp15 < tmp3
    tmp18 = tmp13 & tmp14
    tmp19 = tmp18 & tmp16
    tmp20 = tmp19 & tmp17
    tmp21 = tl.load(in_ptr0 + (6 + x4), tmp20 & xmask, other=0.0)
    tl.store(out_ptr0 + (x5 + 400 * x2), tmp11, xmask)
    tl.store(out_ptr1 + (x5 + 400 * x2), tmp21, xmask)


@triton.jit
def triton_poi_fused_stack_17(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 256
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x1 = xindex // 4 % 4
    x0 = xindex % 4
    x4 = xindex
    x2 = xindex // 16
    x5 = xindex % 16
    tmp0 = 2 + x1
    tmp1 = tl.full([1], 0, tl.int64)
    tmp2 = tmp0 >= tmp1
    tmp3 = tl.full([1], 4, tl.int64)
    tmp4 = tmp0 < tmp3
    tmp5 = -1 + x0
    tmp6 = tmp5 >= tmp1
    tmp7 = tmp5 < tmp3
    tmp8 = tmp2 & tmp4
    tmp9 = tmp8 & tmp6
    tmp10 = tmp9 & tmp7
    tmp11 = tl.load(in_ptr0 + (7 + x4), tmp10 & xmask, other=0.0)
    tl.store(out_ptr0 + (x5 + 400 * x2), tmp11, xmask)


@triton.jit
def triton_poi_fused_stack_18(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 256
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x1 = xindex // 4 % 4
    x0 = xindex % 4
    x3 = xindex
    x2 = xindex // 16
    x4 = xindex % 16
    tmp0 = 2 + x1
    tmp1 = tl.full([1], 0, tl.int64)
    tmp2 = tmp0 >= tmp1
    tmp3 = tl.full([1], 4, tl.int64)
    tmp4 = tmp0 < tmp3
    tmp5 = x0
    tmp6 = tmp5 >= tmp1
    tmp7 = tmp5 < tmp3
    tmp8 = tmp2 & tmp4
    tmp9 = tmp8 & tmp6
    tmp10 = tmp9 & tmp7
    tmp11 = tl.load(in_ptr0 + (8 + x3), tmp10 & xmask, other=0.0)
    tl.store(out_ptr0 + (x4 + 400 * x2), tmp11, xmask)


@triton.jit
def triton_poi_fused_stack_19(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 256
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x1 = xindex // 4 % 4
    x0 = xindex % 4
    x4 = xindex
    x2 = xindex // 16
    x5 = xindex % 16
    tmp0 = 2 + x1
    tmp1 = tl.full([1], 0, tl.int64)
    tmp2 = tmp0 >= tmp1
    tmp3 = tl.full([1], 4, tl.int64)
    tmp4 = tmp0 < tmp3
    tmp5 = 1 + x0
    tmp6 = tmp5 >= tmp1
    tmp7 = tmp5 < tmp3
    tmp8 = tmp2 & tmp4
    tmp9 = tmp8 & tmp6
    tmp10 = tmp9 & tmp7
    tmp11 = tl.load(in_ptr0 + (9 + x4), tmp10 & xmask, other=0.0)
    tl.store(out_ptr0 + (x5 + 400 * x2), tmp11, xmask)


@triton.jit
def triton_poi_fused_stack_20(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 256
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x1 = xindex // 4 % 4
    x0 = xindex % 4
    x4 = xindex
    x2 = xindex // 16
    x5 = xindex % 16
    tmp0 = 2 + x1
    tmp1 = tl.full([1], 0, tl.int64)
    tmp2 = tmp0 >= tmp1
    tmp3 = tl.full([1], 4, tl.int64)
    tmp4 = tmp0 < tmp3
    tmp5 = 2 + x0
    tmp6 = tmp5 >= tmp1
    tmp7 = tmp5 < tmp3
    tmp8 = tmp2 & tmp4
    tmp9 = tmp8 & tmp6
    tmp10 = tmp9 & tmp7
    tmp11 = tl.load(in_ptr0 + (10 + x4), tmp10 & xmask, other=0.0)
    tl.store(out_ptr0 + (x5 + 400 * x2), tmp11, xmask)


@triton.jit
def triton_per_fused_div_mul_sum_21(in_out_ptr0, in_ptr0, in_ptr1, xnumel,
    rnumel, XBLOCK: tl.constexpr):
    xnumel = 256
    rnumel = 25
    RBLOCK: tl.constexpr = 32
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:, None]
    xmask = xindex < xnumel
    rindex = tl.arange(0, RBLOCK)[None, :]
    rmask = rindex < rnumel
    x3 = xindex
    r2 = rindex
    x0 = xindex % 16
    x1 = xindex // 16
    tmp0 = tl.load(in_ptr0 + x3, xmask, eviction_policy='evict_last')
    tmp1 = tl.load(in_ptr1 + (x0 + 16 * r2 + 400 * x1), rmask & xmask,
        other=0.0)
    tmp2 = tmp0 * tmp1
    tmp3 = tl.broadcast_to(tmp2, [XBLOCK, RBLOCK])
    tmp5 = tl.where(rmask & xmask, tmp3, 0)
    tmp6 = tl.sum(tmp5, 1)[:, None]
    tmp7 = 1.0
    tmp8 = tmp6 / tmp7
    tmp9 = tmp8 * tmp7
    tl.debug_barrier()
    tl.store(in_out_ptr0 + x3, tmp9, xmask)


def call(args):
    arg0_1, arg1_1 = args
    args.clear()
    assert_size_stride(arg0_1, (4, 4, 4, 4), (64, 16, 4, 1))
    assert_size_stride(arg1_1, (4, 4, 4, 4), (64, 16, 4, 1))
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        buf25 = empty_strided_cuda((4, 4, 25, 4, 4), (1600, 400, 16, 4, 1),
            torch.float32)
        buf0 = reinterpret_tensor(buf25, (4, 4, 1, 4, 4), (1600, 400, 16, 4,
            1), 0)
        get_raw_stream(0)
        triton_poi_fused_stack_0[grid(256)](arg0_1, buf0, 256, XBLOCK=256,
            num_warps=4, num_stages=1)
        buf1 = reinterpret_tensor(buf25, (4, 4, 1, 4, 4), (1600, 400, 16, 4,
            1), 16)
        triton_poi_fused_stack_1[grid(256)](arg0_1, buf1, 256, XBLOCK=256,
            num_warps=4, num_stages=1)
        buf2 = reinterpret_tensor(buf25, (4, 4, 1, 4, 4), (1600, 400, 16, 4,
            1), 32)
        triton_poi_fused_stack_2[grid(256)](arg0_1, buf2, 256, XBLOCK=256,
            num_warps=4, num_stages=1)
        buf3 = reinterpret_tensor(buf25, (4, 4, 1, 4, 4), (1600, 400, 16, 4,
            1), 48)
        triton_poi_fused_stack_3[grid(256)](arg0_1, buf3, 256, XBLOCK=256,
            num_warps=4, num_stages=1)
        buf4 = reinterpret_tensor(buf25, (4, 4, 1, 4, 4), (1600, 400, 16, 4,
            1), 64)
        buf5 = reinterpret_tensor(buf25, (4, 4, 1, 4, 4), (1600, 400, 16, 4,
            1), 80)
        triton_poi_fused_stack_4[grid(256)](arg0_1, buf4, buf5, 256, XBLOCK
            =128, num_warps=4, num_stages=1)
        buf6 = reinterpret_tensor(buf25, (4, 4, 1, 4, 4), (1600, 400, 16, 4,
            1), 96)
        triton_poi_fused_stack_5[grid(256)](arg0_1, buf6, 256, XBLOCK=256,
            num_warps=4, num_stages=1)
        buf7 = reinterpret_tensor(buf25, (4, 4, 1, 4, 4), (1600, 400, 16, 4,
            1), 112)
        triton_poi_fused_stack_6[grid(256)](arg0_1, buf7, 256, XBLOCK=128,
            num_warps=4, num_stages=1)
        buf8 = reinterpret_tensor(buf25, (4, 4, 1, 4, 4), (1600, 400, 16, 4,
            1), 128)
        triton_poi_fused_stack_7[grid(256)](arg0_1, buf8, 256, XBLOCK=256,
            num_warps=4, num_stages=1)
        buf9 = reinterpret_tensor(buf25, (4, 4, 1, 4, 4), (1600, 400, 16, 4,
            1), 144)
        buf10 = reinterpret_tensor(buf25, (4, 4, 1, 4, 4), (1600, 400, 16, 
            4, 1), 160)
        triton_poi_fused_stack_8[grid(256)](arg0_1, buf9, buf10, 256,
            XBLOCK=256, num_warps=4, num_stages=1)
        buf11 = reinterpret_tensor(buf25, (4, 4, 1, 4, 4), (1600, 400, 16, 
            4, 1), 176)
        triton_poi_fused_stack_9[grid(256)](arg0_1, buf11, 256, XBLOCK=128,
            num_warps=4, num_stages=1)
        buf12 = reinterpret_tensor(buf25, (4, 4, 1, 4, 4), (1600, 400, 16, 
            4, 1), 192)
        triton_poi_fused_stack_10[grid(256)](arg0_1, buf12, 256, XBLOCK=256,
            num_warps=4, num_stages=1)
        buf13 = reinterpret_tensor(buf25, (4, 4, 1, 4, 4), (1600, 400, 16, 
            4, 1), 208)
        triton_poi_fused_stack_11[grid(256)](arg0_1, buf13, 256, XBLOCK=256,
            num_warps=4, num_stages=1)
        buf14 = reinterpret_tensor(buf25, (4, 4, 1, 4, 4), (1600, 400, 16, 
            4, 1), 224)
        buf15 = reinterpret_tensor(buf25, (4, 4, 1, 4, 4), (1600, 400, 16, 
            4, 1), 240)
        triton_poi_fused_stack_12[grid(256)](arg0_1, buf14, buf15, 256,
            XBLOCK=128, num_warps=4, num_stages=1)
        buf16 = reinterpret_tensor(buf25, (4, 4, 1, 4, 4), (1600, 400, 16, 
            4, 1), 256)
        triton_poi_fused_stack_13[grid(256)](arg0_1, buf16, 256, XBLOCK=256,
            num_warps=4, num_stages=1)
        buf17 = reinterpret_tensor(buf25, (4, 4, 1, 4, 4), (1600, 400, 16, 
            4, 1), 272)
        triton_poi_fused_stack_14[grid(256)](arg0_1, buf17, 256, XBLOCK=256,
            num_warps=4, num_stages=1)
        buf18 = reinterpret_tensor(buf25, (4, 4, 1, 4, 4), (1600, 400, 16, 
            4, 1), 288)
        triton_poi_fused_stack_15[grid(256)](arg0_1, buf18, 256, XBLOCK=256,
            num_warps=4, num_stages=1)
        buf19 = reinterpret_tensor(buf25, (4, 4, 1, 4, 4), (1600, 400, 16, 
            4, 1), 304)
        buf20 = reinterpret_tensor(buf25, (4, 4, 1, 4, 4), (1600, 400, 16, 
            4, 1), 320)
        triton_poi_fused_stack_16[grid(256)](arg0_1, buf19, buf20, 256,
            XBLOCK=128, num_warps=4, num_stages=1)
        buf21 = reinterpret_tensor(buf25, (4, 4, 1, 4, 4), (1600, 400, 16, 
            4, 1), 336)
        triton_poi_fused_stack_17[grid(256)](arg0_1, buf21, 256, XBLOCK=128,
            num_warps=4, num_stages=1)
        buf22 = reinterpret_tensor(buf25, (4, 4, 1, 4, 4), (1600, 400, 16, 
            4, 1), 352)
        triton_poi_fused_stack_18[grid(256)](arg0_1, buf22, 256, XBLOCK=128,
            num_warps=4, num_stages=1)
        buf23 = reinterpret_tensor(buf25, (4, 4, 1, 4, 4), (1600, 400, 16, 
            4, 1), 368)
        triton_poi_fused_stack_19[grid(256)](arg0_1, buf23, 256, XBLOCK=128,
            num_warps=4, num_stages=1)
        buf24 = reinterpret_tensor(buf25, (4, 4, 1, 4, 4), (1600, 400, 16, 
            4, 1), 384)
        triton_poi_fused_stack_20[grid(256)](arg0_1, buf24, 256, XBLOCK=256,
            num_warps=4, num_stages=1)
        del arg0_1
        buf26 = empty_strided_cuda((4, 4, 1, 4, 4), (64, 16, 16, 4, 1),
            torch.float32)
        buf27 = reinterpret_tensor(buf26, (4, 4, 4, 4), (64, 16, 4, 1), 0)
        del buf26
        triton_per_fused_div_mul_sum_21[grid(256)](buf27, arg1_1, buf25, 
            256, 25, XBLOCK=8, num_warps=2, num_stages=1)
        del arg1_1
        del buf0
        del buf1
        del buf10
        del buf11
        del buf12
        del buf13
        del buf14
        del buf15
        del buf16
        del buf17
        del buf18
        del buf19
        del buf2
        del buf20
        del buf21
        del buf22
        del buf23
        del buf24
        del buf25
        del buf3
        del buf4
        del buf5
        del buf6
        del buf7
        del buf8
        del buf9
    return buf27,


class KernelConvNew(nn.Module):
    """
    the class of computing prediction
    """

    def __init__(self, kernel_size=[5], sep_conv=False, core_bias=False):
        super(KernelConvNew, self).__init__()
        self.kernel_size = sorted(kernel_size)
        self.sep_conv = sep_conv
        self.core_bias = core_bias

    def _sep_conv_core(self, core, batch_size, N, color, height, width):
        """
        convert the sep_conv core to conv2d core
        2p --> p^2
        :param core: shape: batch*(N*2*K)*height*width
        :return:
        """
        kernel_total = sum(self.kernel_size)
        core = core.view(batch_size, N, -1, color, height, width)
        if not self.core_bias:
            core_1, core_2 = torch.split(core, kernel_total, dim=2)
        else:
            core_1, core_2, core_3 = torch.split(core, kernel_total, dim=2)
        core_out = {}
        cur = 0
        for K in self.kernel_size:
            t1 = core_1[:, :, cur:cur + K, ...].view(batch_size, N, K, 1, 3,
                height, width)
            t2 = core_2[:, :, cur:cur + K, ...].view(batch_size, N, 1, K, 3,
                height, width)
            core_out[K] = torch.einsum('ijklno,ijlmno->ijkmno', [t1, t2]).view(
                batch_size, N, K * K, color, height, width)
            cur += K
        return core_out, None if not self.core_bias else core_3.squeeze()

    def _convert_dict(self, core, batch_size, N, color, height, width):
        """
        make sure the core to be a dict, generally, only one kind of kernel size is suitable for the func.
        :param core: shape: batch_size*(N*K*K)*height*width
        :return: core_out, a dict
        """
        core_out = {}
        core = core.view(batch_size, N, -1, color, height, width)
        core_out[self.kernel_size[0]] = core[:, :, 0:self.kernel_size[0] **
            2, ...]
        bias = None if not self.core_bias else core[:, :, -1, ...]
        return core_out, bias

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


Measured latency_ms: 0.015114