# AOT ID: ['0_forward']
from ctypes import c_void_p, c_long, c_int
import torch
import math
import random
import os
import tempfile
from math import inf, nan
from torch._inductor.hooks import run_intermediate_hooks
from torch._inductor.utils import maybe_profile
from torch._inductor.codegen.memory_planning import _align as align
from torch import device, empty_strided
from torch._inductor.async_compile import AsyncCompile
from torch._inductor.select_algorithm import extern_kernels
from torch._inductor.codegen.multi_kernel import MultiKernelCall
import triton
import triton.language as tl
from torch._inductor.runtime.triton_heuristics import grid, split_scan_grid, grid_combo_kernels, start_graph, end_graph
from torch._C import _cuda_getCurrentRawStream as get_raw_stream

aten = torch.ops.aten
inductor_ops = torch.ops.inductor
_quantized = torch.ops._quantized
assert_size_stride = torch._C._dynamo.guards.assert_size_stride
empty_strided_cpu = torch._C._dynamo.guards._empty_strided_cpu
empty_strided_cuda = torch._C._dynamo.guards._empty_strided_cuda
empty_strided_xpu = torch._C._dynamo.guards._empty_strided_xpu
reinterpret_tensor = torch._C._dynamo.guards._reinterpret_tensor
alloc_from_pool = torch.ops.inductor._alloc_from_pool
async_compile = AsyncCompile()


# kernel path: runs/run_shard_7/inductor_cache/gy/cgyfxbdoyn6aqocgkacshriylcks5c6kt3gfzmwpealzqpvxe2qq.py
# Topologically Sorted Source Nodes: [contiguous_1], Original ATen: [aten.clone]
# Source node to ATen node mapping:
#   contiguous_1 => clone_1
# Graph fragment:
#   %clone_1 : [num_users=1] = call_function[target=torch.ops.aten.clone.default](args = (%permute_4,), kwargs = {memory_format: torch.contiguous_format})
triton_poi_fused_clone_0 = async_compile.triton('triton_poi_fused_clone_0', '''
import triton
import triton.language as tl
from triton.compiler.compiler import AttrsDescriptor

from torch._inductor.runtime import triton_helpers, triton_heuristics
from torch._inductor.runtime.triton_helpers import libdevice, math as tl_math
from torch._inductor.runtime.hints import AutotuneHint, ReductionHint, TileHint, instance_descriptor, DeviceProperties

@triton_heuristics.pointwise(
    size_hints=[512], 
    filename=__file__,
    triton_meta={'signature': {0: '*fp32', 1: '*fp32', 2: 'i32'}, 'device': DeviceProperties(type='cuda', index=0, cc=80, major=8, regs_per_multiprocessor=65536, max_threads_per_multi_processor=2048, multi_processor_count=108), 'constants': {}, 'configs': [AttrsDescriptor(divisible_by_16=(0, 1, 2), equal_to_1=())]},
    inductor_meta={'autotune_hints': set(), 'kernel_name': 'triton_poi_fused_clone_0', 'mutated_arg_names': [], 'no_x_dim': False, 'num_load': 1, 'num_reduction': 0, 'backend_hash': 'A9C866B4A14FD3277824029365D703C2427B2E685E54EC9B3EF4ADC8D1EEAC1D', 'are_deterministic_algorithms_enabled': False, 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': None, 'force_disable_caches': False, 'dynamic_scale_rblock': True, 'max_autotune': False, 'max_autotune_pointwise': False, 'min_split_scan_rblock': 256, 'spill_threshold': 16, 'store_cubin': False},
    min_elem_per_thread=0
)
@triton.jit
def triton_poi_fused_clone_0(in_ptr0, out_ptr0, xnumel, XBLOCK : tl.constexpr):
    xnumel = 512
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex % 4
    x1 = (xindex // 4) % 4
    x2 = (xindex // 16) % 8
    x3 = (xindex // 128)
    x4 = xindex
    tmp0 = tl.load(in_ptr0 + (x0 + (4*x2) + (32*x1) + (128*x3)), xmask)
    tl.store(out_ptr0 + (x4), tmp0, xmask)
''', device_str='cuda')


