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
from torch import nn
import torch.autograd
assert_size_stride = torch._C._dynamo.guards.assert_size_stride
empty_strided_cuda = torch._C._dynamo.guards._empty_strided_cuda


@triton.jit
def triton_poi_fused_add_mul_sigmoid_0(in_ptr0, in_ptr1, in_ptr2, in_ptr3,
    in_ptr4, in_ptr5, in_ptr6, in_ptr7, in_ptr8, in_ptr9, in_ptr10,
    out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 256
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex % 64
    x2 = xindex
    tmp0 = tl.load(in_ptr0 + x0, xmask, eviction_policy='evict_last')
    tmp1 = tl.load(in_ptr1 + x2, xmask)
    tmp2 = tl.load(in_ptr2 + x0, xmask, eviction_policy='evict_last')
    tmp5 = tl.load(in_ptr3 + x2, xmask)
    tmp6 = tl.load(in_ptr4 + x0, xmask, eviction_policy='evict_last')
    tmp9 = tl.load(in_ptr5 + x0, xmask, eviction_policy='evict_last')
    tmp13 = tl.load(in_ptr6 + x0, xmask, eviction_policy='evict_last')
    tmp14 = tl.load(in_ptr7 + x0, xmask, eviction_policy='evict_last')
    tmp15 = tl.load(in_ptr8 + x0, xmask, eviction_policy='evict_last')
    tmp18 = tl.load(in_ptr9 + x0, xmask, eviction_policy='evict_last')
    tmp21 = tl.load(in_ptr10 + x0, xmask, eviction_policy='evict_last')
    tmp3 = tmp1 * tmp2
    tmp4 = tmp0 + tmp3
    tmp7 = tmp5 * tmp6
    tmp8 = tmp4 + tmp7
    tmp10 = tmp5 * tmp9
    tmp11 = tmp1 * tmp10
    tmp12 = tmp8 + tmp11
    tmp16 = tmp1 * tmp15
    tmp17 = tmp14 + tmp16
    tmp19 = tmp5 * tmp18
    tmp20 = tmp17 + tmp19
    tmp22 = tmp5 * tmp21
    tmp23 = tmp1 * tmp22
    tmp24 = tmp20 + tmp23
    tmp25 = tl.sigmoid(tmp24)
    tmp26 = tmp13 * tmp25
    tmp27 = tmp12 + tmp26
    tl.store(out_ptr0 + x2, tmp27, xmask)


def call(args):
    (primals_1, primals_2, primals_3, primals_4, primals_5, primals_6,
        primals_7, primals_8, primals_9, primals_10, primals_11) = args
    args.clear()
    assert_size_stride(primals_1, (4, 4, 4, 4), (64, 16, 4, 1))
    assert_size_stride(primals_2, (4, 4, 4, 4), (64, 16, 4, 1))
    assert_size_stride(primals_3, (4, 4, 4), (16, 4, 1))
    assert_size_stride(primals_4, (4, 4, 4), (16, 4, 1))
    assert_size_stride(primals_5, (4, 4, 4), (16, 4, 1))
    assert_size_stride(primals_6, (4, 4, 4), (16, 4, 1))
    assert_size_stride(primals_7, (4, 4, 4), (16, 4, 1))
    assert_size_stride(primals_8, (4, 4, 4), (16, 4, 1))
    assert_size_stride(primals_9, (4, 4, 4), (16, 4, 1))
    assert_size_stride(primals_10, (4, 4, 4), (16, 4, 1))
    assert_size_stride(primals_11, (4, 4, 4), (16, 4, 1))
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        buf0 = empty_strided_cuda((4, 4, 4, 4), (64, 16, 4, 1), torch.float32)
        get_raw_stream(0)
        triton_poi_fused_add_mul_sigmoid_0[grid(256)](primals_3, primals_1,
            primals_4, primals_2, primals_5, primals_6, primals_7,
            primals_8, primals_9, primals_10, primals_11, buf0, 256, XBLOCK
            =256, num_warps=4, num_stages=1)
        del primals_3
        del primals_4
        del primals_5
        del primals_6
    return (buf0, primals_1, primals_2, primals_7, primals_8, primals_9,
        primals_10, primals_11)


class CombinatorNew(nn.Module):
    """
    The vanilla combinator function g() that combines vertical and
    lateral connections as explained in Pezeshki et al. (2016).
    The weights are initialized as described in Eq. 17
    and the g() is defined in Eq. 16.
    """

    def __init__(self, n_channels, length, data_type='2d'):
        super(CombinatorNew, self).__init__()
        if data_type == '2d':
            zeros = torch.zeros(n_channels, length, length)
            ones = torch.ones(n_channels, length, length)
        elif data_type == '1d':
            zeros = torch.zeros(n_channels, length)
            ones = torch.ones(n_channels, length)
        else:
            raise ValueError
        self.b0 = nn.Parameter(zeros)
        self.w0z = nn.Parameter(ones)
        self.w0u = nn.Parameter(zeros)
        self.w0zu = nn.Parameter(ones)
        self.b1 = nn.Parameter(zeros)
        self.w1z = nn.Parameter(ones)
        self.w1u = nn.Parameter(zeros)
        self.w1zu = nn.Parameter(zeros)
        self.wsig = nn.Parameter(ones)

    def forward(self, input_0, input_1):
        primals_3 = self.b0
        primals_4 = self.w0z
        primals_5 = self.w0u
        primals_6 = self.w0zu
        primals_7 = self.b1
        primals_8 = self.w1z
        primals_9 = self.w1u
        primals_10 = self.w1zu
        primals_11 = self.wsig
        primals_1 = input_0
        primals_2 = input_1
        output = call([primals_1, primals_2, primals_3, primals_4,
            primals_5, primals_6, primals_7, primals_8, primals_9,
            primals_10, primals_11])
        return output[0]


Measured latency_ms: 0.022587