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._inductor.runtime.triton_helpers import math as tl_math
import torch.utils.data
import torch.nn as nn
from torch.nn.modules.utils import _pair
from torch.nn.modules.utils import _quadruple
from torch import optim as optim
import torch.nn.parallel
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_0(in_ptr0, out_ptr0, ynumel, xnumel, YBLOCK: tl.
    constexpr, XBLOCK: tl.constexpr):
    ynumel = 64
    xnumel = 9
    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
    x3 = xindex % 3
    x4 = xindex // 3
    y0 = yindex % 2
    y1 = yindex // 2 % 2
    y2 = yindex // 4
    x6 = xindex
    y5 = yindex
    tmp0 = tl.load(in_ptr0 + (15 + -1 * tl_math.abs(-3 + x3 + y0) + -4 *
        tl_math.abs(-3 + x4 + y1) + 16 * y2), xmask & ymask,
        eviction_policy='evict_last')
    tl.store(out_ptr0 + (x6 + 9 * y5), tmp0, xmask & ymask)


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, 2, 2, 3, 3), (144, 36, 18, 9, 3, 1
            ), torch.float32)
        get_raw_stream(0)
        triton_poi_fused_clone_0[grid(64, 9)](arg0_1, buf0, 64, 9, XBLOCK=
            16, YBLOCK=64, num_warps=4, num_stages=1)
        del arg0_1
        buf1 = torch.ops.aten.median.dim(reinterpret_tensor(buf0, (4, 4, 2,
            2, 9), (144, 36, 18, 9, 1), 0), -1)
        del buf0
        buf2 = buf1[0]
        del buf1
    return buf2,


class MedianPool2dNew(nn.Module):
    """ Median pool (usable as median filter when stride=1) module.

    Args:
         kernel_size: size of pooling kernel, int or 2-tuple
         stride: pool stride, int or 2-tuple
         padding: pool padding, int or 4-tuple (l, r, t, b) as in pytorch F.pad
         same: override padding and enforce same padding, boolean
    """

    def __init__(self, kernel_size=3, stride=1, padding=0, same=False):
        super(MedianPool2dNew, self).__init__()
        self.k = _pair(kernel_size)
        self.stride = _pair(stride)
        self.padding = _quadruple(padding)
        self.same = same

    def _padding(self, x):
        if self.same:
            ih, iw = x.size()[2:]
            if ih % self.stride[0] == 0:
                ph = max(self.k[0] - self.stride[0], 0)
            else:
                ph = max(self.k[0] - ih % self.stride[0], 0)
            if iw % self.stride[1] == 0:
                pw = max(self.k[1] - self.stride[1], 0)
            else:
                pw = max(self.k[1] - iw % self.stride[1], 0)
            pl = pw // 2
            pr = pw - pl
            pt = ph // 2
            pb = ph - pt
            padding = pl, pr, pt, pb
        else:
            padding = self.padding
        return padding

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


Measured latency_ms: 0.014628