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 math
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_add_0(in_ptr0, in_ptr1, in_ptr2, in_ptr3, out_ptr0,
    xnumel, XBLOCK: tl.constexpr):
    xnumel = 256
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x4 = xindex
    x0 = xindex % 4
    x3 = xindex // 64
    x2 = xindex // 16 % 4
    x7 = xindex % 16
    tmp0 = tl.load(in_ptr0 + x4, xmask)
    tmp1 = tl.load(in_ptr1 + (x0 + 4 * x3), xmask, eviction_policy='evict_last'
        )
    tmp3 = tl.load(in_ptr2 + (x0 + 4 * x2), xmask, eviction_policy='evict_last'
        )
    tmp5 = tl.load(in_ptr3 + x7, xmask, eviction_policy='evict_last')
    tmp2 = tmp0 + tmp1
    tmp4 = tmp2 + tmp3
    tmp6 = tmp4 + tmp5
    tl.store(out_ptr0 + x4, tmp6, xmask)


def call(args):
    primals_1, primals_2, primals_3, primals_4 = args
    args.clear()
    assert_size_stride(primals_1, (1, 4, 1, 1, 4), (16, 4, 4, 4, 1))
    assert_size_stride(primals_2, (4, 4, 4, 4), (64, 16, 4, 1))
    assert_size_stride(primals_3, (1, 1, 4, 1, 4), (16, 16, 4, 4, 1))
    assert_size_stride(primals_4, (1, 1, 1, 4, 4), (16, 16, 16, 4, 1))
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        buf0 = empty_strided_cuda((1, 4, 4, 4, 4), (256, 64, 16, 4, 1),
            torch.float32)
        get_raw_stream(0)
        triton_poi_fused_add_0[grid(256)](primals_2, primals_1, primals_3,
            primals_4, buf0, 256, XBLOCK=128, num_warps=4, num_stages=1)
        del primals_1
        del primals_2
        del primals_3
        del primals_4
    return buf0,


class PositionalEncodingImageNew(nn.Module):
    """
    Learning positional embeddings for images.
    Embeddings for channel, height and width are added to form the full positional embedding.
    These encodings correspond to the ones from Sparse Transformers (https://arxiv.org/abs/1904.10509).

    Args:
        image_shape: Iterable, the shape of the image.
        embedding_dim: int, the size of each embedding vector.
    """

    def __init__(self, image_shape, embedding_dim):
        super(PositionalEncodingImageNew, self).__init__()
        assert len(image_shape
            ) == 3, 'image_shape should have length 3: (C,H,W)'
        self.image_shape = image_shape
        self.embedding_dim = embedding_dim
        c, h, w = image_shape
        self.encode_c = nn.Parameter(torch.Tensor(1, c, 1, 1, embedding_dim))
        self.encode_h = nn.Parameter(torch.Tensor(1, 1, h, 1, embedding_dim))
        self.encode_w = nn.Parameter(torch.Tensor(1, 1, 1, w, embedding_dim))
        self.reset_parameters()

    def reset_parameters(self):
        nn.init.normal_(self.encode_c, std=0.125 / math.sqrt(3 * self.
            embedding_dim))
        nn.init.normal_(self.encode_h, std=0.125 / math.sqrt(3 * self.
            embedding_dim))
        nn.init.normal_(self.encode_w, std=0.125 / math.sqrt(3 * self.
            embedding_dim))

    def forward(self, input_0):
        primals_1 = self.encode_c
        primals_3 = self.encode_h
        primals_4 = self.encode_w
        primals_2 = input_0
        output = call([primals_1, primals_2, primals_3, primals_4])
        return output[0]


Measured latency_ms: 0.014364