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
from torch._inductor.runtime import triton_helpers
from torch._inductor.runtime.triton_helpers import math as tl_math
import numpy as np
from torch import nn
from typing import Optional
assert_size_stride = torch._C._dynamo.guards.assert_size_stride
empty_strided_cuda = torch._C._dynamo.guards._empty_strided_cuda


@triton.jit
def triton_per_fused_clamp_gather_log_mean_neg_0(in_out_ptr0, in_ptr0,
    in_ptr1, xnumel, rnumel, XBLOCK: tl.constexpr):
    RBLOCK: tl.constexpr = 16
    xoffset = tl.program_id(0) * XBLOCK
    xoffset + tl.arange(0, XBLOCK)[:, None]
    tl.full([XBLOCK, RBLOCK], True, tl.int1)
    rindex = tl.arange(0, RBLOCK)[None, :]
    tl.full([XBLOCK, RBLOCK], True, tl.int1)
    r2 = rindex
    r1 = rindex // 4
    tmp0 = tl.load(in_ptr0 + r2, None)
    tmp1 = tl.full([XBLOCK, RBLOCK], 4, tl.int32)
    tmp2 = tmp0 + tmp1
    tmp3 = tmp0 < 0
    tmp4 = tl.where(tmp3, tmp2, tmp0)
    tl.device_assert((0 <= tmp4) & (tmp4 < 4),
        'index out of bounds: 0 <= tmp4 < 4')
    tmp6 = tl.load(in_ptr1 + (tmp4 + 4 * r1), None, eviction_policy=
        'evict_last')
    tmp7 = tmp6.to(tl.float32)
    tmp8 = 1e-15
    tmp9 = triton_helpers.maximum(tmp7, tmp8)
    tmp10 = 0.999999999999999
    tmp11 = triton_helpers.minimum(tmp9, tmp10)
    tmp12 = tmp11.to(tl.int64)
    tmp13 = tmp12.to(tl.float32)
    tmp14 = tl_math.log(tmp13)
    tmp15 = -tmp14
    tmp16 = tl.broadcast_to(tmp15, [XBLOCK, RBLOCK])
    tmp18 = tl.sum(tmp16, 1)[:, None]
    tmp19 = 16.0
    tmp20 = tmp18 / tmp19
    tl.debug_barrier()
    tl.store(in_out_ptr0 + tl.full([XBLOCK, 1], 0, tl.int32), tmp20, None)


def call(args):
    arg0_1, arg1_1 = args
    args.clear()
    assert_size_stride(arg0_1, (4, 4), (4, 1))
    assert_size_stride(arg1_1, (4, 4), (4, 1))
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        buf0 = empty_strided_cuda((), (), torch.float32)
        buf1 = buf0
        del buf0
        get_raw_stream(0)
        triton_per_fused_clamp_gather_log_mean_neg_0[grid(1)](buf1, arg1_1,
            arg0_1, 1, 16, XBLOCK=1, num_warps=2, num_stages=1)
        del arg0_1
        del arg1_1
    return buf1,


def _reduction(loss: 'torch.Tensor', reduction: 'str') ->torch.Tensor:
    """
    Reduce loss

    Parameters
    ----------
    loss : torch.Tensor, [batch_size, num_classes]
        Batch losses.
    reduction : str
        Method for reducing the loss. Options include 'elementwise_mean',
        'none', and 'sum'.

    Returns
    -------
    loss : torch.Tensor
        Reduced loss.

    """
    if reduction == 'elementwise_mean':
        return loss.mean()
    elif reduction == 'none':
        return loss
    elif reduction == 'sum':
        return loss.sum()
    else:
        raise ValueError(f'{reduction} is not a valid reduction')


def cumulative_link_loss(y_pred: 'torch.Tensor', y_true: 'torch.Tensor',
    reduction: 'str'='elementwise_mean', class_weights:
    'Optional[np.ndarray]'=None) ->torch.Tensor:
    """
    Calculates the negative log likelihood using the logistic cumulative link
    function.

    See "On the consistency of ordinal regression methods", Pedregosa et. al.
    for more details. While this paper is not the first to introduce this, it
    is the only one that I could find that was easily readable outside of
    paywalls.

    Parameters
    ----------
    y_pred : torch.Tensor, [batch_size, num_classes]
        Predicted target class probabilities. float dtype.
    y_true : torch.Tensor, [batch_size, 1]
        True target classes. long dtype.
    reduction : str
        Method for reducing the loss. Options include 'elementwise_mean',
        'none', and 'sum'.
    class_weights : np.ndarray, [num_classes] optional (default=None)
        An array of weights for each class. If included, then for each sample,
        look up the true class and multiply that sample's loss by the weight in
        this array.

    Returns
    -------
    loss: torch.Tensor

    """
    eps = 1e-15
    likelihoods = torch.clamp(torch.gather(y_pred, 1, y_true), eps, 1 - eps)
    neg_log_likelihood = -torch.log(likelihoods)
    if class_weights is not None:
        class_weights = torch.as_tensor(class_weights, dtype=
            neg_log_likelihood.dtype, device=neg_log_likelihood.device)
        neg_log_likelihood *= class_weights[y_true]
    loss = _reduction(neg_log_likelihood, reduction)
    return loss


class CumulativeLinkLossNew(nn.Module):
    """
    Module form of cumulative_link_loss() loss function

    Parameters
    ----------
    reduction : str
        Method for reducing the loss. Options include 'elementwise_mean',
        'none', and 'sum'.
    class_weights : np.ndarray, [num_classes] optional (default=None)
        An array of weights for each class. If included, then for each sample,
        look up the true class and multiply that sample's loss by the weight in
        this array.

    """

    def __init__(self, reduction: 'str'='elementwise_mean', class_weights:
        'Optional[torch.Tensor]'=None) ->None:
        super().__init__()
        self.class_weights = class_weights
        self.reduction = reduction

    def forward(self, input_0, input_1):
        arg0_1 = input_0
        arg1_1 = input_1
        output = call([arg0_1, arg1_1])
        return output[0]


Measured latency_ms: 0.015137