# AOT ID: ['0_inference']
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_4/inductor_cache/kq/ckqk4ahydv37mmcfwbfoz46szfvxyyipbmsfv4nxhr5bpnwyl3ln.py
# Topologically Sorted Source Nodes: [pair_matrix], Original ATen: [aten.cat]
# Source node to ATen node mapping:
#   pair_matrix => cat
# Graph fragment:
#   %cat : [num_users=1] = call_function[target=torch.ops.aten.cat.default](args = ([%view_1, %permute, %mul], -1), kwargs = {})
triton_poi_fused_cat_0 = async_compile.triton('triton_poi_fused_cat_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=[1024], 
    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_cat_0', 'mutated_arg_names': [], 'no_x_dim': False, 'num_load': 4, '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_cat_0(in_ptr0, out_ptr0, xnumel, XBLOCK : tl.constexpr):
    xnumel = 768
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex % 12
    x4 = (xindex // 48)
    x1 = (xindex // 12) % 4
    x3 = (xindex // 192)
    x5 = xindex
    tmp0 = x0
    tmp1 = tl.full([1], 0, tl.int64)
    tmp2 = tmp0 >= tmp1
    tmp3 = tl.full([1], 4, tl.int64)
    tmp4 = tmp0 < tmp3
    tmp5 = tl.load(in_ptr0 + ((4*x4) + x0), tmp4 & xmask, eviction_policy='evict_last', other=0.0)
    tmp6 = tmp0 >= tmp3
    tmp7 = tl.full([1], 8, tl.int64)
    tmp8 = tmp0 < tmp7
    tmp9 = tmp6 & tmp8
    tmp10 = tl.load(in_ptr0 + ((4*x1) + (16*x3) + ((-4) + x0)), tmp9 & xmask, eviction_policy='evict_last', other=0.0)
    tmp11 = tmp0 >= tmp7
    tmp12 = tl.full([1], 12, tl.int64)
    tmp13 = tmp0 < tmp12
    tmp14 = tl.load(in_ptr0 + ((4*x4) + ((-8) + x0)), tmp11 & xmask, eviction_policy='evict_last', other=0.0)
    tmp15 = tl.load(in_ptr0 + ((4*x1) + (16*x3) + ((-8) + x0)), tmp11 & xmask, eviction_policy='evict_last', other=0.0)
    tmp16 = tmp14 * tmp15
    tmp17 = tl.full(tmp16.shape, 0.0, tmp16.dtype)
    tmp18 = tl.where(tmp11, tmp16, tmp17)
    tmp19 = tl.where(tmp9, tmp10, tmp18)
    tmp20 = tl.where(tmp4, tmp5, tmp19)
    tl.store(out_ptr0 + (x5), tmp20, xmask)
''', device_str='cuda')


async_compile.wait(globals())
del async_compile

def call(args):
    arg0_1, = args
    args.clear()
    assert_size_stride(arg0_1, (4, 4, 4), (16, 4, 1))
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        buf0 = empty_strided_cuda((4, 4, 4, 12), (192, 48, 12, 1), torch.float32)
        # Topologically Sorted Source Nodes: [pair_matrix], Original ATen: [aten.cat]
        stream0 = get_raw_stream(0)
        triton_poi_fused_cat_0.run(arg0_1, buf0, 768, grid=grid(768), stream=stream0)
        del arg0_1
    return (buf0, )


def benchmark_compiled_module(times=10, repeat=10):
    from torch._dynamo.testing import rand_strided
    from torch._inductor.utils import print_performance
    arg0_1 = rand_strided((4, 4, 4), (16, 4, 1), device='cuda:0', dtype=torch.float32)
    fn = lambda: call([arg0_1])
    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 torch
import numpy as np
from torch import Tensor
from torch.functional import Tensor
from typing import Union


class Pairer(torch.nn.Module):
    """

    To predict links between segments we will find all possible pairs and estimate the probability that they are linked.
    We do this by creating a matrix where each row in each sample represent a pair of segments. Let say we want to represent a pair
    the embedding of the first and second member, as well as the dot product of the member and lastly a positional encoding. Then
    a row would be 

       r =  [hi; hj; hi*hj, φ(h, m)] 

    where i, j are indexes of the segments, h is the representation of a segment and φ is a one-hot encodings
    representing the relative distance betwen the segments i,j (j-i)

    when we have such a matrix we score the all the pairs and select the pair which scores highest using Maximum Spanning Tree.
    I.e. simply taking the argmax along axis 2.

    To create that matrix following steps will be taken. In this example let us assume that the max 
    segments in our datset is 3.
    
    For each sample we want the following matrix:    
    s = [
                [
                [h0, h0, h0*h0, 0, 0, 1, 0, 0],
                [h0, h1, h0*h1, 0, 0, 0, 1, 0], 
                [h0, h2, h0*h2, 0, 0, 0, 0, 1],
                ],
                [
                [h1, h0, h1*h0, 0, 1, 0, 0, 0],
                [h1, h1, h1*h1, 0, 0, 1, 0, 0],
                [h1, h2, h1*h2, 0, 0, 0, 1, 0],
                ],
                [
                [h2, h0, h2*h0, 1, 0, 0, 0, 0],
                [h2, h1, h2*h1, 0, 1, 0, 0, 0],
                [h2, h2, h2*h2, 0, 0, 1, 0, 0],
                ]
            ]
        
    
    NOTE: that each row here is a 1d vector, the rows with h are vector represent a span of values

    THE LAST 5 DIMENSIONS is the one-hot encoding representing the relative position of the paired segment. 5 dimensions 
    is due to the possible relative relations, stretching from ( -(max_segments -1), (max_segments-1)). In this example
    the max_segments==3, which means that a segment can either be related to 2 segments behind itself (-2), 1 segment behind
    itself (-1) , 0 segments behind itself (0)  and so on, hence leaving us with a 5 dim one-hot-encodings.

    h = segment embedding
    
    A sample in out input matrix will look like:

        S = [
            [h0,h1,h2],
            ],

    1) duplicating the matrix at dim==1 nr times equal to max_segments ( dim==1)
    
    2) then we reshape the result of 1) so for each segment in each sample is filled with copies of itself.
       this will columns 0 (0:h).
       
    3) then we transpose output of 2) to create columns 1 (h:h+h)
    
    4) then we create the relattive positional one-hot-encodings. These encodings will follow a strict diagonal
       pattern as we can seen in the example above. 
       
    5) concatenate output from steps 2,3 and 4, along with the multiplication of output of step 2 and 3 which 
       creates the columns 2 (h*2:(h*2)).
    """

    def __init__(self, input_size: 'int', mode: 'str'='cat+multi',
        n_rel_pos: 'int'=None):
        super().__init__()
        self.mode = mode.split('+')
        self.n_rel_pos = n_rel_pos
        self.output_size = 0
        if n_rel_pos:
            self.output_size += self.n_rel_pos * 2 - 1
        if 'cat' in self.mode:
            self.output_size += input_size * 2
        if 'multi' in self.mode:
            self.output_size += input_size
        if 'sum' in self.mode:
            self.output_size += input_size
        if 'mean' in self.mode:
            self.output_size += input_size

    def forward(self, input: 'Tensor', device: 'Union[str, torch.device]'='cpu'
        ) ->Tensor:
        batch_size = input.shape[0]
        dim1 = input.shape[1]
        shape = batch_size, dim1, dim1, input.shape[-1]
        m = torch.reshape(torch.repeat_interleave(input, dim1, dim=1), shape)
        mT = m.transpose(2, 1)
        to_cat = []
        if 'cat' in self.mode:
            to_cat.append(m)
            to_cat.append(mT)
        if 'multi' in self.mode:
            to_cat.append(m * mT)
        if 'mean' in self.mode:
            to_cat.append(m + mT / 2)
        if 'sum' in self.mode:
            to_cat.append(m + mT)
        if self.n_rel_pos:
            rel_one_hot_dim = self.n_rel_pos * 2 - 1
            rel_one_hots = torch.tensor([np.diag(np.ones(rel_one_hot_dim),
                i)[:dim1, :rel_one_hot_dim] for i in range(dim1 - 1, -1, -1
                )], dtype=torch.uint8, device=device)
            rel_one_hots = rel_one_hots.repeat_interleave(batch_size, dim=0)
            rel_one_hots = rel_one_hots.view((batch_size, dim1, dim1,
                rel_one_hot_dim))
            to_cat.append(rel_one_hots)
        pair_matrix = torch.cat(to_cat, axis=-1)
        return pair_matrix


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


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

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
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_cat_0(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 768
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex % 12
    x4 = xindex // 48
    x1 = xindex // 12 % 4
    x3 = xindex // 192
    x5 = xindex
    tmp0 = x0
    tl.full([1], 0, tl.int64)
    tmp3 = tl.full([1], 4, tl.int64)
    tmp4 = tmp0 < tmp3
    tmp5 = tl.load(in_ptr0 + (4 * x4 + x0), tmp4 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp6 = tmp0 >= tmp3
    tmp7 = tl.full([1], 8, tl.int64)
    tmp8 = tmp0 < tmp7
    tmp9 = tmp6 & tmp8
    tmp10 = tl.load(in_ptr0 + (4 * x1 + 16 * x3 + (-4 + x0)), tmp9 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp11 = tmp0 >= tmp7
    tl.full([1], 12, tl.int64)
    tmp14 = tl.load(in_ptr0 + (4 * x4 + (-8 + x0)), tmp11 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp15 = tl.load(in_ptr0 + (4 * x1 + 16 * x3 + (-8 + x0)), tmp11 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp16 = tmp14 * tmp15
    tmp17 = tl.full(tmp16.shape, 0.0, tmp16.dtype)
    tmp18 = tl.where(tmp11, tmp16, tmp17)
    tmp19 = tl.where(tmp9, tmp10, tmp18)
    tmp20 = tl.where(tmp4, tmp5, tmp19)
    tl.store(out_ptr0 + x5, tmp20, xmask)


def call(args):
    arg0_1, = args
    args.clear()
    assert_size_stride(arg0_1, (4, 4, 4), (16, 4, 1))
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        buf0 = empty_strided_cuda((4, 4, 4, 12), (192, 48, 12, 1), torch.
            float32)
        get_raw_stream(0)
        triton_poi_fused_cat_0[grid(768)](arg0_1, buf0, 768, XBLOCK=256,
            num_warps=4, num_stages=1)
        del arg0_1
    return buf0,


class PairerNew(torch.nn.Module):
    """

    To predict links between segments we will find all possible pairs and estimate the probability that they are linked.
    We do this by creating a matrix where each row in each sample represent a pair of segments. Let say we want to represent a pair
    the embedding of the first and second member, as well as the dot product of the member and lastly a positional encoding. Then
    a row would be 

       r =  [hi; hj; hi*hj, φ(h, m)] 

    where i, j are indexes of the segments, h is the representation of a segment and φ is a one-hot encodings
    representing the relative distance betwen the segments i,j (j-i)

    when we have such a matrix we score the all the pairs and select the pair which scores highest using Maximum Spanning Tree.
    I.e. simply taking the argmax along axis 2.

    To create that matrix following steps will be taken. In this example let us assume that the max 
    segments in our datset is 3.
    
    For each sample we want the following matrix:    
    s = [
                [
                [h0, h0, h0*h0, 0, 0, 1, 0, 0],
                [h0, h1, h0*h1, 0, 0, 0, 1, 0], 
                [h0, h2, h0*h2, 0, 0, 0, 0, 1],
                ],
                [
                [h1, h0, h1*h0, 0, 1, 0, 0, 0],
                [h1, h1, h1*h1, 0, 0, 1, 0, 0],
                [h1, h2, h1*h2, 0, 0, 0, 1, 0],
                ],
                [
                [h2, h0, h2*h0, 1, 0, 0, 0, 0],
                [h2, h1, h2*h1, 0, 1, 0, 0, 0],
                [h2, h2, h2*h2, 0, 0, 1, 0, 0],
                ]
            ]
        
    
    NOTE: that each row here is a 1d vector, the rows with h are vector represent a span of values

    THE LAST 5 DIMENSIONS is the one-hot encoding representing the relative position of the paired segment. 5 dimensions 
    is due to the possible relative relations, stretching from ( -(max_segments -1), (max_segments-1)). In this example
    the max_segments==3, which means that a segment can either be related to 2 segments behind itself (-2), 1 segment behind
    itself (-1) , 0 segments behind itself (0)  and so on, hence leaving us with a 5 dim one-hot-encodings.

    h = segment embedding
    
    A sample in out input matrix will look like:

        S = [
            [h0,h1,h2],
            ],

    1) duplicating the matrix at dim==1 nr times equal to max_segments ( dim==1)
    
    2) then we reshape the result of 1) so for each segment in each sample is filled with copies of itself.
       this will columns 0 (0:h).
       
    3) then we transpose output of 2) to create columns 1 (h:h+h)
    
    4) then we create the relattive positional one-hot-encodings. These encodings will follow a strict diagonal
       pattern as we can seen in the example above. 
       
    5) concatenate output from steps 2,3 and 4, along with the multiplication of output of step 2 and 3 which 
       creates the columns 2 (h*2:(h*2)).
    """

    def __init__(self, input_size: 'int', mode: 'str'='cat+multi',
        n_rel_pos: 'int'=None):
        super().__init__()
        self.mode = mode.split('+')
        self.n_rel_pos = n_rel_pos
        self.output_size = 0
        if n_rel_pos:
            self.output_size += self.n_rel_pos * 2 - 1
        if 'cat' in self.mode:
            self.output_size += input_size * 2
        if 'multi' in self.mode:
            self.output_size += input_size
        if 'sum' in self.mode:
            self.output_size += input_size
        if 'mean' in self.mode:
            self.output_size += input_size

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