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 torch.nn as nn
from torch.nn.parameter import Parameter
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_add_mul_0(in_out_ptr0, in_ptr0, in_ptr1, in_ptr2,
    in_ptr3, xnumel, XBLOCK: tl.constexpr):
    xnumel = 400
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x3 = xindex
    x4 = xindex % 25
    x0 = xindex % 5
    x2 = xindex // 25
    tmp0 = tl.load(in_out_ptr0 + x3, xmask)
    tmp1 = tl.load(in_ptr0 + x4, xmask, eviction_policy='evict_last')
    tmp3 = tl.load(in_ptr1 + (x0 + 5 * x2), xmask, eviction_policy='evict_last'
        )
    tmp4 = tl.load(in_ptr2 + x0, xmask, eviction_policy='evict_last')
    tmp7 = tl.load(in_ptr3 + x2, xmask, eviction_policy='evict_last')
    tmp2 = tmp0 + tmp1
    tmp5 = tmp3 + tmp4
    tmp6 = tmp2 + tmp5
    tmp8 = tmp6 * tmp7
    tl.store(in_out_ptr0 + x3, tmp8, 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, (5, 4), (4, 1))
    assert_size_stride(primals_3, (5,), (1,))
    assert_size_stride(primals_4, (25, 4), (4, 1))
    assert_size_stride(primals_5, (25,), (1,))
    assert_size_stride(primals_6, (4, 4), (4, 1))
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        buf0 = empty_strided_cuda((16, 5), (5, 1), torch.float32)
        extern_kernels.mm(reinterpret_tensor(primals_1, (16, 4), (4, 1), 0),
            reinterpret_tensor(primals_2, (4, 5), (1, 4), 0), out=buf0)
        del primals_2
        buf1 = empty_strided_cuda((16, 25), (25, 1), torch.float32)
        extern_kernels.mm(reinterpret_tensor(primals_1, (16, 4), (4, 1), 0),
            reinterpret_tensor(primals_4, (4, 25), (1, 4), 0), out=buf1)
        del primals_4
        buf2 = reinterpret_tensor(buf1, (4, 4, 5, 5), (100, 25, 5, 1), 0)
        del buf1
        get_raw_stream(0)
        triton_poi_fused_add_mul_0[grid(400)](buf2, primals_5, buf0,
            primals_3, primals_6, 400, XBLOCK=128, num_warps=4, num_stages=1)
        del buf0
        del primals_3
        del primals_5
    return buf2, primals_6, reinterpret_tensor(primals_1, (16, 4), (4, 1), 0)


def logsumexp(x, dim=None):
    """
    Args:
        x: A pytorch tensor (any dimension will do)
        dim: int or None, over which to perform the summation. `None`, the
             default, performs over all axes.
    Returns: The result of the log(sum(exp(...))) operation.
    """
    if dim is None:
        xmax = x.max()
        xmax_ = x.max()
        return xmax_ + torch.log(torch.exp(x - xmax).sum())
    else:
        xmax, _ = x.max(dim, keepdim=True)
        xmax_, _ = x.max(dim)
        return xmax_ + torch.log(torch.exp(x - xmax).sum(dim))


