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.nn import Module
import math
import numpy as np
from torch.nn.modules.utils import _pair
import torch.nn.functional as F
from torch.nn.parameter import Parameter
from torch.nn import init
from torch._jit_internal import Optional
from torch.nn.modules.module import Module
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_per_fused_mean_sub_0(in_out_ptr0, in_ptr0, out_ptr0, xnumel,
    rnumel, XBLOCK: tl.constexpr):
    xnumel = 16
    RBLOCK: tl.constexpr = 16
    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)
    r1 = rindex
    x0 = xindex
    tmp0 = tl.load(in_ptr0 + (r1 + 16 * x0), xmask, other=0.0)
    tmp1 = tl.broadcast_to(tmp0, [XBLOCK, RBLOCK])
    tmp3 = tl.where(xmask, tmp1, 0)
    tmp4 = tl.sum(tmp3, 1)[:, None]
    tmp5 = 16.0
    tmp6 = tmp4 / tmp5
    tmp7 = tmp0 - tmp6
    tl.debug_barrier()
    tl.store(in_out_ptr0 + x0, tmp6, xmask)
    tl.store(out_ptr0 + (r1 + 16 * x0), tmp7, xmask)


@triton.jit
def triton_poi_fused_add_1(in_ptr0, in_ptr1, out_ptr0, xnumel, XBLOCK: tl.
    constexpr):
    xnumel = 1024
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex
    tmp0 = tl.load(in_ptr0 + x0, xmask)
    tmp1 = tl.load(in_ptr1 + x0, xmask)
    tmp2 = tmp0 + tmp1
    tl.store(out_ptr0 + x0, tmp2, xmask)


@triton.jit
def triton_poi_fused_add_2(in_out_ptr0, in_ptr0, in_ptr1, xnumel, XBLOCK:
    tl.constexpr):
    xnumel = 256
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x3 = xindex
    x0 = xindex % 4
    x2 = xindex // 64
    tmp0 = tl.load(in_out_ptr0 + x3, xmask)
    tmp1 = tl.load(in_ptr0 + (x2 + 4 * x0), xmask, eviction_policy='evict_last'
        )
    tmp2 = tl.load(in_ptr1 + 0)
    tmp3 = tl.broadcast_to(tmp2, [XBLOCK])
    tmp4 = tmp1 * tmp3
    tmp5 = tmp0 + tmp4
    tl.store(in_out_ptr0 + x3, tmp5, xmask)


@triton.jit
def triton_poi_fused_convolution_3(in_ptr0, out_ptr0, ynumel, xnumel,
    YBLOCK: tl.constexpr, XBLOCK: tl.constexpr):
    ynumel = 4
    xnumel = 64
    yoffset = tl.program_id(1) * YBLOCK
    yindex = yoffset + tl.arange(0, YBLOCK)[None, :]
    ymask = yindex < ynumel
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:, None]
    xmask = xindex < xnumel
    x1 = xindex
    y0 = yindex
    tmp0 = tl.load(in_ptr0 + (y0 + 4 * x1), xmask & ymask, eviction_policy=
        'evict_last')
    tl.store(out_ptr0 + (x1 + 64 * y0), tmp0, xmask & ymask)


