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 typing
from torch import Tensor
from collections import Counter
from typing import List
from typing import Optional
from typing import Union
from torch.utils.data import Dataset
import torch.utils.data.dataloader
from torch import nn
import torch.nn
from torch.utils.data.dataset import Dataset
from torch.utils.data import DataLoader
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_neg_stack_0(in_out_ptr0, in_ptr0, in_ptr1, xnumel,
    XBLOCK: tl.constexpr):
    xnumel = 256
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x1 = xindex // 4 % 16
    x0 = xindex % 4
    x2 = xindex // 64
    x3 = xindex
    tmp0 = x1
    tl.full([1], 0, tl.int64)
    tmp3 = tl.full([1], 4, tl.int64)
    tmp4 = tmp0 < tmp3
    tmp5 = tl.load(in_ptr0 + (x0 + 4 * x1 + 64 * x2), tmp4 & xmask, other=0.0)
    tmp6 = tl.load(in_ptr1 + x0, tmp4 & xmask, eviction_policy='evict_last',
        other=0.0)
    tmp7 = tmp5 - tmp6
    tmp8 = tmp7 * tmp7
    tmp9 = tl.load(in_ptr0 + (16 + x0 + 4 * x1 + 64 * x2), tmp4 & xmask,
        other=0.0)
    tmp10 = tmp9 - tmp6
    tmp11 = tmp10 * tmp10
    tmp12 = tmp8 + tmp11
    tmp13 = tl.load(in_ptr0 + (32 + x0 + 4 * x1 + 64 * x2), tmp4 & xmask,
        other=0.0)
    tmp14 = tmp13 - tmp6
    tmp15 = tmp14 * tmp14
    tmp16 = tmp12 + tmp15
    tmp17 = tl.load(in_ptr0 + (48 + x0 + 4 * x1 + 64 * x2), tmp4 & xmask,
        other=0.0)
    tmp18 = tmp17 - tmp6
    tmp19 = tmp18 * tmp18
    tmp20 = tmp16 + tmp19
    tmp21 = tl.full(tmp20.shape, 0.0, tmp20.dtype)
    tmp22 = tl.where(tmp4, tmp20, tmp21)
    tmp23 = tmp0 >= tmp3
    tmp24 = tl.full([1], 8, tl.int64)
    tmp25 = tmp0 < tmp24
    tmp26 = tmp23 & tmp25
    tmp27 = tl.load(in_ptr0 + (x0 + 4 * (-4 + x1) + 64 * x2), tmp26 & xmask,
        other=0.0)
    tmp28 = tl.load(in_ptr1 + (4 + x0), tmp26 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp29 = tmp27 - tmp28
    tmp30 = tmp29 * tmp29
    tmp31 = tl.load(in_ptr0 + (16 + x0 + 4 * (-4 + x1) + 64 * x2), tmp26 &
        xmask, other=0.0)
    tmp32 = tmp31 - tmp28
    tmp33 = tmp32 * tmp32
    tmp34 = tmp30 + tmp33
    tmp35 = tl.load(in_ptr0 + (32 + x0 + 4 * (-4 + x1) + 64 * x2), tmp26 &
        xmask, other=0.0)
    tmp36 = tmp35 - tmp28
    tmp37 = tmp36 * tmp36
    tmp38 = tmp34 + tmp37
    tmp39 = tl.load(in_ptr0 + (48 + x0 + 4 * (-4 + x1) + 64 * x2), tmp26 &
        xmask, other=0.0)
    tmp40 = tmp39 - tmp28
    tmp41 = tmp40 * tmp40
    tmp42 = tmp38 + tmp41
    tmp43 = tl.full(tmp42.shape, 0.0, tmp42.dtype)
    tmp44 = tl.where(tmp26, tmp42, tmp43)
    tmp45 = tmp0 >= tmp24
    tmp46 = tl.full([1], 12, tl.int64)
    tmp47 = tmp0 < tmp46
    tmp48 = tmp45 & tmp47
    tmp49 = tl.load(in_ptr0 + (x0 + 4 * (-8 + x1) + 64 * x2), tmp48 & xmask,
        other=0.0)
    tmp50 = tl.load(in_ptr1 + (8 + x0), tmp48 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp51 = tmp49 - tmp50
    tmp52 = tmp51 * tmp51
    tmp53 = tl.load(in_ptr0 + (16 + x0 + 4 * (-8 + x1) + 64 * x2), tmp48 &
        xmask, other=0.0)
    tmp54 = tmp53 - tmp50
    tmp55 = tmp54 * tmp54
    tmp56 = tmp52 + tmp55
    tmp57 = tl.load(in_ptr0 + (32 + x0 + 4 * (-8 + x1) + 64 * x2), tmp48 &
        xmask, other=0.0)
    tmp58 = tmp57 - tmp50
    tmp59 = tmp58 * tmp58
    tmp60 = tmp56 + tmp59
    tmp61 = tl.load(in_ptr0 + (48 + x0 + 4 * (-8 + x1) + 64 * x2), tmp48 &
        xmask, other=0.0)
    tmp62 = tmp61 - tmp50
    tmp63 = tmp62 * tmp62
    tmp64 = tmp60 + tmp63
    tmp65 = tl.full(tmp64.shape, 0.0, tmp64.dtype)
    tmp66 = tl.where(tmp48, tmp64, tmp65)
    tmp67 = tmp0 >= tmp46
    tl.full([1], 16, tl.int64)
    tmp70 = tl.load(in_ptr0 + (x0 + 4 * (-12 + x1) + 64 * x2), tmp67 &
        xmask, other=0.0)
    tmp71 = tl.load(in_ptr1 + (12 + x0), tmp67 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp72 = tmp70 - tmp71
    tmp73 = tmp72 * tmp72
    tmp74 = tl.load(in_ptr0 + (16 + x0 + 4 * (-12 + x1) + 64 * x2), tmp67 &
        xmask, other=0.0)
    tmp75 = tmp74 - tmp71
    tmp76 = tmp75 * tmp75
    tmp77 = tmp73 + tmp76
    tmp78 = tl.load(in_ptr0 + (32 + x0 + 4 * (-12 + x1) + 64 * x2), tmp67 &
        xmask, other=0.0)
    tmp79 = tmp78 - tmp71
    tmp80 = tmp79 * tmp79
    tmp81 = tmp77 + tmp80
    tmp82 = tl.load(in_ptr0 + (48 + x0 + 4 * (-12 + x1) + 64 * x2), tmp67 &
        xmask, other=0.0)
    tmp83 = tmp82 - tmp71
    tmp84 = tmp83 * tmp83
    tmp85 = tmp81 + tmp84
    tmp86 = tl.full(tmp85.shape, 0.0, tmp85.dtype)
    tmp87 = tl.where(tmp67, tmp85, tmp86)
    tmp88 = tl.where(tmp48, tmp66, tmp87)
    tmp89 = tl.where(tmp26, tmp44, tmp88)
    tmp90 = tl.where(tmp4, tmp22, tmp89)
    tmp91 = -tmp90
    tl.store(in_out_ptr0 + x3, tmp91, xmask)


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


def dot_product(a: 'torch.Tensor', b: 'torch.Tensor', normalize=False):
    """
    Computes dot product for pairs of vectors.
    :param normalize: Vectors are normalized (leads to cosine similarity)
    :return: Matrix with res[i][j]  = dot_product(a[i], b[j])
    """
    if len(a.shape) == 1:
        a = a.unsqueeze(0)
    if len(b.shape) == 1:
        b = b.unsqueeze(0)
    if normalize:
        a = torch.nn.functional.normalize(a, p=2, dim=1)
        b = torch.nn.functional.normalize(b, p=2, dim=1)
    return torch.mm(a, b.transpose(0, 1))


def arccosh(x):
    """Compute the arcosh, numerically stable."""
    x = torch.clamp(x, min=1 + EPSILON)
    a = torch.log(x)
    b = torch.log1p(torch.sqrt(x * x - 1) / x)
    return a + b


def _iter_dataset(dataset: 'Optional[Dataset]') ->typing.Iterable:
    if dataset is None:
        return []
    return map(lambda x: x[0], DataLoader(dataset, batch_size=1, num_workers=0)
        )


def identify_dynamic_embeddings(data_point: 'DataPoint'):
    dynamic_embeddings = []
    if isinstance(data_point, Sentence):
        first_token = data_point[0]
        for name, vector in first_token._embeddings.items():
            if vector.requires_grad:
                dynamic_embeddings.append(name)
    for name, vector in data_point._embeddings.items():
        if vector.requires_grad:
            dynamic_embeddings.append(name)
    return dynamic_embeddings


def store_embeddings(data_points: 'Union[List[DT], Dataset]', storage_mode:
    'str', dynamic_embeddings: 'Optional[List[str]]'=None):
    if isinstance(data_points, Dataset):
        data_points = list(_iter_dataset(data_points))
    if storage_mode == 'none':
        dynamic_embeddings = None
    elif not dynamic_embeddings:
        dynamic_embeddings = identify_dynamic_embeddings(data_points[0])
    for data_point in data_points:
        data_point.clear_embeddings(dynamic_embeddings)
    if storage_mode == 'cpu':
        str(flair.device) != 'cpu'
        for data_point in data_points:
            data_point


def mdot(x, y):
    """Compute the inner product."""
    m = x.new_ones(1, x.size(1))
    m[0, 0] = -1
    return torch.sum(m * x * y, 1, keepdim=True)


def dist(x, y):
    """Get the hyperbolic distance between x and y."""
    return arccosh(-mdot(x, y))


class CosineDistance(torch.nn.Module):

    def forward(self, a, b):
        return -dot_product(a, b, normalize=True)


class EuclideanDistance(nn.Module):
    """Implement a EuclideanDistance object."""

    def forward(self, mat_1: 'Tensor', mat_2: 'Tensor') ->Tensor:
        """Returns the squared euclidean distance between each
        element in mat_1 and each element in mat_2.

        Parameters
        ----------
        mat_1: torch.Tensor
            matrix of shape (n_1, n_features)
        mat_2: torch.Tensor
            matrix of shape (n_2, n_features)

        Returns
        -------
        dist: torch.Tensor
            distance matrix of shape (n_1, n_2)

        """
        _dist = [torch.sum((mat_1 - mat_2[i]) ** 2, dim=1) for i in range(
            mat_2.size(0))]
        dist = torch.stack(_dist, dim=1)
        return dist


class HyperbolicDistance(nn.Module):
    """Implement a HyperbolicDistance object."""

    def forward(self, mat_1: 'Tensor', mat_2: 'Tensor') ->Tensor:
        """Returns the squared euclidean distance between each
        element in mat_1 and each element in mat_2.

        Parameters
        ----------
        mat_1: torch.Tensor
            matrix of shape (n_1, n_features)
        mat_2: torch.Tensor
            matrix of shape (n_2, n_features)

        Returns
        -------
        dist: torch.Tensor
            distance matrix of shape (n_1, n_2)

        """
        mat_1_x_0 = torch.sqrt(1 + mat_1.pow(2).sum(dim=1, keepdim=True))
        mat_2_x_0 = torch.sqrt(1 + mat_2.pow(2).sum(dim=1, keepdim=True))
        left = mat_1_x_0.mm(mat_2_x_0.t())
        right = mat_1[:, 1:].mm(mat_2[:, 1:].t())
        return arccosh(left - right).pow(2)


class LogitCosineDistance(torch.nn.Module):

    def forward(self, a, b):
        return torch.logit(0.5 - 0.5 * dot_product(a, b, normalize=True))


class NegativeScaledDotProduct(torch.nn.Module):

    def forward(self, a, b):
        sqrt_d = torch.sqrt(torch.tensor(a.size(-1)))
        return -dot_product(a, b, normalize=False) / sqrt_d


class PrototypicalDecoderNew(torch.nn.Module):

    def __init__(self, num_prototypes: 'int', embeddings_size: 'int',
        prototype_size: 'Optional[int]'=None, distance_function: 'str'=
        'euclidean', use_radius: 'Optional[bool]'=False, min_radius:
        'Optional[int]'=0, unlabeled_distance: 'Optional[float]'=None,
        unlabeled_idx: 'Optional[int]'=None, learning_mode: 'Optional[str]'
        ='joint', normal_distributed_initial_prototypes: 'bool'=False):
        super().__init__()
        if not prototype_size:
            prototype_size = embeddings_size
        self.prototype_size = prototype_size
        self.metric_space_decoder: 'Optional[torch.nn.Linear]' = None
        if prototype_size != embeddings_size:
            self.metric_space_decoder = torch.nn.Linear(embeddings_size,
                prototype_size)
            torch.nn.init.xavier_uniform_(self.metric_space_decoder.weight)
        self.prototype_vectors = torch.nn.Parameter(torch.ones(
            num_prototypes, prototype_size), requires_grad=True)
        if normal_distributed_initial_prototypes:
            self.prototype_vectors = torch.nn.Parameter(torch.normal(torch.
                zeros(num_prototypes, prototype_size)))
        self.prototype_radii: 'Optional[torch.nn.Parameter]' = None
        if use_radius:
            self.prototype_radii = torch.nn.Parameter(torch.ones(
                num_prototypes), requires_grad=True)
        self.min_radius = min_radius
        self.learning_mode = learning_mode
        assert (unlabeled_idx is None) == (unlabeled_distance is None
            ), "'unlabeled_idx' and 'unlabeled_distance' should either both be set or both not be set."
        self.unlabeled_idx = unlabeled_idx
        self.unlabeled_distance = unlabeled_distance
        self._distance_function = distance_function
        self.distance: 'Optional[torch.nn.Module]' = None
        if distance_function.lower() == 'hyperbolic':
            self.distance = HyperbolicDistance()
        elif distance_function.lower() == 'cosine':
            self.distance = CosineDistance()
        elif distance_function.lower() == 'logit_cosine':
            self.distance = LogitCosineDistance()
        elif distance_function.lower() == 'euclidean':
            self.distance = EuclideanDistance()
        elif distance_function.lower() == 'dot_product':
            self.distance = NegativeScaledDotProduct()
        else:
            raise KeyError(f'Distance function {distance_function} not found.')
        self

    @property
    def num_prototypes(self):
        return self.prototype_vectors.size(0)

    def enable_expectation_maximization(self, data: 'FlairDataset', encoder:
        'DefaultClassifier', exempt_labels: 'List[str]'=[], mini_batch_size:
        'int'=8):
        """Applies monkey-patch to train method (which sets the train flag).

        This allows for computation of average prototypes after a training
        sequence."""
        decoder = self
        unpatched_train = encoder.train

        def patched_train(mode: 'bool'=True):
            unpatched_train(mode=mode)
            if mode:
                logger.info('recalculating prototypes')
                with torch.no_grad():
                    decoder.calculate_prototypes(data=data, encoder=encoder,
                        exempt_labels=exempt_labels, mini_batch_size=
                        mini_batch_size)
        encoder.train = patched_train

    def calculate_prototypes(self, data: 'FlairDataset', encoder:
        'DefaultClassifier', exempt_labels: 'List[str]'=[], mini_batch_size=32
        ):
        """
        Function that calclues a prototype for each class based on the euclidean average embedding over the whole dataset
        :param data: dataset for which to calculate prototypes
        :param encoder: encoder to use
        :param exempt_labels: labels to exclude
        :param mini_batch_size: number of sentences to embed at same time
        :return:
        """
        with torch.no_grad():
            dataloader = DataLoader(data, batch_size=mini_batch_size)
            new_prototypes = torch.zeros(self.num_prototypes, self.
                prototype_size, device=flair.device)
            counter: 'Counter' = Counter()
            for batch in tqdm(dataloader):
                logits, labels = encoder.forward_pass(batch)
                if len(labels) > 0:
                    if self.metric_space_decoder is not None:
                        logits = self.metric_space_decoder(logits)
                    for logit, label in zip(logits, labels):
                        counter.update(label)
                        idx = encoder.label_dictionary.get_idx_for_item(label
                            [0])
                        new_prototypes[idx] += logit
                store_embeddings(batch, storage_mode='none')
            for label, count in counter.most_common():
                average_prototype = new_prototypes[encoder.label_dictionary
                    .get_idx_for_item(label)] / count
                new_prototypes[encoder.label_dictionary.get_idx_for_item(label)
                    ] = average_prototype
            for label in exempt_labels:
                label_idx = encoder.label_dictionary.get_idx_for_item(label)
                new_prototypes[label_idx] = self.prototype_vectors[label_idx]
            self.prototype_vectors.data = new_prototypes

    def forward(self, input_0):
        primals_2 = self.prototype_vectors
        primals_1 = input_0
        output = call([primals_1, primals_2])
        return output[0]


Measured latency_ms: 0.015812