class ChainCRFNew(nn.Module):

    def __init__(self, input_size, num_labels, bigram=True):
        """
        Args:
            input_size: int
                the dimension of the input.
            num_labels: int
                the number of labels of the crf layer
            bigram: bool
                if apply bi-gram parameter.
        """
        super(ChainCRFNew, self).__init__()
        self.input_size = input_size
        self.num_labels = num_labels + 1
        self.pad_label_id = num_labels
        self.bigram = bigram
        self.state_nn = nn.Linear(input_size, self.num_labels)
        if bigram:
            self.trans_nn = nn.Linear(input_size, self.num_labels * self.
                num_labels)
            self.register_parameter('trans_matrix', None)
        else:
            self.trans_nn = None
            self.trans_matrix = Parameter(torch.Tensor(self.num_labels,
                self.num_labels))
        self.reset_parameters()

    def reset_parameters(self):
        nn.init.constant_(self.state_nn.bias, 0.0)
        if self.bigram:
            nn.init.xavier_uniform_(self.trans_nn.weight)
            nn.init.constant_(self.trans_nn.bias, 0.0)
        else:
            nn.init.normal_(self.trans_matrix)

    def loss(self, energy, target, mask):
        """
        Args:
            energy: Tensor
                the energy tensor with shape = [batch, length, num_label, num_label]
            target: Tensor
                the tensor of target labels with shape [batch, length]
            mask:Tensor
                the mask tensor with shape = [batch, length]

        Returns: Tensor
                A 1D tensor for minus log likelihood loss
        """
        batch, length = target.size()
        energy_transpose = energy.transpose(0, 1)
        target_transpose = target.transpose(0, 1)
        mask_transpose = mask.unsqueeze(2).transpose(0, 1)
        partition = None
        batch_index = torch.arange(0, batch, dtype=torch.long, device=
            target.device)
        prev_label = torch.zeros(batch, dtype=torch.long, device=target.device)
        prev_label = prev_label.fill_(self.num_labels - 1)
        tgt_energy = torch.zeros(batch, device=target.device)
        for t in range(length):
            curr_energy = energy_transpose[t]
            mask_t = mask_transpose[t]
            if t == 0:
                partition = curr_energy[:, -1, :]
            else:
                partition_new = logsumexp(curr_energy + partition.unsqueeze
                    (2), dim=1)
                partition = partition + (partition_new - partition) * mask_t
            tgt_energy += curr_energy[batch_index, prev_label,
                target_transpose[t].data]
            prev_label_new = target_transpose[t].data
            prev_label = prev_label + (prev_label_new - prev_label
                ) * mask_t.squeeze(1).long()
        return (logsumexp(partition, dim=1) - tgt_energy).mean()

    def decode(self, energy, mask):
        """
        Args:
            energy: Tensor
                the energy tensor with shape = [batch, length, num_label, num_label]
            mask: Tensor
                the mask tensor with shape = [batch, length]

        Returns: Tensor
            decoding results in shape [batch, length]
        """
        energy_transpose = energy.transpose(0, 1)
        mask_transpose = mask.unsqueeze(2).transpose(0, 1).long()
        energy_transpose = energy_transpose[:, :, :-1, :-1]
        length, batch_size, num_label, _ = energy_transpose.size()
        batch_index = torch.arange(0, batch_size, dtype=torch.long, device=
            energy.device)
        pi = torch.zeros([length, batch_size, num_label, 1], device=energy.
            device)
        pointer = torch.zeros([length, batch_size, num_label], dtype=torch.
            long, device=energy.device)
        dummy_pointer = torch.arange(self.num_labels - 1, device=energy.device)
        back_pointer = torch.zeros([length, batch_size], dtype=torch.long,
            device=energy.device)
        pi[0] = energy[:, 0, -1, :-1].unsqueeze(2)
        pointer[0] = -1
        for t in range(1, length):
            pi_prev = pi[t - 1]
            mask_t = mask_transpose[t]
            pi_t, pointer_t = torch.max(energy_transpose[t] + pi_prev, dim=1)
            pointer[t] = pointer_t * mask_t + dummy_pointer * (1 - mask_t)
            pi[t] = (pi_t * mask_t).unsqueeze(2) + pi[t - 1] * (1 - mask_t.
                unsqueeze(2))
        _, back_pointer[-1] = torch.max(pi[-1].squeeze(2), dim=1)
        for t in reversed(range(length - 1)):
            pointer_last = pointer[t + 1]
            back_pointer[t] = pointer_last[batch_index, back_pointer[t + 1]]
        return back_pointer.transpose(0, 1)

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


Measured latency_ms: 0.041929