# kernel path: runs/run_shard_7/inductor_cache/sd/csdzc7ihn23c3cq4w4usapmdgqsgbmeydp7obuzkly4qrdi7ghk5.py
# Topologically Sorted Source Nodes: [weights], Original ATen: [aten._softmax]
# Source node to ATen node mapping:
#   weights => exp
# Graph fragment:
#   %mul_tensor : [num_users=2] = call_function[target=torch.ops.aten.mul.Tensor](args = (%bmm, 1), kwargs = {})
#   %amax_default : [num_users=1] = call_function[target=torch.ops.aten.amax.default](args = (%mul_tensor, [2], True), kwargs = {})
#   %sub_tensor : [num_users=1] = call_function[target=torch.ops.aten.sub.Tensor](args = (%mul_tensor, %amax_default), kwargs = {})
#   %div_tensor : [num_users=1] = call_function[target=torch.ops.aten.div.Tensor](args = (%sub_tensor, 2.0), kwargs = {})
#   %exp : [num_users=2] = call_function[target=torch.ops.aten.exp.default](args = (%div_tensor,), kwargs = {})
triton_poi_fused__softmax_1 = async_compile.triton('triton_poi_fused__softmax_1', '''
import triton
import triton.language as tl
from triton.compiler.compiler import AttrsDescriptor

from torch._inductor.runtime import triton_helpers, triton_heuristics
from torch._inductor.runtime.triton_helpers import libdevice, math as tl_math
from torch._inductor.runtime.hints import AutotuneHint, ReductionHint, TileHint, instance_descriptor, DeviceProperties

@triton_heuristics.pointwise(
    size_hints=[512], 
    filename=__file__,
    triton_meta={'signature': {0: '*fp32', 1: '*fp32', 2: 'i32'}, 'device': DeviceProperties(type='cuda', index=0, cc=80, major=8, regs_per_multiprocessor=65536, max_threads_per_multi_processor=2048, multi_processor_count=108), 'constants': {}, 'configs': [AttrsDescriptor(divisible_by_16=(0, 1, 2), equal_to_1=())]},
    inductor_meta={'autotune_hints': set(), 'kernel_name': 'triton_poi_fused__softmax_1', 'mutated_arg_names': [], 'no_x_dim': False, 'num_load': 5, 'num_reduction': 0, 'backend_hash': 'A9C866B4A14FD3277824029365D703C2427B2E685E54EC9B3EF4ADC8D1EEAC1D', 'are_deterministic_algorithms_enabled': False, 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': None, 'force_disable_caches': False, 'dynamic_scale_rblock': True, 'max_autotune': False, 'max_autotune_pointwise': False, 'min_split_scan_rblock': 256, 'spill_threshold': 16, 'store_cubin': False},
    min_elem_per_thread=0
)
@triton.jit
def triton_poi_fused__softmax_1(in_ptr0, out_ptr0, xnumel, XBLOCK : tl.constexpr):
    xnumel = 512
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x2 = xindex
    x1 = (xindex // 4)
    tmp0 = tl.load(in_ptr0 + (x2), xmask)
    tmp3 = tl.load(in_ptr0 + (4*x1), xmask, eviction_policy='evict_last')
    tmp5 = tl.load(in_ptr0 + (1 + (4*x1)), xmask, eviction_policy='evict_last')
    tmp8 = tl.load(in_ptr0 + (2 + (4*x1)), xmask, eviction_policy='evict_last')
    tmp11 = tl.load(in_ptr0 + (3 + (4*x1)), xmask, eviction_policy='evict_last')
    tmp1 = 1.0
    tmp2 = tmp0 * tmp1
    tmp4 = tmp3 * tmp1
    tmp6 = tmp5 * tmp1
    tmp7 = triton_helpers.maximum(tmp4, tmp6)
    tmp9 = tmp8 * tmp1
    tmp10 = triton_helpers.maximum(tmp7, tmp9)
    tmp12 = tmp11 * tmp1
    tmp13 = triton_helpers.maximum(tmp10, tmp12)
    tmp14 = tmp2 - tmp13
    tmp15 = 0.5
    tmp16 = tmp14 * tmp15
    tmp17 = tl_math.exp(tmp16)
    tl.store(out_ptr0 + (x2), tmp17, xmask)
''', device_str='cuda')


