import torch
from torch._inductor.select_algorithm import extern_kernels
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.triton_helpers import math as tl_math
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_sub_0(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 8
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x2 = xindex
    x0 = xindex % 4
    tmp0 = tl.load(in_ptr0 + (8 + x2), xmask)
    tmp1 = tl.load(in_ptr0 + (4 + x0), xmask, eviction_policy='evict_last')
    tmp2 = tmp0 - tmp1
    tl.store(out_ptr0 + x2, tmp2, xmask)


@triton.jit
def triton_per_fused_exp_sum_1(in_ptr0, out_ptr0, xnumel, rnumel, XBLOCK:
    tl.constexpr):
    RBLOCK: tl.constexpr = 2
    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)
    tmp1 = tl_math.exp(tmp0)
    tmp2 = tl.broadcast_to(tmp1, [XBLOCK, RBLOCK])
    tmp4 = tl.sum(tmp2, 1)[:, None]
    tl.store(out_ptr0 + tl.full([XBLOCK, 1], 0, tl.int32), tmp4, None)


@triton.jit
def triton_poi_fused_sub_2(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 8
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x2 = xindex
    x0 = xindex % 4
    tmp0 = tl.load(in_ptr0 + (24 + x2), xmask)
    tmp1 = tl.load(in_ptr0 + (20 + x0), xmask, eviction_policy='evict_last')
    tmp2 = tmp0 - tmp1
    tl.store(out_ptr0 + x2, tmp2, xmask)


@triton.jit
def triton_poi_fused_sub_3(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 8
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x2 = xindex
    x0 = xindex % 4
    tmp0 = tl.load(in_ptr0 + (40 + x2), xmask)
    tmp1 = tl.load(in_ptr0 + (36 + x0), xmask, eviction_policy='evict_last')
    tmp2 = tmp0 - tmp1
    tl.store(out_ptr0 + x2, tmp2, xmask)


@triton.jit
def triton_poi_fused_sub_4(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 8
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x2 = xindex
    x0 = xindex % 4
    tmp0 = tl.load(in_ptr0 + (56 + x2), xmask)
    tmp1 = tl.load(in_ptr0 + (52 + x0), xmask, eviction_policy='evict_last')
    tmp2 = tmp0 - tmp1
    tl.store(out_ptr0 + x2, tmp2, xmask)


@triton.jit
def triton_per_fused_add_pow_stack_sum_5(in_ptr0, out_ptr1, xnumel, rnumel,
    XBLOCK: tl.constexpr):
    RBLOCK: tl.constexpr = 4
    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_ptr0 + (4 + r0), None)
    tmp1 = tmp0 * tmp0
    tmp3 = tmp2 * tmp2
    tmp4 = tmp1 + tmp3
    tmp5 = tl.broadcast_to(tmp4, [XBLOCK, RBLOCK])
    tmp7 = tl.sum(tmp5, 1)[:, None]
    tl.store(out_ptr1 + tl.full([XBLOCK, 1], 0, tl.int32), tmp7, None)


@triton.jit
def triton_per_fused_add_pow_stack_sum_6(in_ptr0, out_ptr1, xnumel, rnumel,
    XBLOCK: tl.constexpr):
    RBLOCK: tl.constexpr = 4
    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 + (16 + r0), None)
    tmp2 = tl.load(in_ptr0 + (20 + r0), None)
    tmp1 = tmp0 * tmp0
    tmp3 = tmp2 * tmp2
    tmp4 = tmp1 + tmp3
    tmp5 = tl.broadcast_to(tmp4, [XBLOCK, RBLOCK])
    tmp7 = tl.sum(tmp5, 1)[:, None]
    tl.store(out_ptr1 + tl.full([XBLOCK, 1], 0, tl.int32), tmp7, None)


@triton.jit
def triton_per_fused_add_pow_stack_sum_7(in_ptr0, out_ptr1, xnumel, rnumel,
    XBLOCK: tl.constexpr):
    RBLOCK: tl.constexpr = 4
    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 + (32 + r0), None)
    tmp2 = tl.load(in_ptr0 + (36 + r0), None)
    tmp1 = tmp0 * tmp0
    tmp3 = tmp2 * tmp2
    tmp4 = tmp1 + tmp3
    tmp5 = tl.broadcast_to(tmp4, [XBLOCK, RBLOCK])
    tmp7 = tl.sum(tmp5, 1)[:, None]
    tl.store(out_ptr1 + tl.full([XBLOCK, 1], 0, tl.int32), tmp7, None)


@triton.jit
def triton_per_fused_add_pow_stack_sum_8(in_ptr0, out_ptr1, xnumel, rnumel,
    XBLOCK: tl.constexpr):
    RBLOCK: tl.constexpr = 4
    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 + (48 + r0), None)
    tmp2 = tl.load(in_ptr0 + (52 + r0), None)
    tmp1 = tmp0 * tmp0
    tmp3 = tmp2 * tmp2
    tmp4 = tmp1 + tmp3
    tmp5 = tl.broadcast_to(tmp4, [XBLOCK, RBLOCK])
    tmp7 = tl.sum(tmp5, 1)[:, None]
    tl.store(out_ptr1 + tl.full([XBLOCK, 1], 0, tl.int32), tmp7, None)


