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
import torch.nn as nn
import torch.nn.functional
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_cat_0(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, out_ptr0, xnumel, XBLOCK: tl.
    constexpr):
    xnumel = 64
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x2 = xindex // 16
    x1 = xindex // 4 % 4
    x0 = xindex % 4
    x3 = xindex
    tmp0 = x2
    tl.full([1], 0, tl.int64)
    tmp3 = tl.full([1], 1, tl.int64)
    tmp4 = tmp0 < tmp3
    tmp5 = x1
    tmp7 = tmp5 < tmp3
    tmp8 = tmp7 & tmp4
    tmp9 = tl.load(in_ptr0 + x0, tmp8 & xmask, eviction_policy='evict_last',
        other=0.0)
    tmp10 = tmp5 >= tmp3
    tmp11 = tl.full([1], 2, tl.int64)
    tmp12 = tmp5 < tmp11
    tmp13 = tmp10 & tmp12
    tmp14 = tmp13 & tmp4
    tmp15 = tl.load(in_ptr1 + x0, tmp14 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp16 = tmp5 >= tmp11
    tmp17 = tl.full([1], 3, tl.int64)
    tmp18 = tmp5 < tmp17
    tmp19 = tmp16 & tmp18
    tmp20 = tmp19 & tmp4
    tmp21 = tl.load(in_ptr2 + x0, tmp20 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp22 = tmp5 >= tmp17
    tl.full([1], 4, tl.int64)
    tmp25 = tmp22 & tmp4
    tmp26 = tl.load(in_ptr3 + x0, tmp25 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp27 = tl.where(tmp19, tmp21, tmp26)
    tmp28 = tl.where(tmp13, tmp15, tmp27)
    tmp29 = tl.where(tmp7, tmp9, tmp28)
    tmp30 = tl.full(tmp29.shape, 0.0, tmp29.dtype)
    tmp31 = tl.where(tmp4, tmp29, tmp30)
    tmp32 = tmp0 >= tmp3
    tmp33 = tmp0 < tmp11
    tmp34 = tmp32 & tmp33
    tmp35 = tmp7 & tmp34
    tmp36 = tl.load(in_ptr4 + x0, tmp35 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp37 = tmp13 & tmp34
    tmp38 = tl.load(in_ptr5 + x0, tmp37 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp39 = tmp19 & tmp34
    tmp40 = tl.load(in_ptr6 + x0, tmp39 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp41 = tmp22 & tmp34
    tmp42 = tl.load(in_ptr7 + x0, tmp41 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp43 = tl.where(tmp19, tmp40, tmp42)
    tmp44 = tl.where(tmp13, tmp38, tmp43)
    tmp45 = tl.where(tmp7, tmp36, tmp44)
    tmp46 = tl.full(tmp45.shape, 0.0, tmp45.dtype)
    tmp47 = tl.where(tmp34, tmp45, tmp46)
    tmp48 = tmp0 >= tmp11
    tmp49 = tmp0 < tmp17
    tmp50 = tmp48 & tmp49
    tmp51 = tmp7 & tmp50
    tmp52 = tl.load(in_ptr8 + x0, tmp51 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp53 = tmp13 & tmp50
    tmp54 = tl.load(in_ptr9 + x0, tmp53 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp55 = tmp19 & tmp50
    tmp56 = tl.load(in_ptr10 + x0, tmp55 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp57 = tmp22 & tmp50
    tmp58 = tl.load(in_ptr11 + x0, tmp57 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp59 = tl.where(tmp19, tmp56, tmp58)
    tmp60 = tl.where(tmp13, tmp54, tmp59)
    tmp61 = tl.where(tmp7, tmp52, tmp60)
    tmp62 = tl.full(tmp61.shape, 0.0, tmp61.dtype)
    tmp63 = tl.where(tmp50, tmp61, tmp62)
    tmp64 = tmp0 >= tmp17
    tmp66 = tmp7 & tmp64
    tmp67 = tl.load(in_ptr12 + x0, tmp66 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp68 = tmp13 & tmp64
    tmp69 = tl.load(in_ptr13 + x0, tmp68 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp70 = tmp19 & tmp64
    tmp71 = tl.load(in_ptr14 + x0, tmp70 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp72 = tmp22 & tmp64
    tmp73 = tl.load(in_ptr15 + x0, tmp72 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp74 = tl.where(tmp19, tmp71, tmp73)
    tmp75 = tl.where(tmp13, tmp69, tmp74)
    tmp76 = tl.where(tmp7, tmp67, tmp75)
    tmp77 = tl.full(tmp76.shape, 0.0, tmp76.dtype)
    tmp78 = tl.where(tmp64, tmp76, tmp77)
    tmp79 = tl.where(tmp50, tmp63, tmp78)
    tmp80 = tl.where(tmp34, tmp47, tmp79)
    tmp81 = tl.where(tmp4, tmp31, tmp80)
    tl.store(out_ptr0 + x3, tmp81, xmask)


@triton.jit
def triton_poi_fused_clone_1(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 64
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex % 4
    x1 = xindex // 4 % 4
    x2 = xindex // 16
    x3 = xindex
    tmp0 = tl.load(in_ptr0 + (x0 + 4 * x2 + 16 * x1), xmask)
    tl.store(out_ptr0 + x3, tmp0, xmask)


@triton.jit
def triton_poi_fused_add_2(in_out_ptr0, in_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 64
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x2 = xindex
    x1 = xindex // 16
    tmp0 = tl.load(in_out_ptr0 + x2, xmask)
    tmp1 = tl.load(in_ptr0 + x1, xmask, eviction_policy='evict_last')
    tmp2 = tmp0 + tmp1
    tl.store(in_out_ptr0 + x2, tmp2, xmask)


def call(args):
    primals_1, primals_2, primals_3, primals_4 = args
    args.clear()
    assert_size_stride(primals_1, (4, 4), (4, 1))
    assert_size_stride(primals_2, (4, 4, 4), (16, 4, 1))
    assert_size_stride(primals_3, (4, 4), (4, 1))
    assert_size_stride(primals_4, (1, 4, 1), (4, 1, 1))
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        buf0 = empty_strided_cuda((4, 1), (1, 1), torch.float32)
        extern_kernels.mm(primals_1, reinterpret_tensor(primals_2, (4, 1),
            (1, 1), 0), out=buf0)
        buf1 = empty_strided_cuda((4, 1), (1, 1), torch.float32)
        extern_kernels.mm(primals_1, reinterpret_tensor(primals_2, (4, 1),
            (1, 1), 4), out=buf1)
        buf2 = empty_strided_cuda((4, 1), (1, 1), torch.float32)
        extern_kernels.mm(primals_1, reinterpret_tensor(primals_2, (4, 1),
            (1, 1), 8), out=buf2)
        buf3 = empty_strided_cuda((4, 1), (1, 1), torch.float32)
        extern_kernels.mm(primals_1, reinterpret_tensor(primals_2, (4, 1),
            (1, 1), 12), out=buf3)
        buf4 = empty_strided_cuda((4, 1), (1, 1), torch.float32)
        extern_kernels.mm(primals_1, reinterpret_tensor(primals_2, (4, 1),
            (1, 1), 16), out=buf4)
        buf5 = empty_strided_cuda((4, 1), (1, 1), torch.float32)
        extern_kernels.mm(primals_1, reinterpret_tensor(primals_2, (4, 1),
            (1, 1), 20), out=buf5)
        buf6 = empty_strided_cuda((4, 1), (1, 1), torch.float32)
        extern_kernels.mm(primals_1, reinterpret_tensor(primals_2, (4, 1),
            (1, 1), 24), out=buf6)
        buf7 = empty_strided_cuda((4, 1), (1, 1), torch.float32)
        extern_kernels.mm(primals_1, reinterpret_tensor(primals_2, (4, 1),
            (1, 1), 28), out=buf7)
        buf8 = empty_strided_cuda((4, 1), (1, 1), torch.float32)
        extern_kernels.mm(primals_1, reinterpret_tensor(primals_2, (4, 1),
            (1, 1), 32), out=buf8)
        buf9 = empty_strided_cuda((4, 1), (1, 1), torch.float32)
        extern_kernels.mm(primals_1, reinterpret_tensor(primals_2, (4, 1),
            (1, 1), 36), out=buf9)
        buf10 = empty_strided_cuda((4, 1), (1, 1), torch.float32)
        extern_kernels.mm(primals_1, reinterpret_tensor(primals_2, (4, 1),
            (1, 1), 40), out=buf10)
        buf11 = empty_strided_cuda((4, 1), (1, 1), torch.float32)
        extern_kernels.mm(primals_1, reinterpret_tensor(primals_2, (4, 1),
            (1, 1), 44), out=buf11)
        buf12 = empty_strided_cuda((4, 1), (1, 1), torch.float32)
        extern_kernels.mm(primals_1, reinterpret_tensor(primals_2, (4, 1),
            (1, 1), 48), out=buf12)
        buf13 = empty_strided_cuda((4, 1), (1, 1), torch.float32)
        extern_kernels.mm(primals_1, reinterpret_tensor(primals_2, (4, 1),
            (1, 1), 52), out=buf13)
        buf14 = empty_strided_cuda((4, 1), (1, 1), torch.float32)
        extern_kernels.mm(primals_1, reinterpret_tensor(primals_2, (4, 1),
            (1, 1), 56), out=buf14)
        buf15 = empty_strided_cuda((4, 1), (1, 1), torch.float32)
        extern_kernels.mm(primals_1, reinterpret_tensor(primals_2, (4, 1),
            (1, 1), 60), out=buf15)
        del primals_1
        del primals_2
        buf16 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        get_raw_stream(0)
        triton_poi_fused_cat_0[grid(64)](buf0, buf1, buf2, buf3, buf4, buf5,
            buf6, buf7, buf8, buf9, buf10, buf11, buf12, buf13, buf14,
            buf15, buf16, 64, XBLOCK=64, num_warps=1, num_stages=1)
        del buf0
        del buf1
        del buf10
        del buf11
        del buf12
        del buf13
        del buf14
        del buf15
        del buf2
        del buf3
        del buf4
        del buf5
        del buf6
        del buf7
        del buf8
        del buf9
        buf17 = empty_strided_cuda((4, 4, 4, 1), (16, 4, 1, 1), torch.float32)
        triton_poi_fused_clone_1[grid(64)](buf16, buf17, 64, XBLOCK=64,
            num_warps=1, num_stages=1)
        buf18 = reinterpret_tensor(buf16, (1, 4, 16), (64, 16, 1), 0)
        del buf16
        extern_kernels.bmm(reinterpret_tensor(primals_3, (1, 4, 4), (16, 4,
            1), 0), reinterpret_tensor(buf17, (1, 4, 16), (0, 16, 1), 0),
            out=buf18)
        del primals_3
        buf19 = reinterpret_tensor(buf18, (4, 4, 4), (4, 16, 1), 0)
        del buf18
        triton_poi_fused_add_2[grid(64)](buf19, primals_4, 64, XBLOCK=64,
            num_warps=1, num_stages=1)
        del primals_4
    return buf19, reinterpret_tensor(buf17, (1, 16, 4), (64, 1, 16), 0)


class CoboundaryNew(nn.Module):

    def __init__(self, C_in, C_out, enable_bias=True, variance=1.0):
        super().__init__()
        assert C_in > 0
        assert C_out > 0
        self.C_in = C_in
        self.C_out = C_out
        self.enable_bias = enable_bias
        self.theta = nn.parameter.Parameter(variance * torch.randn((self.
            C_out, self.C_in)))
        if self.enable_bias:
            self.bias = nn.parameter.Parameter(torch.zeros((1, self.C_out, 1)))
        else:
            self.bias = 0.0

    def forward(self, input_0, input_1):
        primals_1 = self.theta
        primals_4 = self.bias
        primals_3 = input_0
        primals_2 = input_1
        output = call([primals_1, primals_2, primals_3, primals_4])
        return output[0]


Measured latency_ms: 0.16247