# kernel path: runs/run_shard_7/inductor_cache/nr/cnr2l3l4x36tjfbtrutwg2z6iijamnnvubvppho2khkq56fs453d.py
# Topologically Sorted Source Nodes: [weights], Original ATen: [aten._softmax]
# Source node to ATen node mapping:
#   weights => div_1, sum_1
# Graph fragment:
#   %sum_1 : [num_users=1] = call_function[target=torch.ops.aten.sum.dim_IntList](args = (%exp, [2], True), kwargs = {})
#   %div_1 : [num_users=2] = call_function[target=torch.ops.aten.div.Tensor](args = (%exp, %sum_1), kwargs = {})
triton_poi_fused__softmax_2 = async_compile.triton('triton_poi_fused__softmax_2', '''
import triton
import triton.language as tl
from triton.compiler.compiler import AttrsDescriptor

from torch._inductor.runtime import triton_helpers, triton_heuristics
from torch._inductor.runtime.triton_helpers import libdevice, math as tl_math
from torch._inductor.runtime.hints import AutotuneHint, ReductionHint, TileHint, instance_descriptor, DeviceProperties

@triton_heuristics.pointwise(
    size_hints=[512], 
    filename=__file__,
    triton_meta={'signature': {0: '*fp32', 1: '*fp32', 2: 'i32'}, 'device': DeviceProperties(type='cuda', index=0, cc=80, major=8, regs_per_multiprocessor=65536, max_threads_per_multi_processor=2048, multi_processor_count=108), 'constants': {}, 'configs': [AttrsDescriptor(divisible_by_16=(0, 1, 2), equal_to_1=())]},
    inductor_meta={'autotune_hints': set(), 'kernel_name': 'triton_poi_fused__softmax_2', 'mutated_arg_names': [], 'no_x_dim': False, 'num_load': 5, 'num_reduction': 0, 'backend_hash': 'A9C866B4A14FD3277824029365D703C2427B2E685E54EC9B3EF4ADC8D1EEAC1D', 'are_deterministic_algorithms_enabled': False, 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': None, 'force_disable_caches': False, 'dynamic_scale_rblock': True, 'max_autotune': False, 'max_autotune_pointwise': False, 'min_split_scan_rblock': 256, 'spill_threshold': 16, 'store_cubin': False},
    min_elem_per_thread=0
)
@triton.jit
def triton_poi_fused__softmax_2(in_ptr0, out_ptr0, xnumel, XBLOCK : tl.constexpr):
    xnumel = 512
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x2 = xindex
    x1 = (xindex // 4)
    tmp0 = tl.load(in_ptr0 + (x2), xmask)
    tmp1 = tl.load(in_ptr0 + (4*x1), xmask, eviction_policy='evict_last')
    tmp2 = tl.load(in_ptr0 + (1 + (4*x1)), xmask, eviction_policy='evict_last')
    tmp4 = tl.load(in_ptr0 + (2 + (4*x1)), xmask, eviction_policy='evict_last')
    tmp6 = tl.load(in_ptr0 + (3 + (4*x1)), xmask, eviction_policy='evict_last')
    tmp3 = tmp1 + tmp2
    tmp5 = tmp3 + tmp4
    tmp7 = tmp5 + tmp6
    tmp8 = tmp0 / tmp7
    tl.store(out_ptr0 + (x2), tmp8, xmask)
''', device_str='cuda')


# kernel path: runs/run_shard_7/inductor_cache/wf/cwfnu7iigx37txvjlhi5dkscdizq3ojcwgdznc7shmckykzg5ba5.py
# Topologically Sorted Source Nodes: [contiguous_3], Original ATen: [aten.clone]
# Source node to ATen node mapping:
#   contiguous_3 => clone_3
# Graph fragment:
#   %clone_3 : [num_users=1] = call_function[target=torch.ops.aten.clone.default](args = (%permute_7,), kwargs = {memory_format: torch.contiguous_format})
triton_poi_fused_clone_3 = async_compile.triton('triton_poi_fused_clone_3', '''
import triton
import triton.language as tl
from triton.compiler.compiler import AttrsDescriptor

from torch._inductor.runtime import triton_helpers, triton_heuristics
from torch._inductor.runtime.triton_helpers import libdevice, math as tl_math
from torch._inductor.runtime.hints import AutotuneHint, ReductionHint, TileHint, instance_descriptor, DeviceProperties

@triton_heuristics.pointwise(
    size_hints=[512], 
    filename=__file__,
    triton_meta={'signature': {0: '*fp32', 1: '*fp32', 2: 'i32'}, 'device': DeviceProperties(type='cuda', index=0, cc=80, major=8, regs_per_multiprocessor=65536, max_threads_per_multi_processor=2048, multi_processor_count=108), 'constants': {}, 'configs': [AttrsDescriptor(divisible_by_16=(0, 1, 2), equal_to_1=())]},
    inductor_meta={'autotune_hints': set(), 'kernel_name': 'triton_poi_fused_clone_3', 'mutated_arg_names': [], 'no_x_dim': False, 'num_load': 1, 'num_reduction': 0, 'backend_hash': 'A9C866B4A14FD3277824029365D703C2427B2E685E54EC9B3EF4ADC8D1EEAC1D', 'are_deterministic_algorithms_enabled': False, 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': None, 'force_disable_caches': False, 'dynamic_scale_rblock': True, 'max_autotune': False, 'max_autotune_pointwise': False, 'min_split_scan_rblock': 256, 'spill_threshold': 16, 'store_cubin': False},
    min_elem_per_thread=0
)
@triton.jit
def triton_poi_fused_clone_3(in_ptr0, out_ptr0, xnumel, XBLOCK : tl.constexpr):
    xnumel = 512
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex % 4
    x1 = (xindex // 4) % 8
    x2 = (xindex // 32) % 4
    x3 = (xindex // 128)
    x4 = xindex
    tmp0 = tl.load(in_ptr0 + (x0 + (4*x2) + (16*x1) + (128*x3)), xmask)
    tl.store(out_ptr0 + (x4), tmp0, xmask)
''', device_str='cuda')


