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
import numpy as np
import torch.nn as nn
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_clone_view_0(in_out_ptr0, in_ptr0, xnumel, XBLOCK: tl.
    constexpr):
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    tl.full([XBLOCK], True, tl.int1)
    x0 = xindex % 32
    x1 = xindex // 32 % 8
    x2 = xindex // 256 % 4
    x3 = xindex // 1024
    x6 = xindex
    tmp0 = (x1 + 8 * x0) // 16
    tmp1 = tmp0.to(tl.float32)
    tmp2 = 0.25
    tmp3 = tmp1 * tmp2
    tmp4 = tmp3.to(tl.int32)
    tmp5 = (x1 + 8 * x0) % 16
    tmp6 = tmp5.to(tl.float32)
    tmp7 = tmp6 * tmp2
    tmp8 = tmp7.to(tl.int32)
    tmp9 = tl.load(in_ptr0 + (tmp8 + 4 * tmp4 + 16 * x2 + 16 * ((x1 + 8 *
        x0) // 256) + 64 * x3 + 64 * ((x1 + 8 * x0 + 256 * x2) // 1024)),
        None, eviction_policy='evict_last')
    tl.store(in_out_ptr0 + x6, tmp9, None)


def call(args):
    arg0_1, = args
    args.clear()
    assert_size_stride(arg0_1, (4, 4, 4, 4), (64, 16, 4, 1))
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        buf0 = empty_strided_cuda((4, 4, 8, 32), (1024, 256, 32, 1), torch.
            float32)
        buf1 = reinterpret_tensor(buf0, (4, 512, 2), (1024, 2, 1), 0)
        del buf0
        get_raw_stream(0)
        triton_poi_fused_clone_view_0[grid(4096)](buf1, arg0_1, 4096,
            XBLOCK=128, num_warps=4, num_stages=1)
        del arg0_1
    return buf1,


class SqueezeLayer(nn.Module):

    def __init__(self, factor):
        super(SqueezeLayer, self).__init__()
        self.factor = factor

    def forward(self, input, logdet=None, reverse=False, **kwargs):
        if not reverse:
            assert input.size(-1) % self.factor == 0
            output = input.view(input.size(0), input.size(1), -1, self.factor)
            output = output.permute(0, 1, 3, 2).contiguous()
            output = output.view(input.size(0), -1, input.size(-1) // self.
                factor)
            return output, logdet
        else:
            assert input.size(1) % self.factor == 0
            output = input.view(input.size(0), -1, self.factor, input.size(-1))
            output = output.permute(0, 1, 3, 2).contiguous()
            output = output.view(input.size(0), input.size(1) // self.
                factor, -1)
            return output, logdet


class UpsampleNetNew(nn.Module):

    def __init__(self, upsample_factor, upsample_method='duplicate',
        squeeze_factor=8):
        super(UpsampleNetNew, self).__init__()
        self.upsample_factor = upsample_factor
        self.upsample_method = upsample_method
        self.squeeze_factor = squeeze_factor
        if upsample_method == 'duplicate':
            upsample_factor = int(np.prod(upsample_factor))
            self.upsample = nn.Upsample(scale_factor=upsample_factor, mode=
                'nearest')
        elif upsample_method == 'transposed_conv2d':
            if not isinstance(upsample_factor, list):
                raise ValueError(
                    'You must specify upsample_factor as a list when used with transposed_conv2d'
                    )
            freq_axis_kernel_size = 3
            self.upsample_conv = nn.ModuleList()
            for s in upsample_factor:
                freq_axis_padding = (freq_axis_kernel_size - 1) // 2
                conv = nn.ConvTranspose2d(1, 1, (freq_axis_kernel_size, 2 *
                    s), padding=(freq_axis_padding, s // 2), dilation=1,
                    stride=(1, s))
                self.upsample_conv.append(conv)
                self.upsample_conv.append(nn.LeakyReLU(negative_slope=0.4,
                    inplace=True))
        else:
            raise ValueError('{} upsampling is not supported'.format(self.
                _upsample_method))
        self.squeeze_layer = SqueezeLayer(squeeze_factor)

    def forward(self, input_0):
        arg0_1 = input_0
        output = call([arg0_1])
        return output[0]


Measured latency_ms: 0.014987