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 torch.nn as nn
from torch.nn.parameter import Parameter
import torch.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_0(in_ptr0, out_ptr0, ynumel, xnumel,
    YBLOCK: tl.constexpr, XBLOCK: tl.constexpr):
    ynumel = 16
    xnumel = 16
    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
    x2 = xindex
    y3 = yindex
    y0 = yindex % 4
    y1 = yindex // 4
    tmp0 = tl.load(in_ptr0 + (x2 + 16 * y3), xmask & ymask)
    tl.store(out_ptr0 + (y0 + 4 * x2 + 64 * y1), tmp0, xmask & ymask)


def call(args):
    arg0_1, arg1_1 = args
    args.clear()
    assert_size_stride(arg0_1, (4, 4, 4, 4), (64, 16, 4, 1))
    assert_size_stride(arg1_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, 4), (64, 1, 16, 4), torch.float32)
        get_raw_stream(0)
        triton_poi_fused_convolution_0[grid(16, 16)](arg1_1, buf0, 16, 16,
            XBLOCK=16, YBLOCK=16, num_warps=4, num_stages=1)
        del arg1_1
        buf1 = empty_strided_cuda((4, 4, 4, 4), (64, 1, 16, 4), torch.float32)
        triton_poi_fused_convolution_0[grid(16, 16)](arg0_1, buf1, 16, 16,
            XBLOCK=16, YBLOCK=16, num_warps=4, num_stages=1)
        del arg0_1
        buf2 = extern_kernels.convolution(buf0, buf1, stride=(1, 1),
            padding=(0, 0), dilation=(1, 1), transposed=False,
            output_padding=(0, 0), groups=1, bias=None)
        assert_size_stride(buf2, (4, 4, 1, 1), (4, 1, 4, 4))
        del buf0
        del buf1
    return buf2,


def to_pair(data):
    """Converts a single or a tuple of data into a pair. If the data is a tuple with more than two elements, it selects
	the first two of them. In case of single data, it duplicates that data into a pair.

	Args:
		data (object or tuple): The input data.

	Returns:
		Tuple: A pair of data.
	"""
    if isinstance(data, tuple):
        return data[0:2]
    return data, data


class ConvolutionNew(nn.Module):
    """Performs a 2D convolution over an input spike-wave composed of several input
	planes. Current version only supports stride of 1 with no padding.

	The input is a 4D tensor with the size :math:`(T, C_{{in}}, H_{{in}}, W_{{in}})` and the crresponsing output
	is of size :math:`(T, C_{{out}}, H_{{out}}, W_{{out}})`, 
	where :math:`T` is the number of time steps, :math:`C` is the number of feature maps (channels), and
	:math:`H`, and :math:`W` are the hight and width of the input/output planes.

	* :attr:`in_channels` controls the number of input planes (channels/feature maps).

	* :attr:`out_channels` controls the number of feature maps in the current layer.

	* :attr:`kernel_size` controls the size of the convolution kernel. It can be a single integer or a tuple of two integers.

	* :attr:`weight_mean` controls the mean of the normal distribution used for initial random weights.

	* :attr:`weight_std` controls the standard deviation of the normal distribution used for initial random weights.

	.. note::

		Since this version of convolution does not support padding, it is the user responsibility to add proper padding
		on the input before applying convolution.

	Args:
		in_channels (int): Number of channels in the input.
		out_channels (int): Number of channels produced by the convolution.
		kernel_size (int or tuple): Size of the convolving kernel.
		weight_mean (float, optional): Mean of the initial random weights. Default: 0.8
		weight_std (float, optional): Standard deviation of the initial random weights. Default: 0.02
	"""

    def __init__(self, in_channels, out_channels, kernel_size, weight_mean=
        0.8, weight_std=0.02):
        super(ConvolutionNew, self).__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = to_pair(kernel_size)
        self.stride = 1
        self.bias = None
        self.dilation = 1
        self.groups = 1
        self.padding = 0
        self.weight = Parameter(torch.Tensor(self.out_channels, self.
            in_channels, *self.kernel_size))
        self.weight.requires_grad_(False)
        self.reset_weight(weight_mean, weight_std)

    def reset_weight(self, weight_mean=0.8, weight_std=0.02):
        """Resets weights to random values based on a normal distribution.

		Args:
			weight_mean (float, optional): Mean of the random weights. Default: 0.8
			weight_std (float, optional): Standard deviation of the random weights. Default: 0.02
		"""
        self.weight.normal_(weight_mean, weight_std)

    def load_weight(self, target):
        """Loads weights with the target tensor.

		Args:
			target (Tensor=): The target tensor.
		"""
        self.weight.copy_(target)

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


Measured latency_ms: 0.037645