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
import torch.nn as nn
import torch.nn.modules
import torch.nn
import torch.nn.init
import torch.optim.lr_scheduler
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_add_clamp_div_mul_pow_rsub_sum_0(in_out_ptr0, in_ptr0,
    in_ptr1, xnumel, rnumel):
    XBLOCK: tl.constexpr = 1
    RBLOCK: tl.constexpr = 256
    xoffset = tl.program_id(0) * XBLOCK
    tl.full([1], xoffset, tl.int32)
    tl.full([RBLOCK], True, tl.int1)
    rindex = tl.arange(0, RBLOCK)[:]
    tl.full([RBLOCK], True, tl.int1)
    r0 = rindex
    tmp0 = tl.load(in_ptr0 + r0, None)
    tmp1 = tl.load(in_ptr1 + r0, None)
    tmp2 = tmp1 * tmp1
    tmp3 = tmp0 * tmp2
    tmp4 = 1.0
    tmp5 = tmp4 - tmp0
    tmp6 = tmp4 - tmp1
    tmp7 = 0.0
    tmp8 = triton_helpers.maximum(tmp6, tmp7)
    tmp9 = tmp8 * tmp8
    tmp10 = tmp5 * tmp9
    tmp11 = tmp3 + tmp10
    tmp12 = tl.broadcast_to(tmp11, [RBLOCK])
    tmp14 = triton_helpers.promote_to_tensor(tl.sum(tmp12, 0))
    tmp15 = 0.5
    tmp16 = tmp14 * tmp15
    tmp17 = 0.25
    tmp18 = tmp16 * tmp17
    tl.debug_barrier()
    tl.store(in_out_ptr0 + tl.full([1], 0, tl.int32), tmp18, 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((), (), torch.float32)
        buf1 = buf0
        del buf0
        get_raw_stream(0)
        triton_per_fused_add_clamp_div_mul_pow_rsub_sum_0[grid(1)](buf1,
            arg1_1, arg0_1, 1, 256, num_warps=2, num_stages=1)
        del arg0_1
        del arg1_1
    return buf1,


class ContrastiveLossNew(nn.Module):
    """
    Contrastive loss function.

    References:
        https://github.com/delijati/pytorch-siamese/blob/master/contrastive.py

    LaTeX:
        Let D be the distance computed between the network layers or the direct
        distance output of the network.

        y is 0 if the pair should be labled as an imposter
        y is 1 if the pair should be labled as genuine

        ContrastiveLoss = ((y * D) ** 2 + ((1 - y) * max(m - D, 0) ** 2)) / 2

        $(y E)^2 + ((1 - y) max(m - E, 0)^2)$

    CommandLine:
        python -m clab.criterions ContrastiveLoss --show


    DisableExample:
        >>> # DISABLE_DOCTEST
        >>> from clab.criterions import *
        >>> import utool as ut
        >>> import numpy as np
        >>> vecs1, vecs2, label = testdata_siam_desc()
        >>> output = torch.nn.PairwiseDistance(p=2)(vecs1, vecs2)
        >>> self = ContrastiveLoss()
        >>> ut.exec_func_src(self.forward, globals())
        >>> func = self.forward
        >>> loss2x, dist = ut.exec_func_src(self.forward, globals(), globals(), keys=['loss2x', 'dist'])
        >>> ut.quit_if_noshow()
        >>> loss2x, dist, label = map(np.array, [loss2x, dist, label])
        >>> label = label.astype(np.bool)
        >>> dist0_l2 = dist[~label]
        >>> dist1_l2 = dist[label]
        >>> loss0 = loss2x[~label] / 2
        >>> loss1 = loss2x[label] / 2
        >>> # xdoc: +REQUIRES(--show)
        >>> import plottool as pt
        >>> pt.plot2(dist0_l2, loss0, 'x', color=pt.FALSE_RED, label='imposter_loss', y_label='loss')
        >>> pt.plot2(dist1_l2, loss1, 'x', color=pt.TRUE_BLUE, label='genuine_loss', y_label='loss')
        >>> pt.gca().set_xlabel('l2-dist')
        >>> pt.legend()
        >>> ut.show_if_requested()

    Example:
        >>> # DISABLE_DOCTEST
        >>> import torch
        >>> import netharn as nh
        >>> xpu = nh.XPU(None)
        >>> imgs1 = xpu.variable(torch.rand(1, 3, 224, 244))
        >>> imgs2 = xpu.variable(torch.rand(1, 3, 224, 244))
        >>> label = (xpu.variable(torch.rand(3)) * 2).long()

        >>> model = SiameseLP(input_shape=imgs1.shape[1:])
        >>> output = model(imgs1, imgs2)
        >>> self = ContrastiveLoss(margin=10)
        >>> self.forward(output, label)
    """

    def __init__(self, weight=None, margin=1.0):
        super(ContrastiveLossNew, self).__init__()
        self.weight = weight
        self.margin = margin
        self.neg_label = 0
        self.pos_label = 1

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