# AOT ID: ['0_forward']
from ctypes import c_void_p, c_long, c_int
import torch
import math
import random
import os
import tempfile
from math import inf, nan
from torch._inductor.hooks import run_intermediate_hooks
from torch._inductor.utils import maybe_profile
from torch._inductor.codegen.memory_planning import _align as align
from torch import device, empty_strided
from torch._inductor.async_compile import AsyncCompile
from torch._inductor.select_algorithm import extern_kernels
from torch._inductor.codegen.multi_kernel import MultiKernelCall
import triton
import triton.language as tl
from torch._inductor.runtime.triton_heuristics import grid, split_scan_grid, grid_combo_kernels, start_graph, end_graph
from torch._C import _cuda_getCurrentRawStream as get_raw_stream

aten = torch.ops.aten
inductor_ops = torch.ops.inductor
_quantized = torch.ops._quantized
assert_size_stride = torch._C._dynamo.guards.assert_size_stride
empty_strided_cpu = torch._C._dynamo.guards._empty_strided_cpu
empty_strided_cuda = torch._C._dynamo.guards._empty_strided_cuda
empty_strided_xpu = torch._C._dynamo.guards._empty_strided_xpu
reinterpret_tensor = torch._C._dynamo.guards._reinterpret_tensor
alloc_from_pool = torch.ops.inductor._alloc_from_pool
async_compile = AsyncCompile()


# kernel path: runs/run_shard_8/inductor_cache/mm/cmm76txwg4trotpajgsjdvrbrz3ev5bvf7flp45qexvyjgbe2j6u.py
# Topologically Sorted Source Nodes: [X], Original ATen: [aten.cat]
# Source node to ATen node mapping:
#   X => cat_4
# Graph fragment:
#   %cat_4 : [num_users=1] = call_function[target=torch.ops.aten.cat.default](args = ([%unsqueeze_4, %unsqueeze_9, %unsqueeze_14, %unsqueeze_19],), kwargs = {})
triton_poi_fused_cat_0 = async_compile.triton('triton_poi_fused_cat_0', '''
import triton
import triton.language as tl
from triton.compiler.compiler import AttrsDescriptor

from torch._inductor.runtime import triton_helpers, triton_heuristics
from torch._inductor.runtime.triton_helpers import libdevice, math as tl_math
from torch._inductor.runtime.hints import AutotuneHint, ReductionHint, TileHint, instance_descriptor, DeviceProperties

@triton_heuristics.pointwise(
    size_hints=[64], 
    filename=__file__,
    triton_meta={'signature': {0: '*fp32', 1: '*fp32', 2: '*fp32', 3: '*fp32', 4: '*fp32', 5: '*fp32', 6: '*fp32', 7: '*fp32', 8: '*fp32', 9: '*fp32', 10: '*fp32', 11: '*fp32', 12: '*fp32', 13: '*fp32', 14: '*fp32', 15: '*fp32', 16: '*fp32', 17: 'i32'}, 'device': DeviceProperties(type='cuda', index=0, cc=80, major=8, regs_per_multiprocessor=65536, max_threads_per_multi_processor=2048, multi_processor_count=108), 'constants': {}, 'configs': [AttrsDescriptor(divisible_by_16=(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17), equal_to_1=())]},
    inductor_meta={'autotune_hints': set(), 'kernel_name': 'triton_poi_fused_cat_0', 'mutated_arg_names': [], 'no_x_dim': False, 'num_load': 16, 'num_reduction': 0, 'backend_hash': 'A9C866B4A14FD3277824029365D703C2427B2E685E54EC9B3EF4ADC8D1EEAC1D', 'are_deterministic_algorithms_enabled': False, 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': None, 'force_disable_caches': False, 'dynamic_scale_rblock': True, 'max_autotune': False, 'max_autotune_pointwise': False, 'min_split_scan_rblock': 256, 'spill_threshold': 16, 'store_cubin': False},
    min_elem_per_thread=0
)
@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
    tmp1 = tl.full([1], 0, tl.int64)
    tmp2 = tmp0 >= tmp1
    tmp3 = tl.full([1], 1, tl.int64)
    tmp4 = tmp0 < tmp3
    tmp5 = x1
    tmp6 = tmp5 >= tmp1
    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
    tmp23 = tl.full([1], 4, tl.int64)
    tmp24 = tmp5 < tmp23
    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
    tmp65 = tmp0 < tmp23
    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)
''', device_str='cuda')