@triton.jit
def triton_per_fused_add_mean_mul_stack_9(in_out_ptr0, in_ptr0, in_ptr1,
    in_ptr2, in_ptr3, in_ptr4, xnumel, rnumel, XBLOCK: tl.constexpr):
    RBLOCK: tl.constexpr = 4
    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
    tmp5 = tl.load(in_ptr0 + 0)
    tmp6 = tl.broadcast_to(tmp5, [XBLOCK, RBLOCK])
    tmp16 = tl.load(in_ptr1 + 0)
    tmp17 = tl.broadcast_to(tmp16, [XBLOCK, RBLOCK])
    tmp26 = tl.load(in_ptr2 + 0)
    tmp27 = tl.broadcast_to(tmp26, [XBLOCK, RBLOCK])
    tmp35 = tl.load(in_ptr3 + 0)
    tmp36 = tl.broadcast_to(tmp35, [XBLOCK, RBLOCK])
    tmp44 = tl.load(in_ptr4 + r0, None)
    tmp0 = r0
    tl.full([1, 1], 0, tl.int64)
    tmp3 = tl.full([1, 1], 1, tl.int64)
    tmp4 = tmp0 < tmp3
    tmp7 = 1.0
    tmp8 = tmp6 + tmp7
    tmp9 = tl_math.log(tmp8)
    tmp10 = tl.full(tmp9.shape, 0.0, tmp9.dtype)
    tmp11 = tl.where(tmp4, tmp9, tmp10)
    tmp12 = tmp0 >= tmp3
    tmp13 = tl.full([1, 1], 2, tl.int64)
    tmp14 = tmp0 < tmp13
    tmp15 = tmp12 & tmp14
    tmp18 = tmp17 + tmp7
    tmp19 = tl_math.log(tmp18)
    tmp20 = tl.full(tmp19.shape, 0.0, tmp19.dtype)
    tmp21 = tl.where(tmp15, tmp19, tmp20)
    tmp22 = tmp0 >= tmp13
    tmp23 = tl.full([1, 1], 3, tl.int64)
    tmp24 = tmp0 < tmp23
    tmp25 = tmp22 & tmp24
    tmp28 = tmp27 + tmp7
    tmp29 = tl_math.log(tmp28)
    tmp30 = tl.full(tmp29.shape, 0.0, tmp29.dtype)
    tmp31 = tl.where(tmp25, tmp29, tmp30)
    tmp32 = tmp0 >= tmp23
    tl.full([1, 1], 4, tl.int64)
    tmp37 = tmp36 + tmp7
    tmp38 = tl_math.log(tmp37)
    tmp39 = tl.full(tmp38.shape, 0.0, tmp38.dtype)
    tmp40 = tl.where(tmp32, tmp38, tmp39)
    tmp41 = tl.where(tmp25, tmp31, tmp40)
    tmp42 = tl.where(tmp15, tmp21, tmp41)
    tmp43 = tl.where(tmp4, tmp11, tmp42)
    tmp45 = tl.broadcast_to(tmp44, [XBLOCK, RBLOCK])
    tmp47 = tl.sum(tmp45, 1)[:, None]
    tmp48 = 4.0
    tmp49 = tmp47 / tmp48
    tmp50 = 0.05
    tmp51 = tmp49 * tmp50
    tmp52 = tmp43 + tmp51
    tmp53 = tl.broadcast_to(tmp52, [XBLOCK, RBLOCK])
    tmp55 = tl.sum(tmp53, 1)[:, None]
    tmp56 = tmp55 / tmp48
    tl.debug_barrier()
    tl.store(in_out_ptr0 + tl.full([XBLOCK, 1], 0, tl.int32), tmp56, None)


