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.nn 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_poi_fused_mul_sum_0(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
    x0 = xindex
    tmp0 = tl.load(in_ptr0 + 16 * x0, xmask, eviction_policy='evict_last')
    tmp1 = tl.load(in_ptr0 + (4 + 16 * x0), xmask, eviction_policy='evict_last'
        )
    tmp3 = tl.load(in_ptr0 + (8 + 16 * x0), xmask, eviction_policy='evict_last'
        )
    tmp5 = tl.load(in_ptr0 + (12 + 16 * x0), xmask, eviction_policy=
        'evict_last')
    tmp7 = tl.load(in_ptr0 + (1 + 16 * x0), xmask, eviction_policy='evict_last'
        )
    tmp8 = tl.load(in_ptr0 + (5 + 16 * x0), xmask, eviction_policy='evict_last'
        )
    tmp10 = tl.load(in_ptr0 + (9 + 16 * x0), xmask, eviction_policy=
        'evict_last')
    tmp12 = tl.load(in_ptr0 + (13 + 16 * x0), xmask, eviction_policy=
        'evict_last')
    tmp15 = tl.load(in_ptr0 + (2 + 16 * x0), xmask, eviction_policy=
        'evict_last')
    tmp16 = tl.load(in_ptr0 + (6 + 16 * x0), xmask, eviction_policy=
        'evict_last')
    tmp18 = tl.load(in_ptr0 + (10 + 16 * x0), xmask, eviction_policy=
        'evict_last')
    tmp20 = tl.load(in_ptr0 + (14 + 16 * x0), xmask, eviction_policy=
        'evict_last')
    tmp23 = tl.load(in_ptr0 + (3 + 16 * x0), xmask, eviction_policy=
        'evict_last')
    tmp24 = tl.load(in_ptr0 + (7 + 16 * x0), xmask, eviction_policy=
        'evict_last')
    tmp26 = tl.load(in_ptr0 + (11 + 16 * x0), xmask, eviction_policy=
        'evict_last')
    tmp28 = tl.load(in_ptr0 + (15 + 16 * x0), xmask, eviction_policy=
        'evict_last')
    tmp2 = tmp0 + tmp1
    tmp4 = tmp2 + tmp3
    tmp6 = tmp4 + tmp5
    tmp9 = tmp7 + tmp8
    tmp11 = tmp9 + tmp10
    tmp13 = tmp11 + tmp12
    tmp14 = tmp6 + tmp13
    tmp17 = tmp15 + tmp16
    tmp19 = tmp17 + tmp18
    tmp21 = tmp19 + tmp20
    tmp22 = tmp14 + tmp21
    tmp25 = tmp23 + tmp24
    tmp27 = tmp25 + tmp26
    tmp29 = tmp27 + tmp28
    tmp30 = tmp22 + tmp29
    tmp31 = 0.0625
    tmp32 = tmp30 * tmp31
    tl.store(in_out_ptr0 + x0, tmp32, xmask)


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), (4, 1), torch.float32)
        buf1 = buf0
        del buf0
        get_raw_stream(0)
        triton_poi_fused_mul_sum_0[grid(16)](buf1, arg0_1, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        del arg0_1
    return reinterpret_tensor(buf1, (4, 4, 1, 1), (4, 1, 1, 1), 0),


class AvgPool2dNew(Module):
    """
    This class is the beginning of an exact python port of the torch.nn.AvgPool2d
    module. Because PySyft cannot hook into layers which are implemented in C++,
    our special functionalities (such as encrypted computation) do not work with
    torch.nn.AvgPool2d and so we must have python ports available for all layer types
    which we seek to use.

    Note that this module has been tested to ensure that it outputs the exact output
    values that the main module outputs in the same order that the main module does.

    However, there is often some rounding error of unknown origin, usually less than
    1e-6 in magnitude.

    This module has not yet been tested with GPUs but should work out of the box.
    """

    def __init__(self, kernel_size, stride=None, padding=0, ceil_mode=False,
        count_include_pad=True, divisor_override=None):
        """For information on the constructor arguments, please see PyTorch's
        documentation in torch.nn.AvgPool2d"""
        super().__init__()
        assert padding == 0
        assert ceil_mode is False
        assert count_include_pad is True
        assert divisor_override is None
        if stride is None:
            stride = kernel_size
        self.kernel_size = kernel_size
        self.stride = stride
        self.padding = padding
        self.ceil_mode = ceil_mode
        self.count_include_pad = count_include_pad
        self.divisor_override = divisor_override
        self._one_over_kernel_size = 1 / (self.kernel_size * self.kernel_size)

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


Measured latency_ms: 0.014748