async_compile.wait(globals())
del async_compile

def call(args):
    primals_1, primals_2, primals_3, primals_4, primals_5, primals_6 = args
    args.clear()
    assert_size_stride(primals_1, (4, 4, 4), (16, 4, 1))
    assert_size_stride(primals_2, (32, 4), (4, 1))
    assert_size_stride(primals_3, (32, 4), (4, 1))
    assert_size_stride(primals_4, (32, 4), (4, 1))
    assert_size_stride(primals_5, (4, 32), (32, 1))
    assert_size_stride(primals_6, (4, ), (1, ))
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        buf0 = empty_strided_cuda((16, 32), (32, 1), torch.float32)
        # Topologically Sorted Source Nodes: [linear], Original ATen: [aten.mm]
        extern_kernels.mm(reinterpret_tensor(primals_1, (16, 4), (4, 1), 0), reinterpret_tensor(primals_2, (4, 32), (1, 4), 0), out=buf0)
        del primals_2
        buf1 = empty_strided_cuda((16, 32), (32, 1), torch.float32)
        # Topologically Sorted Source Nodes: [linear_1], Original ATen: [aten.mm]
        extern_kernels.mm(reinterpret_tensor(primals_1, (16, 4), (4, 1), 0), reinterpret_tensor(primals_3, (4, 32), (1, 4), 0), out=buf1)
        del primals_3
        buf2 = empty_strided_cuda((16, 32), (32, 1), torch.float32)
        # Topologically Sorted Source Nodes: [linear_2], Original ATen: [aten.mm]
        extern_kernels.mm(reinterpret_tensor(primals_1, (16, 4), (4, 1), 0), reinterpret_tensor(primals_4, (4, 32), (1, 4), 0), out=buf2)
        del primals_4
        buf3 = empty_strided_cuda((4, 8, 4, 4), (128, 16, 4, 1), torch.float32)
        # Topologically Sorted Source Nodes: [contiguous_1], Original ATen: [aten.clone]
        stream0 = get_raw_stream(0)
        triton_poi_fused_clone_0.run(buf1, buf3, 512, grid=grid(512), stream=stream0)
        buf4 = reinterpret_tensor(buf1, (4, 8, 4, 4), (128, 16, 4, 1), 0); del buf1  # reuse
        # Topologically Sorted Source Nodes: [contiguous], Original ATen: [aten.clone]
        triton_poi_fused_clone_0.run(buf0, buf4, 512, grid=grid(512), stream=stream0)
        buf5 = reinterpret_tensor(buf0, (32, 4, 4), (16, 4, 1), 0); del buf0  # reuse
        # Topologically Sorted Source Nodes: [raw_weights], Original ATen: [aten.bmm]
        extern_kernels.bmm(reinterpret_tensor(buf3, (32, 4, 4), (16, 4, 1), 0), reinterpret_tensor(buf4, (32, 4, 4), (16, 1, 4), 0), out=buf5)
        buf6 = empty_strided_cuda((32, 4, 4), (16, 4, 1), torch.float32)
        # Topologically Sorted Source Nodes: [weights], Original ATen: [aten._softmax]
        triton_poi_fused__softmax_1.run(buf5, buf6, 512, grid=grid(512), stream=stream0)
        buf7 = buf5; del buf5  # reuse
        # Topologically Sorted Source Nodes: [weights], Original ATen: [aten._softmax]
        triton_poi_fused__softmax_2.run(buf6, buf7, 512, grid=grid(512), stream=stream0)
        buf8 = reinterpret_tensor(buf6, (4, 8, 4, 4), (128, 16, 4, 1), 0); del buf6  # reuse
        # Topologically Sorted Source Nodes: [contiguous_2], Original ATen: [aten.clone]
        triton_poi_fused_clone_0.run(buf2, buf8, 512, grid=grid(512), stream=stream0)
        buf9 = reinterpret_tensor(buf2, (32, 4, 4), (16, 4, 1), 0); del buf2  # reuse
        # Topologically Sorted Source Nodes: [bmm_1], Original ATen: [aten.bmm]
        extern_kernels.bmm(buf7, reinterpret_tensor(buf8, (32, 4, 4), (16, 4, 1), 0), out=buf9)
        buf10 = empty_strided_cuda((4, 4, 8, 4), (128, 32, 4, 1), torch.float32)
        # Topologically Sorted Source Nodes: [contiguous_3], Original ATen: [aten.clone]
        triton_poi_fused_clone_3.run(buf9, buf10, 512, grid=grid(512), stream=stream0)
        del buf9
        buf11 = empty_strided_cuda((16, 4), (4, 1), torch.float32)
        # Topologically Sorted Source Nodes: [out_2], Original ATen: [aten.addmm]
        extern_kernels.addmm(primals_6, reinterpret_tensor(buf10, (16, 32), (32, 1), 0), reinterpret_tensor(primals_5, (32, 4), (1, 32), 0), alpha=1, beta=1, out=buf11)
        del primals_6
    return (reinterpret_tensor(buf11, (4, 4, 4), (16, 4, 1), 0), buf7, reinterpret_tensor(primals_1, (16, 4), (4, 1), 0), buf7, reinterpret_tensor(buf10, (16, 32), (32, 1), 0), primals_5, reinterpret_tensor(buf8, (32, 4, 4), (16, 1, 4), 0), reinterpret_tensor(buf3, (32, 4, 4), (16, 1, 4), 0), reinterpret_tensor(buf4, (32, 4, 4), (16, 4, 1), 0), )


