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._inductor.runtime import triton_helpers
import torch.nn as nn
from torch.autograd import Variable
from torch.nn import Parameter
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_relu_0(in_out_ptr0, in_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_out_ptr0 + x0, xmask)
    tmp1 = tl.load(in_ptr0 + x0, xmask)
    tmp2 = tmp0 + tmp1
    tmp3 = tl.full([1], 0, tl.int32)
    tmp4 = triton_helpers.maximum(tmp3, tmp2)
    tl.store(in_out_ptr0 + x0, tmp4, xmask)


@triton.jit
def triton_poi_fused_relu_1(in_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_out_ptr0 + x0, xmask)
    tmp1 = tl.full([1], 0, tl.int32)
    tmp2 = triton_helpers.maximum(tmp1, tmp0)
    tl.store(in_out_ptr0 + x0, tmp2, xmask)


@triton.jit
def triton_poi_fused_convolution_relu_threshold_backward_2(in_out_ptr0,
    in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 512
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x3 = xindex
    x1 = xindex // 64 % 2
    tmp0 = tl.load(in_out_ptr0 + x3, xmask)
    tmp1 = tl.load(in_ptr0 + x1, xmask, eviction_policy='evict_last')
    tmp2 = tmp0 + tmp1
    tmp3 = tl.full([1], 0, tl.int32)
    tmp4 = triton_helpers.maximum(tmp3, tmp2)
    tmp5 = 0.0
    tmp6 = tmp4 <= tmp5
    tl.store(in_out_ptr0 + x3, tmp4, xmask)
    tl.store(out_ptr0 + x3, tmp6, 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, 1, 1), (1, 1, 1))
    assert_size_stride(primals_2, (4, 1, 64), (64, 64, 1))
    assert_size_stride(primals_3, (4, 1, 1), (1, 1, 1))
    assert_size_stride(primals_4, (4, 4, 1), (4, 1, 1))
    assert_size_stride(primals_5, (2, 4, 1), (4, 1, 1))
    assert_size_stride(primals_6, (2,), (1,))
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        buf0 = extern_kernels.convolution(primals_2, primals_1, stride=(1,),
            padding=(0,), dilation=(1,), transposed=False, output_padding=(
            0,), groups=1, bias=None)
        assert_size_stride(buf0, (4, 4, 64), (256, 64, 1))
        buf1 = extern_kernels.convolution(primals_2, primals_3, stride=(1,),
            padding=(0,), dilation=(1,), transposed=False, output_padding=(
            0,), groups=1, bias=None)
        assert_size_stride(buf1, (4, 4, 64), (256, 64, 1))
        buf2 = buf0
        del buf0
        get_raw_stream(0)
        triton_poi_fused_add_relu_0[grid(1024)](buf2, buf1, 1024, XBLOCK=
            128, num_warps=4, num_stages=1)
        del buf1
        buf3 = extern_kernels.convolution(buf2, primals_4, stride=(1,),
            padding=(0,), dilation=(1,), transposed=False, output_padding=(
            0,), groups=1, bias=None)
        assert_size_stride(buf3, (4, 4, 64), (256, 64, 1))
        buf4 = buf3
        del buf3
        triton_poi_fused_relu_1[grid(1024)](buf4, 1024, XBLOCK=128,
            num_warps=4, num_stages=1)
        buf5 = extern_kernels.convolution(buf4, primals_5, stride=(1,),
            padding=(0,), dilation=(1,), transposed=False, output_padding=(
            0,), groups=1, bias=None)
        assert_size_stride(buf5, (4, 2, 64), (128, 64, 1))
        buf6 = buf5
        del buf5
        buf7 = empty_strided_cuda((4, 2, 64), (128, 64, 1), torch.bool)
        triton_poi_fused_convolution_relu_threshold_backward_2[grid(512)](buf6,
            primals_6, buf7, 512, XBLOCK=256, num_warps=4, num_stages=1)
        del primals_6
    return (buf6, primals_1, primals_2, primals_3, primals_4, primals_5,
        buf2, buf4, buf7)


