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 numpy as np
from typing import Iterable
from typing import Optional
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
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 = 1024
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex % 64
    x2 = xindex // 256
    x3 = xindex
    tmp0 = tl.load(in_ptr0 + (x0 + 64 * x2), xmask, eviction_policy=
        'evict_last')
    tl.store(out_ptr0 + x3, tmp0, xmask)


@triton.jit
def triton_poi_fused_clone_1(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 1024
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex % 256
    x2 = xindex
    tmp0 = tl.load(in_ptr0 + x0, xmask, eviction_policy='evict_last')
    tl.store(out_ptr0 + x2, tmp0, xmask)


@triton.jit
def triton_poi_fused_add_div_2(in_out_ptr0, in_ptr0, in_ptr1, xnumel,
    XBLOCK: tl.constexpr):
    xnumel = 1024
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex % 256
    x2 = xindex
    tmp0 = tl.load(in_ptr0 + x0, xmask, eviction_policy='evict_last')
    tmp1 = tl.load(in_out_ptr0 + x2, xmask)
    tmp2 = tl.load(in_ptr1 + x2, xmask)
    tmp3 = tmp1 + tmp2
    tmp4 = 0.5
    tmp5 = tmp3 * tmp4
    tmp6 = tmp0 + tmp5
    tl.store(in_out_ptr0 + x2, tmp6, xmask)


def call(args):
    (primals_1, primals_2, primals_3, primals_4, primals_5, primals_6,
        primals_7, primals_8, primals_9, primals_10, primals_11) = args
    args.clear()
    assert_size_stride(primals_1, (1, 4), (4, 1))
    assert_size_stride(primals_2, (1,), (1,))
    assert_size_stride(primals_3, (4, 4, 4, 4), (64, 16, 4, 1))
    assert_size_stride(primals_4, (1, 4), (4, 1))
    assert_size_stride(primals_5, (1,), (1,))
    assert_size_stride(primals_6, (4, 4, 4, 4), (64, 16, 4, 1))
    assert_size_stride(primals_7, (4, 4, 4, 4), (64, 16, 4, 1))
    assert_size_stride(primals_8, (4, 4), (4, 1))
    assert_size_stride(primals_9, (4,), (1,))
    assert_size_stride(primals_10, (4, 4), (4, 1))
    assert_size_stride(primals_11, (4,), (1,))
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        buf1 = empty_strided_cuda((64, 1), (1, 1), torch.float32)
        extern_kernels.addmm(primals_2, reinterpret_tensor(primals_3, (64, 
            4), (4, 1), 0), reinterpret_tensor(primals_1, (4, 1), (1, 4), 0
            ), alpha=1, beta=1, out=buf1)
        del primals_1
        del primals_2
        buf3 = empty_strided_cuda((64, 1), (1, 1), torch.float32)
        extern_kernels.addmm(primals_5, reinterpret_tensor(primals_3, (64, 
            4), (4, 1), 0), reinterpret_tensor(primals_4, (4, 1), (1, 4), 0
            ), alpha=1, beta=1, out=buf3)
        del primals_4
        del primals_5
        buf4 = empty_strided_cuda((4, 4, 4, 4, 4), (256, 64, 16, 4, 1),
            torch.float32)
        get_raw_stream(0)
        triton_poi_fused_clone_0[grid(1024)](primals_6, buf4, 1024, XBLOCK=
            256, num_warps=4, num_stages=1)
        del primals_6
        buf5 = empty_strided_cuda((4, 4, 4, 4, 4), (256, 64, 16, 4, 1),
            torch.float32)
        triton_poi_fused_clone_1[grid(1024)](primals_3, buf5, 1024, XBLOCK=
            256, num_warps=4, num_stages=1)
        buf6 = empty_strided_cuda((64, 4, 4), (16, 4, 1), torch.float32)
        extern_kernels.bmm(reinterpret_tensor(buf4, (64, 4, 4), (16, 4, 1),
            0), reinterpret_tensor(buf5, (64, 4, 4), (16, 4, 1), 0), out=buf6)
        buf7 = buf4
        del buf4
        triton_poi_fused_clone_0[grid(1024)](primals_7, buf7, 1024, XBLOCK=
            256, num_warps=4, num_stages=1)
        del primals_7
        buf8 = empty_strided_cuda((64, 4, 4), (16, 4, 1), torch.float32)
        extern_kernels.bmm(reinterpret_tensor(buf7, (64, 4, 4), (16, 4, 1),
            0), reinterpret_tensor(buf5, (64, 4, 4), (16, 4, 1), 0), out=buf8)
        del buf5
        buf9 = reinterpret_tensor(buf6, (4, 4, 4, 4, 4), (256, 64, 16, 4, 1), 0
            )
        del buf6
        triton_poi_fused_add_div_2[grid(1024)](buf9, primals_3, buf8, 1024,
            XBLOCK=256, num_warps=4, num_stages=1)
        buf10 = reinterpret_tensor(buf8, (256, 4), (4, 1), 0)
        del buf8
        extern_kernels.addmm(primals_9, reinterpret_tensor(buf9, (256, 4),
            (4, 1), 0), reinterpret_tensor(primals_8, (4, 4), (1, 4), 0),
            alpha=1, beta=1, out=buf10)
        del primals_8
        del primals_9
        buf11 = reinterpret_tensor(buf7, (256, 4), (4, 1), 0)
        del buf7
        extern_kernels.addmm(primals_11, reinterpret_tensor(buf9, (256, 4),
            (4, 1), 0), reinterpret_tensor(primals_10, (4, 4), (1, 4), 0),
            alpha=1, beta=1, out=buf11)
        del primals_10
        del primals_11
    return reinterpret_tensor(buf1, (4, 4, 4), (16, 4, 1), 0
        ), reinterpret_tensor(buf3, (4, 4, 4), (16, 4, 1), 0
        ), reinterpret_tensor(buf10, (4, 4, 4, 4, 4), (256, 64, 16, 4, 1), 0
        ), reinterpret_tensor(buf11, (4, 4, 4, 4, 4), (256, 64, 16, 4, 1), 0
        ), reinterpret_tensor(primals_3, (64, 4), (4, 1), 0
        ), reinterpret_tensor(buf9, (256, 4), (4, 1), 0)


def find_closest_span_pairs(head: 'Iterable', tail: 'Iterable', backtrace:
    'Optional[bool]'=True):
    """
    Find all span pairs.

    Args:
        head: list of start position predictions, either 1 or 0
        tail: list of end position predictions, either 1 or 0
        backtrace: if there are more tail predictions than head predictions,
            then backtrace to find a closest head position to get a span pair

    Examples:
        >>> head = torch.tensor([1, 0, 0, 1, 0, 0, 1], dtype=torch.long)
        >>> tail = torch.tensor([0, 1, 0, 1, 0, 1, 1], dtype=torch.long)
        >>> find_closest_span_pairs(head, tail, backtrace=False)
        [(0, 1), (3, 3), (6, 6)]
        >>> find_closest_span_pairs(head, tail, backtrace=True)
        [(0, 1), (3, 3), (6, 6), (3, 5)]
    """
    if isinstance(head, torch.Tensor):
        head = head.detach().cpu()
    if isinstance(tail, torch.Tensor):
        tail = tail.detach().cpu()
    head_valid_poses = np.where(head == 1)[0]
    tail_valid_poses = np.where(tail == 1)[0]
    tail_used_poses = {pos: (False) for pos in tail_valid_poses.tolist()}
    pairs = []
    for head_i in head_valid_poses:
        tail_js = tail_valid_poses[tail_valid_poses >= head_i]
        if len(tail_js) > 0:
            tail_j = tail_js[0]
            tail_used_poses[tail_j] = True
            pairs.append((head_i, tail_j))
    if backtrace:
        for tail_j in tail_used_poses:
            if tail_used_poses[tail_j] is False:
                head_is = head_valid_poses[head_valid_poses <= tail_j]
                if len(head_is) > 0:
                    head_i = head_is[-1]
                    pairs.append((head_i, tail_j))
    return pairs


def find_closest_span_pairs_with_index(heads: 'Iterable', tails: 'Iterable',
    backtrace: 'Optional[bool]'=True):
    """
    Find all possible pairs with indexes,
    useful for object discoveries with class idx.

    Args:
        heads: batch of torch.Tensor
        tails: batch of torch.Tensor
        backtrace: if there are more tail predictions than head predictions,
            then backtrace to find a closest head position to get a span pair

    Examples:
        >>> heads = torch.tensor([[1, 0, 0, 1, 0, 0, 1], [1, 0, 0, 1, 0, 0, 1]], dtype=torch.long)
        >>> tails = torch.tensor([[0, 1, 0, 1, 0, 1, 1], [0, 1, 0, 0, 0, 1, 0]], dtype=torch.long)
        >>> find_closest_span_pairs(heads, tails, backtrace=False)
        [(0, 0, 1), (0, 3, 3), (0, 6, 6), (1, 0, 1), (1, 3, 5)]
        >>> find_closest_span_pairs(heads, tails, backtrace=True)
        [(0, 0, 1), (0, 3, 3), (0, 6, 6), (0, 3, 5), (1, 0, 1), (1, 3, 5)]
    """
    results = []
    for idx, (head, tail) in enumerate(zip(heads, tails)):
        pairs = find_closest_span_pairs(head, tail, backtrace=backtrace)
        for pair in pairs:
            results.append((idx, pair[0], pair[1]))
    return results


class SubjObjSpanNew(nn.Module):
    """
    Inputs:
        hidden: (batch_size, seq_len, hidden_size)
        one_subj_head: object golden head with one subject (batch_size, hidden_size)
        one_subj_tail: object golden tail with one subject (batch_size, hidden_size)
    """

    def __init__(self, hidden_size, num_classes, threshold:
        'Optional[float]'=0.5):
        super().__init__()
        self.threshold = threshold
        self.subj_head_ffnn = nn.Linear(hidden_size, 1)
        self.subj_tail_ffnn = nn.Linear(hidden_size, 1)
        self.obj_head_ffnn = nn.Linear(hidden_size, num_classes)
        self.obj_tail_ffnn = nn.Linear(hidden_size, num_classes)

    def get_objs_for_specific_subj(self, subj_head_mapping,
        subj_tail_mapping, hidden):
        subj_head = torch.matmul(subj_head_mapping, hidden)
        subj_tail = torch.matmul(subj_tail_mapping, hidden)
        sub = (subj_head + subj_tail) / 2
        encoded_text = hidden + sub
        pred_obj_heads = self.obj_head_ffnn(encoded_text)
        pred_obj_tails = self.obj_tail_ffnn(encoded_text)
        return pred_obj_heads, pred_obj_tails

    def build_mapping(self, subj_heads, subj_tails):
        """
        Build head & tail mapping for predicted subjects,
        for each instance in a batch, for a subject in all
        the predicted subjects, return a single subject
        and its corresponding mappings.
        """
        for subj_head, subj_tail in zip(subj_heads, subj_tails):
            subjs = find_closest_span_pairs(subj_head, subj_tail)
            seq_len = subj_head.shape[0]
            for subj in subjs:
                subj_head_mapping = torch.zeros(seq_len, device=subj_head.
                    device)
                subj_tail_mapping = torch.zeros(seq_len, device=subj_tail.
                    device)
                subj_head_mapping[subj[0]] = 1.0
                subj_tail_mapping[subj[1]] = 1.0
                yield subj, subj_head_mapping, subj_tail_mapping

    def build_batch_mapping(self, subj_head, subj_tail):
        """
        Build head & tail mapping for predicted subjects,
        for each instance in a batch, return all the predicted
        subjects and mappings.
        """
        subjs = find_closest_span_pairs(subj_head, subj_tail)
        seq_len = subj_head.shape[0]
        if len(subjs) > 0:
            subjs_head_mapping = torch.zeros(len(subjs), seq_len, device=
                subj_head.device)
            subjs_tail_mapping = torch.zeros(len(subjs), seq_len, device=
                subj_tail.device)
            for subj_idx, subj in enumerate(subjs):
                subjs_head_mapping[subj_idx, subj[0]] = 1.0
                subjs_tail_mapping[subj_idx, subj[1]] = 1.0
            return subjs, subjs_head_mapping, subjs_tail_mapping
        else:
            return None, None, None

    def predict(self, hidden):
        if hidden.shape[0] != 1:
            raise RuntimeError(
                f'eval batch size must be 1 x hidden_size, while hidden is {hidden.shape}'
                )
        subj_head_out = self.subj_head_ffnn(hidden)
        subj_tail_out = self.subj_tail_ffnn(hidden)
        subj_head_out = torch.sigmoid(subj_head_out)
        subj_tail_out = torch.sigmoid(subj_tail_out)
        pred_subj_head = subj_head_out.ge(self.threshold).long()
        pred_subj_tail = subj_tail_out.ge(self.threshold).long()
        triples = []
        subjs, subj_head_mappings, subj_tail_mappings = (self.
            build_batch_mapping(pred_subj_head.squeeze(0).squeeze(-1),
            pred_subj_tail.squeeze(0).squeeze(-1)))
        if subjs:
            obj_head_out, obj_tail_out = self.get_objs_for_specific_subj(
                subj_head_mappings.unsqueeze(1), subj_tail_mappings.
                unsqueeze(1), hidden)
            obj_head_out = torch.sigmoid(obj_head_out)
            obj_tail_out = torch.sigmoid(obj_tail_out)
            obj_head_out = obj_head_out.ge(self.threshold).long()
            obj_tail_out = obj_tail_out.ge(self.threshold).long()
            for subj_idx, subj in enumerate(subjs):
                objs = find_closest_span_pairs_with_index(obj_head_out[
                    subj_idx].permute(1, 0), obj_tail_out[subj_idx].permute
                    (1, 0))
                for relation_idx, obj_pair_start, obj_pair_end in objs:
                    triples.append(((subj[0], subj[1] + 1), relation_idx, (
                        obj_pair_start, obj_pair_end + 1)))
        return [triples]

    def forward(self, input_0, input_1, input_2):
        primals_1 = self.subj_head_ffnn.weight
        primals_2 = self.subj_head_ffnn.bias
        primals_4 = self.subj_tail_ffnn.weight
        primals_5 = self.subj_tail_ffnn.bias
        primals_8 = self.obj_head_ffnn.weight
        primals_9 = self.obj_head_ffnn.bias
        primals_10 = self.obj_tail_ffnn.weight
        primals_11 = self.obj_tail_ffnn.bias
        primals_3 = input_0
        primals_6 = input_1
        primals_7 = input_2
        output = call([primals_1, primals_2, primals_3, primals_4,
            primals_5, primals_6, primals_7, primals_8, primals_9,
            primals_10, primals_11])
        return output[0], output[1], output[2], output[3]


Measured latency_ms: 0.025518