def benchmark_compiled_module(times=10, repeat=10):
    from torch._dynamo.testing import rand_strided
    from torch._inductor.utils import print_performance
    primals_1 = rand_strided((4, 4, 4), (16, 4, 1), device='cuda:0', dtype=torch.float32)
    primals_2 = rand_strided((32, 4), (4, 1), device='cuda:0', dtype=torch.float32)
    primals_3 = rand_strided((32, 4), (4, 1), device='cuda:0', dtype=torch.float32)
    primals_4 = rand_strided((32, 4), (4, 1), device='cuda:0', dtype=torch.float32)
    primals_5 = rand_strided((4, 32), (32, 1), device='cuda:0', dtype=torch.float32)
    primals_6 = rand_strided((4, ), (1, ), device='cuda:0', dtype=torch.float32)
    fn = lambda: call([primals_1, primals_2, primals_3, primals_4, primals_5, primals_6])
    return print_performance(fn, times=times, repeat=repeat)


if __name__ == "__main__":
    from torch._inductor.wrapper_benchmark import compiled_module_main
    compiled_module_main('None', benchmark_compiled_module)

import math
import torch
from torch import nn
import torch.nn.functional as F


class MultiHeadAttention(nn.Module):

    def __init__(self, emb_size, n_heads=8, mask=False):
        """
        Arguments:

        emb_size: Size of input Embeddings
        n_heads: Number of heads for MultiHead Attention

        Layers:

        tokeys: For Keys
        toquery: For Query
        tovalue: For Value

        combine_heads: Convert the ( (number of heads) * (emb_size) ) into (emb_size) for output

        Note:
            The input dimension we give to the Linear Layer (nn.Linear) is the last dimension (no. of columns) in size
            (batch_size ,seq_length, input_dims) ==> input_dims : give to nn.Linear
            So in order for the tensor to pass through linear layer having the last dim equal to input_dim of
            Linear layer is necessary.

        """
        super(MultiHeadAttention, self).__init__()
        self.emb_size = emb_size
        self.n_heads = n_heads
        """
        Explanation Comment:
            For a single head attention the input will be of emb_size and output will be of emb_size Linear(emb_size,emb_size) but in multihead attention
            there are number of heads (n_heads) and each head has it's own key(k), query(q) and value(v) so total number of k,q,v
            will be equal to n_heads.
            We can do that in a single linear layer by giving the size of output layer = n_head * emb_size (here emb_size = output size)
            Same for k,q,v. nn.Linear(emb_size,n_heads*emb_size)

        """
        self.tokeys = nn.Linear(emb_size, n_heads * emb_size, bias=False)
        self.toquery = nn.Linear(emb_size, n_heads * emb_size, bias=False)
        self.tovalue = nn.Linear(emb_size, n_heads * emb_size, bias=False)
        self.combine_heads = nn.Linear(n_heads * emb_size, emb_size)
        self.Attnweights = None
        self.mask = mask

    def __MatrixMask(self, matrix, mask_value=float('-inf'), diagonal=False):
        """
        In self attention mask is used when we are trying to predict the next work based on the previous sequence
        but attention inherently contains information about all the words in the sequence.
        In order to accuretly predict the next word without sort of cheating by looking ahead with attention we set all
        the attention weights of the next words, tokens = -inf or zero, so that there is no input from the sequence ahead
        of the current word


        [ 1 2 3                 [ 1 0 0
          4 5 6    =========>     4 5 0
          7 8 9 ]                 7 8 9 ]

        """
        row, column = matrix.size(-2), matrix.size(-1)
        offset = 0 if diagonal else 1
        upper_triangular_idx = torch.triu_indices(row, column, offset)
        matrix[..., upper_triangular_idx[0], upper_triangular_idx[1]
            ] = mask_value
        return matrix

    def forward(self, x):
        """
        Argument:

        x: shoud be of shape [batch,sequence length, emb_size]

        Local variable:

        b: batch size
        t: length of sequence
        e: embedding size of each individual input token/word
        h: number of heads

        """
        b, t, e = x.size()
        h = self.n_heads
        assert e == self.emb_size, f'Input size expected to be {self.emb_size} but got {e}'
        keys = self.tokeys(x).view(b, t, h, e)
        query = self.toquery(x).view(b, t, h, e)
        value = self.tovalue(x).view(b, t, h, e)
        """
        Pytorch Contiguous() function is used to keep the indexes which are side by side, side by side in memory also

        """
        keys = keys.transpose(1, 2).contiguous().view(b * h, t, e)
        query = query.transpose(1, 2).contiguous().view(b * h, t, e)
        value = value.transpose(1, 2).contiguous().view(b * h, t, e)
        raw_weights = torch.bmm(query, keys.transpose(1, 2))
        raw_weights /= math.sqrt(e)
        assert raw_weights.size() == (b * h, t, t
            ), f'expected shape {b * h, t, t} got {raw_weights.size()}'
        if self.mask:
            raw_weights = self.__MatrixMask(raw_weights)
        weights = F.softmax(raw_weights, dim=2)
        self.Attnweights = weights
        out = torch.bmm(weights, value).view(b, h, t, e)
        out = out.transpose(1, 2).contiguous().view(b, t, h * e)
        out = self.combine_heads(out)
        return out