class Conv1dExt(nn.Conv1d):

    def __init__(self, *args, **kwargs):
        super(Conv1dExt, self).__init__(*args, **kwargs)
        self.init_ncc()
        self.input_tied_modules = []
        self.output_tied_modules = []

    def init_ncc(self):
        w = self.weight.view(self.weight.size(0), -1)
        mean = torch.mean(w, dim=1).unsqueeze(1)
        self.t0_factor = w - mean
        self.t0_norm = torch.norm(w, p=2, dim=1)
        self.start_ncc = Variable(torch.zeros(self.out_channels))
        self.start_ncc = self.normalized_cross_correlation()

    def normalized_cross_correlation(self):
        w = self.weight.view(self.weight.size(0), -1)
        t_norm = torch.norm(w, p=2, dim=1)
        if self.in_channels == 1 & sum(self.kernel_size) == 1:
            ncc = w.squeeze() / torch.norm(self.t0_norm, p=2)
            ncc = ncc - self.start_ncc
            return ncc
        mean = torch.mean(w, dim=1).unsqueeze(1)
        t_factor = w - mean
        h_product = self.t0_factor * t_factor
        cov = torch.sum(h_product, dim=1)
        denom = self.t0_norm * t_norm
        ncc = cov / denom
        ncc = ncc - self.start_ncc
        return ncc

    def split_output_channel(self, channel_i):
        """Split one output channel (a feature) into two, but retain summed value

            Args:
                channel_i: (int) number of channel to be split.  the ith channel
        """
        self.out_channels += 1
        orig_weight = self.weight.data
        split_pos = 2 * torch.rand(self.in_channels, self.kernel_size[0])
        new_weight = torch.zeros(self.out_channels, self.in_channels, self.
            kernel_size[0])
        if channel_i > 0:
            new_weight[:channel_i, :, :] = orig_weight[:channel_i, :, :]
        new_weight[channel_i, :, :] = orig_weight[channel_i, :, :] * split_pos
        new_weight[channel_i + 1, :, :] = orig_weight[channel_i, :, :] * (2 -
            split_pos)
        if channel_i + 2 < self.out_channels:
            new_weight[channel_i + 2, :, :] = orig_weight[channel_i + 1, :, :]
        if self.bias is not None:
            orig_bias = self.bias.data
            new_bias = torch.zeros(self.out_channels)
            new_bias[:channel_i + 1] = orig_bias[:channel_i + 1]
            new_bias[channel_i + 1:] = orig_bias[channel_i:]
            self.bias = Parameter(new_bias)
        self.weight = Parameter(new_weight)
        self.init_ncc()

    def split_input_channel(self, channel_i):
        if channel_i > self.in_channels:
            None
            return
        self.in_channels += 1
        orig_weight = self.weight.data
        dup_slice = orig_weight[:, channel_i, :] * 0.5
        new_weight = torch.zeros(self.out_channels, self.in_channels, self.
            kernel_size[0])
        if channel_i > 0:
            new_weight[:, :channel_i, :] = orig_weight[:, :channel_i, :]
        new_weight[:, channel_i, :] = dup_slice
        new_weight[:, channel_i + 1, :] = dup_slice
        if channel_i + 1 < self.in_channels:
            new_weight[:, channel_i + 2, :] = orig_weight[:, channel_i + 1, :]
        self.weight = Parameter(new_weight)
        self.init_ncc()

    def split_feature(self, feature_i):
        """Splits feature in output and input channels

            Args:
                feature_i: (int)
        """
        self.split_output_channel(channel_i=feature_i)
        for dep in self.input_tied_modules:
            dep.split_input_channel(channel_i=feature_i)
        for dep in self.output_tied_modules:
            dep.split_output_channel(channel_i=feature_i)

    def split_features(self, threshold):
        """Decides which features to split if they are below a specific threshold

            Args:
                threshold: (float?) less than 1.
        """
        ncc = self.normalized_cross_correlation()
        for i, ncc_val in enumerate(ncc):
            if ncc_val < threshold:
                None
                self.split_feature(i)


class NetNew(nn.Module):

    def __init__(self):
        super(NetNew, self).__init__()
        self.conv1 = Conv1dExt(in_channels=1, out_channels=4, kernel_size=1,
            bias=False)
        self.conv2 = Conv1dExt(in_channels=1, out_channels=4, kernel_size=1,
            bias=False)
        self.conv3 = Conv1dExt(in_channels=4, out_channels=4, kernel_size=1,
            bias=False)
        self.conv4 = Conv1dExt(in_channels=4, out_channels=2, kernel_size=1,
            bias=True)
        self.conv1.input_tied_modules = [self.conv3]
        self.conv1.output_tied_modules = [self.conv2]
        self.conv2.input_tied_modules = [self.conv3]
        self.conv2.output_tied_modules = [self.conv1]
        self.conv3.input_tied_modules = [self.conv4]

    def forward(self, input_0):
        primals_1 = self.conv1.weight
        primals_3 = self.conv2.weight
        primals_4 = self.conv3.weight
        primals_5 = self.conv4.weight
        primals_6 = self.conv4.bias
        primals_2 = input_0
        output = call([primals_1, primals_2, primals_3, primals_4,
            primals_5, primals_6])
        return output[0]


Measured latency_ms: 0.077523