@triton.jit
def triton_poi_fused_convolution_4(in_out_ptr0, in_ptr0, xnumel, XBLOCK: tl
    .constexpr):
    xnumel = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x2 = xindex
    x0 = xindex % 4
    tmp0 = tl.load(in_out_ptr0 + x2, xmask)
    tmp1 = tl.load(in_ptr0 + x0, 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, primals_5, primals_6 = args
    args.clear()
    assert_size_stride(primals_1, (4, 16, 16), (256, 16, 1))
    assert_size_stride(primals_2, (4, 16, 16), (256, 16, 1))
    assert_size_stride(primals_3, (1,), (1,))
    assert_size_stride(primals_4, (4, 4, 4, 4), (64, 16, 4, 1))
    assert_size_stride(primals_5, (4,), (1,))
    assert_size_stride(primals_6, (4, 4, 4, 4), (64, 16, 4, 1))
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        buf0 = empty_strided_cuda((4, 4, 1), (4, 1, 16), torch.float32)
        buf1 = reinterpret_tensor(buf0, (4, 4, 1), (4, 1, 1), 0)
        del buf0
        buf3 = empty_strided_cuda((4, 4, 16), (64, 16, 1), torch.float32)
        get_raw_stream(0)
        triton_per_fused_mean_sub_0[grid(16)](buf1, primals_4, buf3, 16, 16,
            XBLOCK=8, num_warps=2, num_stages=1)
        del primals_4
        buf2 = empty_strided_cuda((4, 16, 16), (256, 16, 1), torch.float32)
        triton_poi_fused_add_1[grid(1024)](primals_1, primals_2, buf2, 1024,
            XBLOCK=256, num_warps=4, num_stages=1)
        del primals_1
        del primals_2
        buf4 = empty_strided_cuda((4, 16, 4), (64, 4, 1), torch.float32)
        extern_kernels.bmm(buf2, reinterpret_tensor(buf3, (4, 16, 4), (16, 
            1, 64), 0), out=buf4)
        buf5 = reinterpret_tensor(buf4, (4, 4, 4, 4), (1, 64, 16, 4), 0)
        del buf4
        triton_poi_fused_add_2[grid(256)](buf5, buf1, primals_3, 256,
            XBLOCK=256, num_warps=4, num_stages=1)
        buf6 = empty_strided_cuda((4, 4, 4, 4), (64, 16, 4, 1), torch.float32)
        triton_poi_fused_convolution_3[grid(4, 64)](buf5, buf6, 4, 64,
            XBLOCK=32, YBLOCK=4, num_warps=4, num_stages=1)
        buf7 = extern_kernels.convolution(primals_6, buf6, stride=(1, 1),
            padding=(0, 0), dilation=(1, 1), transposed=False,
            output_padding=(0, 0), groups=1, bias=None)
        assert_size_stride(buf7, (4, 4, 1, 1), (4, 1, 1, 1))
        del buf6
        buf8 = buf7
        del buf7
        triton_poi_fused_convolution_4[grid(16)](buf8, primals_5, 16,
            XBLOCK=16, num_warps=1, num_stages=1)
        del primals_5
    return buf8, primals_3, primals_6, buf1, buf5, reinterpret_tensor(buf2,
        (4, 16, 16), (256, 1, 16), 0), reinterpret_tensor(buf3, (4, 4, 16),
        (16, 64, 1), 0)


class ShapeConv2dNew(Module):
    """
       ShapeConv2d can be used as an alternative for torch.nn.Conv2d.
    """
    __constants__ = ['stride', 'padding', 'dilation', 'groups',
        'padding_mode', 'output_padding', 'in_channels', 'out_channels',
        'kernel_size', 'D_mul']
    __annotations__ = {'bias': Optional[torch.Tensor]}

    def __init__(self, in_channels, out_channels, kernel_size, D_mul=None,
        stride=1, padding=0, dilation=1, groups=1, bias=True, padding_mode=
        'zeros'):
        super(ShapeConv2dNew, self).__init__()
        kernel_size = _pair(kernel_size)
        stride = _pair(stride)
        padding = _pair(padding)
        dilation = _pair(dilation)
        if in_channels % groups != 0:
            raise ValueError('in_channels must be divisible by groups')
        if out_channels % groups != 0:
            raise ValueError('out_channels must be divisible by groups')
        valid_padding_modes = {'zeros', 'reflect', 'replicate', 'circular'}
        if padding_mode not in valid_padding_modes:
            raise ValueError(
                "padding_mode must be one of {}, but got padding_mode='{}'"
                .format(valid_padding_modes, padding_mode))
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = kernel_size
        self.stride = stride
        self.padding = padding
        self.dilation = dilation
        self.groups = groups
        self.padding_mode = padding_mode
        self.testing = not self.training
        self._padding_repeated_twice = tuple(x for x in self.padding for _ in
            range(2))
        M = self.kernel_size[0]
        N = self.kernel_size[1]
        self.D_mul = M * N if D_mul is None or M * N <= 1 else D_mul
        self.weight = Parameter(torch.Tensor(out_channels, in_channels //
            groups, M, N))
        init.kaiming_uniform_(self.weight, a=math.sqrt(5))
        if M * N > 1:
            self.Shape = Parameter(torch.Tensor(in_channels, M * N, self.D_mul)
                )
            self.Base = Parameter(torch.Tensor(1))
            init_zero = np.zeros([in_channels, M * N, self.D_mul], dtype=np
                .float32)
            init_one = np.ones([1], dtype=np.float32)
            self.Shape.data = torch.from_numpy(init_zero)
            self.Base.data = torch.from_numpy(init_one)
            eye = torch.reshape(torch.eye(M * N, dtype=torch.float32), (1, 
                M * N, M * N))
            D_diag = eye.repeat((in_channels, 1, self.D_mul // (M * N)))
            if self.D_mul % (M * N) != 0:
                zeros = torch.zeros([1, M * N, self.D_mul % (M * N)])
                self.D_diag = Parameter(torch.cat([D_diag, zeros], dim=2),
                    requires_grad=False)
            else:
                self.D_diag = Parameter(D_diag, requires_grad=False)
        if bias:
            self.bias = Parameter(torch.Tensor(out_channels))
            fan_in, _ = init._calculate_fan_in_and_fan_out(self.weight)
            bound = 1 / math.sqrt(fan_in)
            init.uniform_(self.bias, -bound, bound)
        else:
            self.register_parameter('bias', None)

    def extra_repr(self):
        s = (
            '{in_channels}, {out_channels}, kernel_size={kernel_size}, stride={stride}'
            )
        if self.padding != (0,) * len(self.padding):
            s += ', padding={padding}'
        if self.dilation != (1,) * len(self.dilation):
            s += ', dilation={dilation}'
        if self.groups != 1:
            s += ', groups={groups}'
        if self.bias is None:
            s += ', bias=False'
        if self.padding_mode != 'zeros':
            s += ', padding_mode={padding_mode}'
        return s.format(**self.__dict__)

    def __setstate__(self, state):
        super(ShapeConv2dNew, self).__setstate__(state)
        if not hasattr(self, 'padding_mode'):
            self.padding_mode = 'zeros'

    def _conv_forward(self, input, weight):
        if self.padding_mode != 'zeros':
            return F.conv2d(F.pad(input, self._padding_repeated_twice, mode
                =self.padding_mode), weight, self.bias, self.stride, _pair(
                0), self.dilation, self.groups)
        return F.conv2d(input, weight, self.bias, self.stride, self.padding,
            self.dilation, self.groups)

    def compute_shape_w(self):
        Shape = self.Shape + self.D_diag
        Base = self.Base
        W = torch.reshape(self.weight, (self.out_channels // self.groups,
            self.in_channels, self.D_mul))
        W_base = torch.mean(W, [2], keepdims=True)
        W_shape = W - W_base
        D_shape = torch.reshape(torch.einsum('ims,ois->oim', Shape, W_shape
            ), self.weight.shape)
        D_base = torch.reshape(W_base * Base, (self.out_channels, self.
            in_channels // self.groups, 1, 1))
        DW = D_shape + D_base
        return DW

    def _load_from_state_dict(self, state_dict, prefix, local_metadata,
        strict, missing_keys, unexpected_keys, error_msgs):
        self.testing = not self.training
        super(ShapeConv2dNew, self)._load_from_state_dict(state_dict,
            prefix, local_metadata, strict, missing_keys, unexpected_keys,
            error_msgs)
        if self.kernel_size[0] * self.kernel_size[1] > 1 and not self.training:
            self.weight.data = self.compute_shape_w()

    def forward(self, input_0):
        primals_4 = self.weight
        primals_1 = self.Shape
        primals_3 = self.Base
        primals_2 = self.D_diag
        primals_5 = self.bias
        primals_6 = input_0
        output = call([primals_1, primals_2, primals_3, primals_4,
            primals_5, primals_6])
        return output[0]


Measured latency_ms: 0.018194