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
from torch import 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_cat_logsumexp_0(in_ptr0, in_ptr1, out_ptr0, out_ptr1,
    out_ptr2, xnumel, XBLOCK: tl.constexpr):
    xnumel = 64
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex
    tmp0 = tl.full([1], 0, tl.int64)
    tmp2 = tl.full([1], 4, tl.int64)
    tmp3 = tmp0 < tmp2
    tmp4 = tl.load(in_ptr0 + (4 * x0 + 0), tmp3 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp5 = 2.0
    tmp6 = tmp4 * tmp5
    tmp7 = 1.0
    tmp8 = tmp7 - tmp6
    tmp9 = tl.load(in_ptr1 + (4 * x0 + 0), tmp3 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp10 = tmp8 * tmp9
    tmp11 = 1000000000000.0
    tmp12 = tmp4 * tmp11
    tmp13 = tmp10 - tmp12
    tmp14 = tl.full(tmp13.shape, 0.0, tmp13.dtype)
    tmp15 = tl.where(tmp3, tmp13, tmp14)
    tmp16 = tmp0 >= tmp2
    tl.full([1], 5, tl.int64)
    tmp19 = 0.0
    tmp20 = tl.full(tmp19.shape, 0.0, tmp19.dtype)
    tmp21 = tl.where(tmp16, tmp19, tmp20)
    tmp22 = tl.where(tmp3, tmp15, tmp21)
    tmp23 = tl.full([1], 1, tl.int64)
    tmp25 = tmp23 < tmp2
    tmp26 = tl.load(in_ptr0 + (4 * x0 + 1), tmp25 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp27 = tmp26 * tmp5
    tmp28 = tmp7 - tmp27
    tmp29 = tl.load(in_ptr1 + (4 * x0 + 1), tmp25 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp30 = tmp28 * tmp29
    tmp31 = tmp26 * tmp11
    tmp32 = tmp30 - tmp31
    tmp33 = tl.full(tmp32.shape, 0.0, tmp32.dtype)
    tmp34 = tl.where(tmp25, tmp32, tmp33)
    tmp35 = tmp23 >= tmp2
    tmp37 = tl.where(tmp35, tmp19, tmp20)
    tmp38 = tl.where(tmp25, tmp34, tmp37)
    tmp39 = triton_helpers.maximum(tmp22, tmp38)
    tmp40 = tl.full([1], 2, tl.int64)
    tmp42 = tmp40 < tmp2
    tmp43 = tl.load(in_ptr0 + (4 * x0 + 2), tmp42 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp44 = tmp43 * tmp5
    tmp45 = tmp7 - tmp44
    tmp46 = tl.load(in_ptr1 + (4 * x0 + 2), tmp42 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp47 = tmp45 * tmp46
    tmp48 = tmp43 * tmp11
    tmp49 = tmp47 - tmp48
    tmp50 = tl.full(tmp49.shape, 0.0, tmp49.dtype)
    tmp51 = tl.where(tmp42, tmp49, tmp50)
    tmp52 = tmp40 >= tmp2
    tmp54 = tl.where(tmp52, tmp19, tmp20)
    tmp55 = tl.where(tmp42, tmp51, tmp54)
    tmp56 = triton_helpers.maximum(tmp39, tmp55)
    tmp57 = tl.full([1], 3, tl.int64)
    tmp59 = tmp57 < tmp2
    tmp60 = tl.load(in_ptr0 + (4 * x0 + 3), tmp59 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp61 = tmp60 * tmp5
    tmp62 = tmp7 - tmp61
    tmp63 = tl.load(in_ptr1 + (4 * x0 + 3), tmp59 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp64 = tmp62 * tmp63
    tmp65 = tmp60 * tmp11
    tmp66 = tmp64 - tmp65
    tmp67 = tl.full(tmp66.shape, 0.0, tmp66.dtype)
    tmp68 = tl.where(tmp59, tmp66, tmp67)
    tmp69 = tmp57 >= tmp2
    tmp71 = tl.where(tmp69, tmp19, tmp20)
    tmp72 = tl.where(tmp59, tmp68, tmp71)
    tmp73 = triton_helpers.maximum(tmp56, tmp72)
    tmp75 = tmp2 < tmp2
    tmp76 = tl.load(in_ptr0 + (4 * x0 + 4), tmp75 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp77 = tmp76 * tmp5
    tmp78 = tmp7 - tmp77
    tmp79 = tl.load(in_ptr1 + (4 * x0 + 4), tmp75 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp80 = tmp78 * tmp79
    tmp81 = tmp76 * tmp11
    tmp82 = tmp80 - tmp81
    tmp83 = tl.full(tmp82.shape, 0.0, tmp82.dtype)
    tmp84 = tl.where(tmp75, tmp82, tmp83)
    tmp85 = tmp2 >= tmp2
    tmp87 = tl.where(tmp85, tmp19, tmp20)
    tmp88 = tl.where(tmp75, tmp84, tmp87)
    tmp89 = triton_helpers.maximum(tmp73, tmp88)
    tmp90 = tl_math.abs(tmp89)
    tmp91 = float('inf')
    tmp92 = tmp90 == tmp91
    tmp93 = tl.where(tmp92, tmp19, tmp89)
    tmp94 = tmp22 - tmp93
    tmp95 = tl_math.exp(tmp94)
    tmp96 = tmp38 - tmp93
    tmp97 = tl_math.exp(tmp96)
    tmp98 = tmp95 + tmp97
    tmp99 = tmp55 - tmp93
    tmp100 = tl_math.exp(tmp99)
    tmp101 = tmp98 + tmp100
    tmp102 = tmp72 - tmp93
    tmp103 = tl_math.exp(tmp102)
    tmp104 = tmp101 + tmp103
    tmp105 = tmp88 - tmp93
    tmp106 = tl_math.exp(tmp105)
    tmp107 = tmp104 + tmp106
    tmp108 = tmp7 - tmp4
    tmp109 = tmp108 * tmp11
    tmp110 = tmp10 - tmp109
    tmp111 = tl.full(tmp110.shape, 0.0, tmp110.dtype)
    tmp112 = tl.where(tmp3, tmp110, tmp111)
    tmp113 = tl.where(tmp3, tmp112, tmp21)
    tmp114 = tmp7 - tmp26
    tmp115 = tmp114 * tmp11
    tmp116 = tmp30 - tmp115
    tmp117 = tl.full(tmp116.shape, 0.0, tmp116.dtype)
    tmp118 = tl.where(tmp25, tmp116, tmp117)
    tmp119 = tl.where(tmp25, tmp118, tmp37)
    tmp120 = triton_helpers.maximum(tmp113, tmp119)
    tmp121 = tmp7 - tmp43
    tmp122 = tmp121 * tmp11
    tmp123 = tmp47 - tmp122
    tmp124 = tl.full(tmp123.shape, 0.0, tmp123.dtype)
    tmp125 = tl.where(tmp42, tmp123, tmp124)
    tmp126 = tl.where(tmp42, tmp125, tmp54)
    tmp127 = triton_helpers.maximum(tmp120, tmp126)
    tmp128 = tmp7 - tmp60
    tmp129 = tmp128 * tmp11
    tmp130 = tmp64 - tmp129
    tmp131 = tl.full(tmp130.shape, 0.0, tmp130.dtype)
    tmp132 = tl.where(tmp59, tmp130, tmp131)
    tmp133 = tl.where(tmp59, tmp132, tmp71)
    tmp134 = triton_helpers.maximum(tmp127, tmp133)
    tmp135 = tmp7 - tmp76
    tmp136 = tmp135 * tmp11
    tmp137 = tmp80 - tmp136
    tmp138 = tl.full(tmp137.shape, 0.0, tmp137.dtype)
    tmp139 = tl.where(tmp75, tmp137, tmp138)
    tmp140 = tl.where(tmp75, tmp139, tmp87)
    tmp141 = triton_helpers.maximum(tmp134, tmp140)
    tl.store(out_ptr0 + x0, tmp89, xmask)
    tl.store(out_ptr1 + x0, tmp107, xmask)
    tl.store(out_ptr2 + x0, tmp141, xmask)


@triton.jit
def triton_poi_fused_cat_logsumexp_1(in_ptr0, in_ptr1, in_ptr2, out_ptr0,
    xnumel, XBLOCK: tl.constexpr):
    xnumel = 320
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex % 5
    x1 = xindex // 5
    x2 = xindex
    tmp25 = tl.load(in_ptr2 + x1, xmask, eviction_policy='evict_last')
    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 = 2.0
    tmp7 = tmp5 * tmp6
    tmp8 = 1.0
    tmp9 = tmp8 - tmp7
    tmp10 = tl.load(in_ptr1 + (4 * x1 + x0), tmp4 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp11 = tmp9 * tmp10
    tmp12 = tmp8 - tmp5
    tmp13 = 1000000000000.0
    tmp14 = tmp12 * tmp13
    tmp15 = tmp11 - tmp14
    tmp16 = tl.full(tmp15.shape, 0.0, tmp15.dtype)
    tmp17 = tl.where(tmp4, tmp15, tmp16)
    tmp18 = tmp0 >= tmp3
    tl.full([1], 5, tl.int64)
    tmp21 = 0.0
    tmp22 = tl.full(tmp21.shape, 0.0, tmp21.dtype)
    tmp23 = tl.where(tmp18, tmp21, tmp22)
    tmp24 = tl.where(tmp4, tmp17, tmp23)
    tmp26 = tl_math.abs(tmp25)
    tmp27 = float('inf')
    tmp28 = tmp26 == tmp27
    tmp29 = tl.where(tmp28, tmp21, tmp25)
    tmp30 = tmp24 - tmp29
    tmp31 = tl_math.exp(tmp30)
    tl.store(out_ptr0 + x2, tmp31, xmask)


@triton.jit
def triton_per_fused_add_logsumexp_mean_2(in_out_ptr0, in_ptr0, in_ptr1,
    in_ptr2, in_ptr3, xnumel, rnumel, XBLOCK: tl.constexpr):
    RBLOCK: tl.constexpr = 64
    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)
    r0 = rindex
    tmp0 = tl.load(in_ptr0 + r0, None)
    tmp2 = tl.load(in_ptr1 + r0, None)
    tmp9 = tl.load(in_ptr2 + 5 * r0, None, eviction_policy='evict_last')
    tmp10 = tl.load(in_ptr2 + (1 + 5 * r0), None, eviction_policy='evict_last')
    tmp12 = tl.load(in_ptr2 + (2 + 5 * r0), None, eviction_policy='evict_last')
    tmp14 = tl.load(in_ptr2 + (3 + 5 * r0), None, eviction_policy='evict_last')
    tmp16 = tl.load(in_ptr2 + (4 + 5 * r0), None, eviction_policy='evict_last')
    tmp19 = tl.load(in_ptr3 + r0, None)
    tmp1 = tl_math.log(tmp0)
    tmp3 = tl_math.abs(tmp2)
    tmp4 = float('inf')
    tmp5 = tmp3 == tmp4
    tmp6 = 0.0
    tmp7 = tl.where(tmp5, tmp6, tmp2)
    tmp8 = tmp1 + tmp7
    tmp11 = tmp9 + tmp10
    tmp13 = tmp11 + tmp12
    tmp15 = tmp13 + tmp14
    tmp17 = tmp15 + tmp16
    tmp18 = tl_math.log(tmp17)
    tmp20 = tl_math.abs(tmp19)
    tmp21 = tmp20 == tmp4
    tmp22 = tl.where(tmp21, tmp6, tmp19)
    tmp23 = tmp18 + tmp22
    tmp24 = tmp8 + tmp23
    tmp25 = tl.broadcast_to(tmp24, [XBLOCK, RBLOCK])
    tmp27 = tl.sum(tmp25, 1)[:, None]
    tmp28 = 64.0
    tmp29 = tmp27 / tmp28
    tl.debug_barrier()
    tl.store(in_out_ptr0 + tl.full([XBLOCK, 1], 0, tl.int32), tmp29, None)


def call(args):
    arg0_1, arg1_1 = args
    args.clear()
    assert_size_stride(arg0_1, (4, 4, 4, 4), (64, 16, 4, 1))
    assert_size_stride(arg1_1, (4, 4, 4, 4), (64, 16, 4, 1))
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        buf0 = empty_strided_cuda((4, 4, 4, 1), (16, 4, 1, 64), torch.float32)
        buf1 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        buf2 = empty_strided_cuda((4, 4, 4, 1), (16, 4, 1, 64), torch.float32)
        get_raw_stream(0)
        triton_poi_fused_cat_logsumexp_0[grid(64)](arg0_1, arg1_1, buf0,
            buf1, buf2, 64, XBLOCK=64, num_warps=1, num_stages=1)
        buf3 = empty_strided_cuda((4, 4, 4, 5), (80, 20, 5, 1), torch.float32)
        triton_poi_fused_cat_logsumexp_1[grid(320)](arg0_1, arg1_1, buf2,
            buf3, 320, XBLOCK=128, num_warps=4, num_stages=1)
        del arg0_1
        del arg1_1
        buf4 = empty_strided_cuda((), (), torch.float32)
        buf5 = buf4
        del buf4
        triton_per_fused_add_logsumexp_mean_2[grid(1)](buf5, buf1, buf0,
            buf3, buf2, 1, 64, XBLOCK=1, num_warps=2, num_stages=1)
        del buf0
        del buf1
        del buf2
        del buf3
    return buf5,


class MultiLabelCircleLossNew(nn.Module):

    def __init__(self, reduction='mean', inf=1000000000000.0):
        """CircleLoss of MultiLabel, 多个目标类的多标签分类场景，希望“每个目标类得分都不小于每个非目标类的得分”
        多标签分类的交叉熵(softmax+crossentropy推广, N选K问题), LSE函数的梯度恰好是softmax函数
        让同类相似度与非同类相似度之间拉开一定的margin。
          - 使同类相似度比最大的非同类相似度更大。
          - 使最小的同类相似度比最大的非同类相似度更大。
          - 所有同类相似度都比所有非同类相似度更大。
        urls: [将“softmax+交叉熵”推广到多标签分类问题](https://spaces.ac.cn/archives/7359)
        args:
            reduction: str, Specifies the reduction to apply to the output, 输出形式. 
                            eg.``'none'`` | ``'mean'`` | ``'sum'``. ``'none'``
            inf: float, Minimum of maths, 无穷大.  eg. 1e12
        returns:
            Tensor of loss.
        examples:
            >>> label, logits = [[1, 1, 1, 1], [0, 0, 0, 1]], [[0, 1, 1, 0], [1, 0, 0, 1],]
            >>> label, logits = torch.tensor(label).float(), torch.tensor(logits).float()
            >>> loss = MultiLabelCircleLoss()(logits, label)
        """
        super(MultiLabelCircleLossNew, self).__init__()
        self.reduction = reduction
        self.inf = inf

    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.046922