def get_inputs():
    return [torch.rand([4, 4, 4])]


def get_init_inputs():
    return [[], {'emb_size': 4}]

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
from torch._inductor.runtime import triton_helpers
from torch._inductor.runtime.triton_helpers import math as tl_math
from torch import nn
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, xnumel, XBLOCK: tl.constexpr):
    xnumel = 512
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex % 4
    x1 = xindex // 4 % 4
    x2 = xindex // 16 % 8
    x3 = xindex // 128
    x4 = xindex
    tmp0 = tl.load(in_ptr0 + (x0 + 4 * x2 + 32 * x1 + 128 * x3), xmask)
    tl.store(out_ptr0 + x4, tmp0, xmask)


@triton.jit
def triton_poi_fused__softmax_1(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr
    ):
    xnumel = 512
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x2 = xindex
    x1 = xindex // 4
    tmp0 = tl.load(in_ptr0 + x2, xmask)
    tmp3 = tl.load(in_ptr0 + 4 * x1, xmask, eviction_policy='evict_last')
    tmp5 = tl.load(in_ptr0 + (1 + 4 * x1), xmask, eviction_policy='evict_last')
    tmp8 = tl.load(in_ptr0 + (2 + 4 * x1), xmask, eviction_policy='evict_last')
    tmp11 = tl.load(in_ptr0 + (3 + 4 * x1), xmask, eviction_policy='evict_last'
        )
    tmp1 = 1.0
    tmp2 = tmp0 * tmp1
    tmp4 = tmp3 * tmp1
    tmp6 = tmp5 * tmp1
    tmp7 = triton_helpers.maximum(tmp4, tmp6)
    tmp9 = tmp8 * tmp1
    tmp10 = triton_helpers.maximum(tmp7, tmp9)
    tmp12 = tmp11 * tmp1
    tmp13 = triton_helpers.maximum(tmp10, tmp12)
    tmp14 = tmp2 - tmp13
    tmp15 = 0.5
    tmp16 = tmp14 * tmp15
    tmp17 = tl_math.exp(tmp16)
    tl.store(out_ptr0 + x2, tmp17, xmask)


