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 math
from typing import List
from typing import Optional
from typing import Union
from typing import Any
from typing import Tuple
from typing import NamedTuple
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


@triton.jit
def triton_poi_fused_convolution_leaky_relu_0(in_ptr0, in_ptr1, out_ptr0,
    out_ptr1, xnumel, XBLOCK: tl.constexpr):
    xnumel = 256
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x3 = xindex
    x1 = xindex // 16 % 4
    tmp0 = tl.load(in_ptr0 + x3, xmask)
    tmp1 = tl.load(in_ptr1 + x1, xmask, eviction_policy='evict_last')
    tmp2 = tmp0 + tmp1
    tmp3 = 0.0
    tmp4 = tmp2 > tmp3
    tmp5 = 0.01
    tmp6 = tmp2 * tmp5
    tmp7 = tl.where(tmp4, tmp2, tmp6)
    tl.store(out_ptr0 + x3, tmp4, xmask)
    tl.store(out_ptr1 + x3, tmp7, xmask)


def call(args):
    primals_1, primals_2, primals_3 = args
    args.clear()
    assert_size_stride(primals_1, (4, 4, 4, 4), (64, 16, 4, 1))
    assert_size_stride(primals_2, (4, 4, 3, 3), (36, 9, 3, 1))
    assert_size_stride(primals_3, (4,), (1,))
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        buf0 = extern_kernels.convolution(primals_1, primals_2, stride=(1, 
            1), padding=(1, 1), dilation=(1, 1), transposed=False,
            output_padding=(0, 0), groups=1, bias=None)
        assert_size_stride(buf0, (4, 4, 4, 4), (64, 16, 4, 1))
        buf1 = empty_strided_cuda((4, 4, 4, 4), (64, 16, 4, 1), torch.bool)
        buf2 = empty_strided_cuda((4, 4, 4, 4), (64, 16, 4, 1), torch.float32)
        get_raw_stream(0)
        triton_poi_fused_convolution_leaky_relu_0[grid(256)](buf0,
            primals_3, buf1, buf2, 256, XBLOCK=128, num_warps=4, num_stages=1)
        del buf0
        del primals_3
    return buf2, primals_1, primals_2, buf1


class PaddedTensor(NamedTuple):
    data: 'torch.Tensor'
    sizes: 'torch.Tensor'

    @classmethod
    def build(cls, data: 'torch.Tensor', sizes: 'torch.Tensor'):
        assert isinstance(data, torch.Tensor)
        assert isinstance(sizes, torch.Tensor)
        assert sizes.dim() == 2, 'PaddedTensor.sizes must have 2 dimensions'
        assert sizes.size(1) in (2, 3
            ), f'PaddedTensor.sizes is incorrect: expected=2 (HxW) or 3 (CxHxW), found={sizes.size(1)}'
        assert data.size(0) == sizes.size(0
            ), f'Batch size {sizes.size(0)} does not match the number of samples in the batch {data.size(0)}'
        return cls(data, sizes)

    def __repr__(self) ->str:
        return (
            f'PaddedTensor(data.size()={list(self.data.size())}, sizes={self.sizes.tolist()}, device={str(self.data.device)})'
            )

    @property
    def device(self) ->torch.device:
        return self.data.device


class ConvBlockNew(nn.Module):

    def __init__(self, in_channels: 'int', out_channels: 'int', kernel_size:
        'Param2d'=3, stride: 'Param2d'=1, dilation: 'Param2d'=1, activation:
        'Optional[nn.Module]'=nn.LeakyReLU, poolsize: 'Param2d'=0, dropout:
        'Optional[float]'=None, batchnorm: 'bool'=False, inplace: 'bool'=
        False, use_masks: 'bool'=False) ->None:
        super().__init__()
        ks, st, di, ps = ConvBlockNew.prepare_dimensional_args(kernel_size,
            stride, dilation, poolsize)
        if ps[0] * ps[1] < 2:
            ps = None
        self.dropout = dropout
        self.in_channels = in_channels
        self.use_masks = use_masks
        self.poolsize = ps
        self.conv = nn.Conv2d(in_channels, out_channels, ks, stride=st,
            padding=tuple((ks[dim] - 1) // 2 * di[dim] for dim in (0, 1)),
            dilation=di, bias=not batchnorm)
        self.batchnorm = nn.BatchNorm2d(out_channels) if batchnorm else None
        self.activation = activation(inplace=inplace) if activation else None
        self.pool = nn.MaxPool2d(ps) if self.poolsize else None

    @staticmethod
    def prepare_dimensional_args(*args: Any, dims: int=2) ->List[Tuple]:
        return [(tuple(arg) if isinstance(arg, (list, tuple)) else (arg,) *
            dims) for arg in args]

    def get_batch_output_size(self, xs: 'torch.Tensor') ->torch.Tensor:
        ys = torch.zeros_like(xs)
        for dim in (0, 1):
            ys[:, dim] = self.get_output_size(size=xs[:, dim], kernel_size=
                self.conv.kernel_size[dim], dilation=self.conv.dilation[dim
                ], stride=self.conv.stride[dim], poolsize=self.poolsize[dim
                ] if self.poolsize else None, padding=self.conv.padding[dim])
        return ys

    @staticmethod
    def get_output_size(size: 'Union[torch.Tensor, int]', kernel_size:
        'int', dilation: 'int', stride: 'int', poolsize: 'int', padding:
        'Optional[int]'=None) ->Union[torch.LongTensor, int]:
        if padding is None:
            padding = (kernel_size - 1) // 2 * dilation
        size = size.float() if isinstance(size, torch.Tensor) else float(size)
        size = (size + 2 * padding - dilation * (kernel_size - 1) - 1
            ) / stride + 1
        size = size.floor() if isinstance(size, torch.Tensor) else math.floor(
            size)
        if poolsize:
            size /= poolsize
        return size.floor().long() if isinstance(size, torch.Tensor) else int(
            math.floor(size))

    def forward(self, input_0):
        primals_2 = self.conv.weight
        primals_3 = self.conv.bias
        primals_1 = input_0
        output = call([primals_1, primals_2, primals_3])
        return output[0]


Measured latency_ms: 0.024678