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
import numpy as np
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


@triton.jit
def triton_poi_fused_repeat_0(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex
    tmp0 = tl.load(in_ptr0 + x0, xmask)
    tl.store(out_ptr0 + x0, tmp0, xmask)


@triton.jit
def triton_poi_fused_repeat_1(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 4
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex
    tmp0 = tl.load(in_ptr0 + x0, xmask)
    tl.store(out_ptr0 + x0, tmp0, xmask)


def call(args):
    primals_1, primals_2 = args
    args.clear()
    assert_size_stride(primals_1, (5, 4), (4, 1))
    assert_size_stride(primals_2, (5, 1), (1, 1))
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        buf0 = empty_strided_cuda((1, 4, 4), (16, 4, 1), torch.float32)
        get_raw_stream(0)
        triton_poi_fused_repeat_0[grid(16)](primals_1, buf0, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        del primals_1
        buf1 = empty_strided_cuda((1, 4, 1), (4, 1, 1), torch.float32)
        triton_poi_fused_repeat_1[grid(4)](primals_2, buf1, 4, XBLOCK=4,
            num_warps=1, num_stages=1)
        del primals_2
    return buf0, buf1


class TransformerLinearXMCHeadNew(nn.Module):
    """XMC head for Transformers

    Containing label weight embeddings and label bias embeddings
    """

    def __init__(self, hidden_size, num_labels):
        super().__init__()
        self.label_pad = num_labels
        self.num_labels = num_labels
        self.W = nn.Embedding(num_labels + 1, hidden_size, padding_idx=self
            .label_pad)
        self.b = nn.Embedding(num_labels + 1, 1, padding_idx=self.label_pad)
        self.random_init()

    @property
    def device(self):
        return self.W.weight.device

    def random_init(self):
        """Initialize the weight and bias embeddings

        Initialize label weight embedding with N(0, 0.02) while keeping PAD
        column to be 0. Initialize label bias embedding with 0.
        """
        mat = 0.02 * np.random.randn(self.label_pad, self.W.weight.shape[1])
        mat = np.hstack([mat, np.zeros([mat.shape[0], 1])])
        self.init_from(mat)

    def inherit(self, prev_head, C):
        prev_W = prev_head.W.weight[:-1, :].detach().numpy()
        prev_b = prev_head.b.weight[:-1, :].detach().numpy()
        cur_W = C * prev_W
        cur_b = C * prev_b
        mat = np.hstack([cur_W, cur_b])
        self.init_from(mat)

    def bootstrap(self, prob, **kwargs):
        """Initialize head with weights learned from linear model using transformer embeddings

        Args:
            prob (MLProblem): the multi-label problem to bootstrap with
            kwargs:
                Cp (float): the weight on positive samples. Default 100.0
                Cn (float): the weight on negative samples. Default 100.0
                threshold (float): the threshold to sparsify the model
        """
        Cp = kwargs.get('Cp', 100.0)
        Cn = kwargs.get('Cn', 100.0)
        threshold = kwargs.get('threshold', 0)
        mat = MLModel.train(prob, threshold=threshold, Cp=Cp, Cn=Cn)
        mat = mat.W.toarray().T
        self.init_from(mat)

    def init_from(self, mat):
        """Initialize the weight and bias embeddings with given matrix

        Args:
            mat (ndarray): matrix used for initialize, shape = (nr_labels, hidden_size + 1)
        """
        if not isinstance(mat, np.ndarray):
            raise ValueError('Expect ndarray to initialize label embedding')
        if mat.shape[0] != self.label_pad:
            raise ValueError('nr_labels mismatch!')
        mat = np.vstack([mat, np.zeros([1, mat.shape[1]])])
        self.W = nn.Embedding.from_pretrained(torch.FloatTensor(mat[:, :-1]
            ), freeze=False, sparse=True, padding_idx=self.label_pad)
        self.b = nn.Embedding.from_pretrained(torch.FloatTensor(mat[:, -1])
            .view((self.label_pad + 1, 1)), freeze=False, sparse=True,
            padding_idx=self.label_pad)

    def forward(self):
        primals_1 = self.W.weight
        primals_2 = self.b.weight
        output = call([primals_1, primals_2])
        return output[0], output[1]


Measured latency_ms: 0.015012