def call(args):
    arg0_1, = args
    args.clear()
    assert_size_stride(arg0_1, (4, 4, 4), (16, 4, 1))
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        buf0 = empty_strided_cuda((2, 4), (4, 1), torch.float32)
        get_raw_stream(0)
        triton_poi_fused_sub_0[grid(8)](arg0_1, buf0, 8, XBLOCK=8,
            num_warps=1, num_stages=1)
        buf1 = empty_strided_cuda((1, 2), (2, 1), torch.float32)
        extern_kernels.mm(reinterpret_tensor(arg0_1, (1, 4), (4, 1), 0),
            reinterpret_tensor(buf0, (4, 2), (1, 4), 0), out=buf1)
        buf2 = empty_strided_cuda((), (), torch.float32)
        triton_per_fused_exp_sum_1[grid(1)](buf1, buf2, 1, 2, XBLOCK=1,
            num_warps=2, num_stages=1)
        buf3 = buf0
        del buf0
        triton_poi_fused_sub_2[grid(8)](arg0_1, buf3, 8, XBLOCK=8,
            num_warps=1, num_stages=1)
        buf4 = buf1
        del buf1
        extern_kernels.mm(reinterpret_tensor(arg0_1, (1, 4), (4, 1), 16),
            reinterpret_tensor(buf3, (4, 2), (1, 4), 0), out=buf4)
        buf5 = empty_strided_cuda((), (), torch.float32)
        triton_per_fused_exp_sum_1[grid(1)](buf4, buf5, 1, 2, XBLOCK=1,
            num_warps=2, num_stages=1)
        buf6 = buf3
        del buf3
        triton_poi_fused_sub_3[grid(8)](arg0_1, buf6, 8, XBLOCK=8,
            num_warps=1, num_stages=1)
        buf7 = buf4
        del buf4
        extern_kernels.mm(reinterpret_tensor(arg0_1, (1, 4), (4, 1), 32),
            reinterpret_tensor(buf6, (4, 2), (1, 4), 0), out=buf7)
        buf8 = empty_strided_cuda((), (), torch.float32)
        triton_per_fused_exp_sum_1[grid(1)](buf7, buf8, 1, 2, XBLOCK=1,
            num_warps=2, num_stages=1)
        buf9 = buf6
        del buf6
        triton_poi_fused_sub_4[grid(8)](arg0_1, buf9, 8, XBLOCK=8,
            num_warps=1, num_stages=1)
        buf10 = buf7
        del buf7
        extern_kernels.mm(reinterpret_tensor(arg0_1, (1, 4), (4, 1), 48),
            reinterpret_tensor(buf9, (4, 2), (1, 4), 0), out=buf10)
        del buf9
        buf11 = empty_strided_cuda((), (), torch.float32)
        triton_per_fused_exp_sum_1[grid(1)](buf10, buf11, 1, 2, XBLOCK=1,
            num_warps=2, num_stages=1)
        del buf10
        buf21 = empty_strided_cuda((4,), (1,), torch.float32)
        buf17 = reinterpret_tensor(buf21, (1,), (1,), 0)
        triton_per_fused_add_pow_stack_sum_5[grid(1)](arg0_1, buf17, 1, 4,
            XBLOCK=1, num_warps=2, num_stages=1)
        buf18 = reinterpret_tensor(buf21, (1,), (1,), 1)
        triton_per_fused_add_pow_stack_sum_6[grid(1)](arg0_1, buf18, 1, 4,
            XBLOCK=1, num_warps=2, num_stages=1)
        buf19 = reinterpret_tensor(buf21, (1,), (1,), 2)
        triton_per_fused_add_pow_stack_sum_7[grid(1)](arg0_1, buf19, 1, 4,
            XBLOCK=1, num_warps=2, num_stages=1)
        buf20 = reinterpret_tensor(buf21, (1,), (1,), 3)
        triton_per_fused_add_pow_stack_sum_8[grid(1)](arg0_1, buf20, 1, 4,
            XBLOCK=1, num_warps=2, num_stages=1)
        del arg0_1
        buf22 = empty_strided_cuda((), (), torch.float32)
        buf23 = buf22
        del buf22
        buf24 = buf23
        del buf23
        triton_per_fused_add_mean_mul_stack_9[grid(1)](buf24, buf2, buf5,
            buf8, buf11, buf21, 1, 4, XBLOCK=1, num_warps=2, num_stages=1)
        del buf11
        del buf17
        del buf18
        del buf19
        del buf2
        del buf20
        del buf21
        del buf5
        del buf8
    return buf24,


class NPairLossNew(torch.nn.Module):

    def __init__(self, l2=0.05):
        """
        Basic N-Pair Loss as proposed in 'Improved Deep Metric Learning with Multi-class N-pair Loss Objective'
        Args:
            l2: float, weighting parameter for weight penality due to embeddings not being normalized.
        Returns:
            Nothing!
        """
        super(NPairLossNew, self).__init__()
        self.l2 = l2

    def npair_distance(self, anchor, positive, negatives):
        """
        Compute basic N-Pair loss.
        Args:
            anchor, positive, negative: torch.Tensor(), resp. embeddings for anchor, positive and negative samples.
        Returns:
            n-pair loss (torch.Tensor())
        """
        return torch.log(1 + torch.sum(torch.exp(anchor.reshape(1, -1).mm((
            negatives - positive).transpose(0, 1)))))

    def weightsum(self, anchor, positive):
        """
        Compute weight penalty.
        NOTE: Only need to penalize anchor and positive since the negatives are created based on these.
        Args:
            anchor, positive: torch.Tensor(), resp. embeddings for anchor and positive samples.
        Returns:
            torch.Tensor(), Weight penalty
        """
        return torch.sum(anchor ** 2 + positive ** 2)

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


Measured latency_ms: 0.18987