# kernel path: runs/run_shard_8/inductor_cache/pu/cpuk3ivxfd5uot4vfqv32lo2lfmbbflxlacsfch7ifj6ctlef2mu.py
# Topologically Sorted Source Nodes: [y], Original ATen: [aten.clone]
# Source node to ATen node mapping:
#   y => clone
# Graph fragment:
#   %clone : [num_users=1] = call_function[target=torch.ops.aten.clone.default](args = (%permute_19,), kwargs = {memory_format: torch.contiguous_format})
triton_poi_fused_clone_1 = async_compile.triton('triton_poi_fused_clone_1', '''
import triton
import triton.language as tl
from triton.compiler.compiler import AttrsDescriptor

from torch._inductor.runtime import triton_helpers, triton_heuristics
from torch._inductor.runtime.triton_helpers import libdevice, math as tl_math
from torch._inductor.runtime.hints import AutotuneHint, ReductionHint, TileHint, instance_descriptor, DeviceProperties

@triton_heuristics.pointwise(
    size_hints=[64], 
    filename=__file__,
    triton_meta={'signature': {0: '*fp32', 1: '*fp32', 2: 'i32'}, 'device': DeviceProperties(type='cuda', index=0, cc=80, major=8, regs_per_multiprocessor=65536, max_threads_per_multi_processor=2048, multi_processor_count=108), 'constants': {}, 'configs': [AttrsDescriptor(divisible_by_16=(0, 1, 2), equal_to_1=())]},
    inductor_meta={'autotune_hints': set(), 'kernel_name': 'triton_poi_fused_clone_1', 'mutated_arg_names': [], 'no_x_dim': False, 'num_load': 1, 'num_reduction': 0, 'backend_hash': 'A9C866B4A14FD3277824029365D703C2427B2E685E54EC9B3EF4ADC8D1EEAC1D', 'are_deterministic_algorithms_enabled': False, 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': None, 'force_disable_caches': False, 'dynamic_scale_rblock': True, 'max_autotune': False, 'max_autotune_pointwise': False, 'min_split_scan_rblock': 256, 'spill_threshold': 16, 'store_cubin': False},
    min_elem_per_thread=0
)
@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)
''', device_str='cuda')


# kernel path: runs/run_shard_8/inductor_cache/fr/cfr54sk6c74tkbbvmlgy4hk7zqy5ljpoivf7jutbtxgzn463tkvb.py
# Topologically Sorted Source Nodes: [add], Original ATen: [aten.add]
# Source node to ATen node mapping:
#   add => add
# Graph fragment:
#   %add : [num_users=1] = call_function[target=torch.ops.aten.add.Tensor](args = (%view_3, %primals_4), kwargs = {})
triton_poi_fused_add_2 = async_compile.triton('triton_poi_fused_add_2', '''
import triton
import triton.language as tl
from triton.compiler.compiler import AttrsDescriptor

from torch._inductor.runtime import triton_helpers, triton_heuristics
from torch._inductor.runtime.triton_helpers import libdevice, math as tl_math
from torch._inductor.runtime.hints import AutotuneHint, ReductionHint, TileHint, instance_descriptor, DeviceProperties

@triton_heuristics.pointwise(
    size_hints=[64], 
    filename=__file__,
    triton_meta={'signature': {0: '*fp32', 1: '*fp32', 2: 'i32'}, 'device': DeviceProperties(type='cuda', index=0, cc=80, major=8, regs_per_multiprocessor=65536, max_threads_per_multi_processor=2048, multi_processor_count=108), 'constants': {}, 'configs': [AttrsDescriptor(divisible_by_16=(0, 1, 2), equal_to_1=())]},
    inductor_meta={'autotune_hints': set(), 'kernel_name': 'triton_poi_fused_add_2', 'mutated_arg_names': ['in_out_ptr0'], 'no_x_dim': False, 'num_load': 2, 'num_reduction': 0, 'backend_hash': 'A9C866B4A14FD3277824029365D703C2427B2E685E54EC9B3EF4ADC8D1EEAC1D', 'are_deterministic_algorithms_enabled': False, 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': None, 'force_disable_caches': False, 'dynamic_scale_rblock': True, 'max_autotune': False, 'max_autotune_pointwise': False, 'min_split_scan_rblock': 256, 'spill_threshold': 16, 'store_cubin': False},
    min_elem_per_thread=0
)
@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)
''', device_str='cuda')


