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 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
reinterpret_tensor = torch._C._dynamo.guards._reinterpret_tensor


@triton.jit
def triton_poi_fused_cat_view_0(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr
    ):
    xnumel = 80
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex % 5
    x1 = xindex // 5
    x2 = 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 * x1 + x0), tmp4 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp6 = tmp0 >= tmp3
    tl.full([1], 5, tl.int64)
    tmp9 = 1.0
    tmp10 = tl.full(tmp9.shape, 0.0, tmp9.dtype)
    tmp11 = tl.where(tmp6, tmp9, tmp10)
    tmp12 = tl.where(tmp4, tmp5, tmp11)
    tl.store(out_ptr0 + x2, tmp12, xmask)


@triton.jit
def triton_poi_fused_add_transpose_1(in_ptr0, in_ptr1, out_ptr1, ynumel,
    xnumel, YBLOCK: tl.constexpr, XBLOCK: tl.constexpr):
    ynumel = 64
    xnumel = 4
    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
    y0 = yindex % 4
    y1 = yindex // 4
    tmp0 = tl.load(in_ptr0 + (y0 + 4 * x2 + 16 * y1), xmask & ymask,
        eviction_policy='evict_last')
    tmp1 = tl.load(in_ptr1 + x2, xmask, eviction_policy='evict_last')
    tmp2 = tmp0 + tmp1
    tl.store(out_ptr1 + (y0 + 4 * x2 + 16 * y1), tmp2, xmask & ymask)


def call(args):
    primals_1, primals_2, primals_3, primals_4 = args
    args.clear()
    assert_size_stride(primals_1, (4, 4, 4), (16, 4, 1))
    assert_size_stride(primals_2, (4, 4, 4), (16, 4, 1))
    assert_size_stride(primals_3, (5, 4, 5), (20, 5, 1))
    assert_size_stride(primals_4, (4,), (1,))
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        buf0 = empty_strided_cuda((16, 5), (5, 1), torch.float32)
        get_raw_stream(0)
        triton_poi_fused_cat_view_0[grid(80)](primals_1, buf0, 80, XBLOCK=
            128, num_warps=4, num_stages=1)
        del primals_1
        buf1 = empty_strided_cuda((16, 20), (20, 1), torch.float32)
        extern_kernels.mm(buf0, reinterpret_tensor(primals_3, (5, 20), (20,
            1), 0), out=buf1)
        del primals_3
        buf2 = empty_strided_cuda((4, 4, 5), (20, 5, 1), torch.float32)
        triton_poi_fused_cat_view_0[grid(80)](primals_2, buf2, 80, XBLOCK=
            128, num_warps=4, num_stages=1)
        del primals_2
        buf3 = empty_strided_cuda((4, 16, 4), (64, 4, 1), torch.float32)
        extern_kernels.bmm(reinterpret_tensor(buf1, (4, 16, 5), (80, 5, 1),
            0), reinterpret_tensor(buf2, (4, 5, 4), (20, 1, 5), 0), out=buf3)
        del buf1
        buf5 = empty_strided_cuda((4, 4, 4, 4), (64, 16, 1, 4), torch.float32)
        triton_poi_fused_add_transpose_1[grid(64, 4)](buf3, primals_4, buf5,
            64, 4, XBLOCK=4, YBLOCK=32, num_warps=4, num_stages=1)
        del buf3
        del primals_4
    return buf5, buf2, reinterpret_tensor(buf0, (5, 16), (1, 5), 0)


class PairwiseBilinear(nn.Module):
    """
    https://github.com/stanfordnlp/stanza/blob/v1.1.1/stanza/models/common/biaffine.py#L5  # noqa
    """

    def __init__(self, in1_features: 'int', in2_features: 'int',
        out_features: 'int', bias: 'bool'=True):
        super().__init__()
        self.in1_features = in1_features
        self.in2_features = in2_features
        self.out_features = out_features
        self.weight = nn.Parameter(torch.Tensor(in1_features, out_features,
            in2_features))
        if bias:
            self.bias = nn.Parameter(torch.Tensor(out_features))
        else:
            self.register_parameter('bias', None)
        self.reset_parameters()

    def reset_parameters(self):
        bound = 1 / math.sqrt(self.weight.size(0))
        nn.init.uniform_(self.weight, -bound, bound)
        if self.bias is not None:
            nn.init.uniform_(self.bias, -bound, bound)

    def forward(self, input1: 'torch.Tensor', input2: 'torch.Tensor'
        ) ->torch.Tensor:
        d1, d2, out = self.in1_features, self.in2_features, self.out_features
        n1, n2 = input1.size(1), input2.size(1)
        x1W = torch.mm(input1.view(-1, d1), self.weight.view(d1, out * d2))
        x1Wx2 = x1W.view(-1, n1 * out, d2).bmm(input2.transpose(1, 2))
        y = x1Wx2.view(-1, n1, self.out_features, n2).transpose(2, 3)
        if self.bias is not None:
            y.add_(self.bias)
        return y

    def extra_repr(self) ->str:
        return ('in1_features={}, in2_features={}, out_features={}, bias={}'
            .format(self.in1_features, self.in2_features, self.out_features,
            self.bias is not None))


class BiaffineNew(nn.Module):

    def __init__(self, in1_features: 'int', in2_features: 'int',
        out_features: 'int'):
        super().__init__()
        self.bilinear = PairwiseBilinear(in1_features + 1, in2_features + 1,
            out_features)
        self.bilinear.weight.data.zero_()
        self.bilinear.bias.data.zero_()

    def forward(self, input_0, input_1):
        primals_3 = self.bilinear.weight
        primals_4 = self.bilinear.bias
        primals_1 = input_0
        primals_2 = input_1
        output = call([primals_1, primals_2, primals_3, primals_4])
        return output[0]


Measured latency_ms: 0.057757