@triton.jit
def triton_poi_fused__softmax_2(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr
    ):
    xnumel = 512
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x2 = xindex
    x1 = xindex // 4
    tmp0 = tl.load(in_ptr0 + x2, xmask)
    tmp1 = tl.load(in_ptr0 + 4 * x1, xmask, eviction_policy='evict_last')
    tmp2 = tl.load(in_ptr0 + (1 + 4 * x1), xmask, eviction_policy='evict_last')
    tmp4 = tl.load(in_ptr0 + (2 + 4 * x1), xmask, eviction_policy='evict_last')
    tmp6 = tl.load(in_ptr0 + (3 + 4 * x1), xmask, eviction_policy='evict_last')
    tmp3 = tmp1 + tmp2
    tmp5 = tmp3 + tmp4
    tmp7 = tmp5 + tmp6
    tmp8 = tmp0 / tmp7
    tl.store(out_ptr0 + x2, tmp8, xmask)


@triton.jit
def triton_poi_fused_clone_3(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 512
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex % 4
    x1 = xindex // 4 % 8
    x2 = xindex // 32 % 4
    x3 = xindex // 128
    x4 = xindex
    tmp0 = tl.load(in_ptr0 + (x0 + 4 * x2 + 16 * x1 + 128 * x3), xmask)
    tl.store(out_ptr0 + x4, tmp0, xmask)


def call(args):
    primals_1, primals_2, primals_3, primals_4, primals_5, primals_6 = args
    args.clear()
    assert_size_stride(primals_1, (4, 4, 4), (16, 4, 1))
    assert_size_stride(primals_2, (32, 4), (4, 1))
    assert_size_stride(primals_3, (32, 4), (4, 1))
    assert_size_stride(primals_4, (32, 4), (4, 1))
    assert_size_stride(primals_5, (4, 32), (32, 1))
    assert_size_stride(primals_6, (4,), (1,))
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        buf0 = empty_strided_cuda((16, 32), (32, 1), torch.float32)
        extern_kernels.mm(reinterpret_tensor(primals_1, (16, 4), (4, 1), 0),
            reinterpret_tensor(primals_2, (4, 32), (1, 4), 0), out=buf0)
        del primals_2
        buf1 = empty_strided_cuda((16, 32), (32, 1), torch.float32)
        extern_kernels.mm(reinterpret_tensor(primals_1, (16, 4), (4, 1), 0),
            reinterpret_tensor(primals_3, (4, 32), (1, 4), 0), out=buf1)
        del primals_3
        buf2 = empty_strided_cuda((16, 32), (32, 1), torch.float32)
        extern_kernels.mm(reinterpret_tensor(primals_1, (16, 4), (4, 1), 0),
            reinterpret_tensor(primals_4, (4, 32), (1, 4), 0), out=buf2)
        del primals_4
        buf3 = empty_strided_cuda((4, 8, 4, 4), (128, 16, 4, 1), torch.float32)
        get_raw_stream(0)
        triton_poi_fused_clone_0[grid(512)](buf1, buf3, 512, XBLOCK=128,
            num_warps=4, num_stages=1)
        buf4 = reinterpret_tensor(buf1, (4, 8, 4, 4), (128, 16, 4, 1), 0)
        del buf1
        triton_poi_fused_clone_0[grid(512)](buf0, buf4, 512, XBLOCK=128,
            num_warps=4, num_stages=1)
        buf5 = reinterpret_tensor(buf0, (32, 4, 4), (16, 4, 1), 0)
        del buf0
        extern_kernels.bmm(reinterpret_tensor(buf3, (32, 4, 4), (16, 4, 1),
            0), reinterpret_tensor(buf4, (32, 4, 4), (16, 1, 4), 0), out=buf5)
        buf6 = empty_strided_cuda((32, 4, 4), (16, 4, 1), torch.float32)
        triton_poi_fused__softmax_1[grid(512)](buf5, buf6, 512, XBLOCK=128,
            num_warps=4, num_stages=1)
        buf7 = buf5
        del buf5
        triton_poi_fused__softmax_2[grid(512)](buf6, buf7, 512, XBLOCK=256,
            num_warps=4, num_stages=1)
        buf8 = reinterpret_tensor(buf6, (4, 8, 4, 4), (128, 16, 4, 1), 0)
        del buf6
        triton_poi_fused_clone_0[grid(512)](buf2, buf8, 512, XBLOCK=128,
            num_warps=4, num_stages=1)
        buf9 = reinterpret_tensor(buf2, (32, 4, 4), (16, 4, 1), 0)
        del buf2
        extern_kernels.bmm(buf7, reinterpret_tensor(buf8, (32, 4, 4), (16, 
            4, 1), 0), out=buf9)
        buf10 = empty_strided_cuda((4, 4, 8, 4), (128, 32, 4, 1), torch.float32
            )
        triton_poi_fused_clone_3[grid(512)](buf9, buf10, 512, XBLOCK=128,
            num_warps=4, num_stages=1)
        del buf9
        buf11 = empty_strided_cuda((16, 4), (4, 1), torch.float32)
        extern_kernels.addmm(primals_6, reinterpret_tensor(buf10, (16, 32),
            (32, 1), 0), reinterpret_tensor(primals_5, (32, 4), (1, 32), 0),
            alpha=1, beta=1, out=buf11)
        del primals_6
    return reinterpret_tensor(buf11, (4, 4, 4), (16, 4, 1), 0
        ), buf7, reinterpret_tensor(primals_1, (16, 4), (4, 1), 0
        ), buf7, reinterpret_tensor(buf10, (16, 32), (32, 1), 0
        ), primals_5, reinterpret_tensor(buf8, (32, 4, 4), (16, 1, 4), 0
        ), reinterpret_tensor(buf3, (32, 4, 4), (16, 1, 4), 0
        ), reinterpret_tensor(buf4, (32, 4, 4), (16, 4, 1), 0)


class MultiHeadAttentionNew(nn.Module):

    def __init__(self, emb_size, n_heads=8, mask=False):
        """
        Arguments:

        emb_size: Size of input Embeddings
        n_heads: Number of heads for MultiHead Attention

        Layers:

        tokeys: For Keys
        toquery: For Query
        tovalue: For Value

        combine_heads: Convert the ( (number of heads) * (emb_size) ) into (emb_size) for output

        Note:
            The input dimension we give to the Linear Layer (nn.Linear) is the last dimension (no. of columns) in size
            (batch_size ,seq_length, input_dims) ==> input_dims : give to nn.Linear
            So in order for the tensor to pass through linear layer having the last dim equal to input_dim of
            Linear layer is necessary.

        """
        super(MultiHeadAttentionNew, self).__init__()
        self.emb_size = emb_size
        self.n_heads = n_heads
        """
        Explanation Comment:
            For a single head attention the input will be of emb_size and output will be of emb_size Linear(emb_size,emb_size) but in multihead attention
            there are number of heads (n_heads) and each head has it's own key(k), query(q) and value(v) so total number of k,q,v
            will be equal to n_heads.
            We can do that in a single linear layer by giving the size of output layer = n_head * emb_size (here emb_size = output size)
            Same for k,q,v. nn.Linear(emb_size,n_heads*emb_size)

        """
        self.tokeys = nn.Linear(emb_size, n_heads * emb_size, bias=False)
        self.toquery = nn.Linear(emb_size, n_heads * emb_size, bias=False)
        self.tovalue = nn.Linear(emb_size, n_heads * emb_size, bias=False)
        self.combine_heads = nn.Linear(n_heads * emb_size, emb_size)
        self.Attnweights = None
        self.mask = mask

    def __MatrixMask(self, matrix, mask_value=float('-inf'), diagonal=False):
        """
        In self attention mask is used when we are trying to predict the next work based on the previous sequence
        but attention inherently contains information about all the words in the sequence.
        In order to accuretly predict the next word without sort of cheating by looking ahead with attention we set all
        the attention weights of the next words, tokens = -inf or zero, so that there is no input from the sequence ahead
        of the current word


        [ 1 2 3                 [ 1 0 0
          4 5 6    =========>     4 5 0
          7 8 9 ]                 7 8 9 ]

        """
        row, column = matrix.size(-2), matrix.size(-1)
        offset = 0 if diagonal else 1
        upper_triangular_idx = torch.triu_indices(row, column, offset)
        matrix[..., upper_triangular_idx[0], upper_triangular_idx[1]
            ] = mask_value
        return matrix

    def forward(self, input_0):
        primals_2 = self.tokeys.weight
        primals_3 = self.toquery.weight
        primals_4 = self.tovalue.weight
        primals_5 = self.combine_heads.weight
        primals_6 = self.combine_heads.bias
        primals_1 = input_0
        output = call([primals_1, primals_2, primals_3, primals_4,
            primals_5, primals_6])
        return output[0]