async_compile.wait(globals())
del async_compile

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)
        # Topologically Sorted Source Nodes: [mm], Original ATen: [aten.mm]
        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)
        # Topologically Sorted Source Nodes: [mm_1], Original ATen: [aten.mm]
        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)
        # Topologically Sorted Source Nodes: [mm_2], Original ATen: [aten.mm]
        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)
        # Topologically Sorted Source Nodes: [mm_3], Original ATen: [aten.mm]
        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)
        # Topologically Sorted Source Nodes: [mm_4], Original ATen: [aten.mm]
        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)
        # Topologically Sorted Source Nodes: [mm_5], Original ATen: [aten.mm]
        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)
        # Topologically Sorted Source Nodes: [mm_6], Original ATen: [aten.mm]
        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)
        # Topologically Sorted Source Nodes: [mm_7], Original ATen: [aten.mm]
        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)
        # Topologically Sorted Source Nodes: [mm_8], Original ATen: [aten.mm]
        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)
        # Topologically Sorted Source Nodes: [mm_9], Original ATen: [aten.mm]
        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)
        # Topologically Sorted Source Nodes: [mm_10], Original ATen: [aten.mm]
        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)
        # Topologically Sorted Source Nodes: [mm_11], Original ATen: [aten.mm]
        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)
        # Topologically Sorted Source Nodes: [mm_12], Original ATen: [aten.mm]
        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)
        # Topologically Sorted Source Nodes: [mm_13], Original ATen: [aten.mm]
        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)
        # Topologically Sorted Source Nodes: [mm_14], Original ATen: [aten.mm]
        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)
        # Topologically Sorted Source Nodes: [mm_15], Original ATen: [aten.mm]
        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)
        # Topologically Sorted Source Nodes: [X], Original ATen: [aten.cat]
        stream0 = get_raw_stream(0)
        triton_poi_fused_cat_0.run(buf0, buf1, buf2, buf3, buf4, buf5, buf6, buf7, buf8, buf9, buf10, buf11, buf12, buf13, buf14, buf15, buf16, 64, grid=grid(64), stream=stream0)
        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)
        # Topologically Sorted Source Nodes: [y], Original ATen: [aten.clone]
        triton_poi_fused_clone_1.run(buf16, buf17, 64, grid=grid(64), stream=stream0)
        buf18 = reinterpret_tensor(buf16, (1, 4, 16), (64, 16, 1), 0); del buf16  # reuse
        # Topologically Sorted Source Nodes: [y], Original ATen: [aten.bmm]
        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  # reuse
        # Topologically Sorted Source Nodes: [add], Original ATen: [aten.add]
        triton_poi_fused_add_2.run(buf19, primals_4, 64, grid=grid(64), stream=stream0)
        del primals_4
    return (buf19, reinterpret_tensor(buf17, (1, 16, 4), (64, 1, 16), 0), )


def benchmark_compiled_module(times=10, repeat=10):
    from torch._dynamo.testing import rand_strided
    from torch._inductor.utils import print_performance
    primals_1 = rand_strided((4, 4), (4, 1), device='cuda:0', dtype=torch.float32)
    primals_2 = rand_strided((4, 4, 4), (16, 4, 1), device='cuda:0', dtype=torch.float32)
    primals_3 = rand_strided((4, 4), (4, 1), device='cuda:0', dtype=torch.float32)
    primals_4 = rand_strided((1, 4, 1), (4, 1, 1), device='cuda:0', dtype=torch.float32)
    fn = lambda: call([primals_1, primals_2, primals_3, primals_4])
    return print_performance(fn, times=times, repeat=repeat)


if __name__ == "__main__":
    from torch._inductor.wrapper_benchmark import compiled_module_main
    compiled_module_main('None', benchmark_compiled_module)

import torch
import torch.nn as nn
import torch.nn.functional


class Coboundary(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, D, x):
        assert len(D.shape) == 2
        B, C_in, M = x.shape
        assert D.shape[1] == M
        assert C_in == self.C_in
        N = D.shape[0]
        X = []
        for b in range(0, B):
            X12 = []
            for c_in in range(0, self.C_in):
                X12.append(D.mm(x[b, c_in, :].unsqueeze(1)).transpose(0, 1))
            X12 = torch.cat(X12, 0)
            assert X12.shape == (self.C_in, N)
            X.append(X12.unsqueeze(0))
        X = torch.cat(X, 0)
        assert X.shape == (B, self.C_in, N)
        y = torch.einsum('oi,bin->bon', (self.theta, X))
        assert y.shape == (B, self.C_out, N)
        return y + self.bias


def get_inputs():
    return [torch.rand([4, 4]), torch.rand([4, 4, 4])]


def get_init_inputs():
    return [[], {'C_in': 4, 'C_out': 4}]

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]
