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 numpy as np
import torch.nn as 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_per_fused_add_div_mul_sum_0(in_out_ptr0, in_ptr0, in_ptr1,
    xnumel, rnumel, XBLOCK: tl.constexpr):
    xnumel = 16
    RBLOCK: tl.constexpr = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:, None]
    xmask = xindex < xnumel
    rindex = tl.arange(0, RBLOCK)[None, :]
    tl.full([XBLOCK, RBLOCK], True, tl.int1)
    r1 = rindex
    x0 = xindex
    tmp0 = tl.load(in_ptr0 + (r1 + 16 * x0), xmask, other=0.0)
    tmp1 = tl.load(in_ptr1 + (r1 + 16 * x0), xmask, other=0.0)
    tmp2 = tmp0 * tmp1
    tmp3 = tl.broadcast_to(tmp2, [XBLOCK, RBLOCK])
    tmp5 = tl.where(xmask, tmp3, 0)
    tmp6 = tl.sum(tmp5, 1)[:, None]
    tmp7 = tmp0 + tmp1
    tmp8 = tl.broadcast_to(tmp7, [XBLOCK, RBLOCK])
    tmp10 = tl.where(xmask, tmp8, 0)
    tmp11 = tl.sum(tmp10, 1)[:, None]
    tmp12 = 2.0
    tmp13 = tmp6 * tmp12
    tmp14 = 0.0
    tmp15 = tmp6 * tmp14
    tmp16 = tmp11 + tmp15
    tmp17 = tmp16 + tmp14
    tmp18 = tmp13 / tmp17
    tl.debug_barrier()
    tl.store(in_out_ptr0 + x0, tmp18, xmask)


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), torch.float32)
        buf2 = buf0
        del buf0
        get_raw_stream(0)
        triton_per_fused_add_div_mul_sum_0[grid(16)](buf2, arg1_1, arg0_1, 
            16, 16, XBLOCK=1, num_warps=2, num_stages=1)
        del arg0_1
        del arg1_1
    return buf2,


class IoU(nn.Module):

    def __init__(self, mode='iou', axis=1, eps=0.0):
        """ Return a matrix of [batch * num_classes]. 
            Note: In order to separate from iou=0, function WILL return NaN if both 
            y_true and y_pred are 0. Need further treatment to remove nan in either 
            loss function or matrix.
        """
        super(IoU, self).__init__()
        assert mode in ['iou', 'dice']
        self.factor = {'iou': -1.0, 'dice': 0.0}[mode]
        self.eps = eps
        self.axis = axis

    def forward(self, y_pred, y_true):
        assert y_pred.shape == y_true.shape
        sum_axis = list(range(1, self.axis)) + list(range(self.axis + 1,
            y_pred.ndim))
        prod = (y_true * y_pred).sum(sum_axis)
        plus = (y_true + y_pred).sum(sum_axis)
        iou = (2 + self.factor) * prod / (plus + self.factor * prod + self.eps)
        return iou


class SoftDiceLossNew(IoU):

    def __init__(self, weight=None, ignore_index=[], reduction='mean', mode
        ='dice', axis=1, eps=0.0):
        super(SoftDiceLossNew, self).__init__(mode, axis, eps)
        self.ignore_index = ignore_index
        self.register_buffer('weight', weight)
        self.reduction = {'none': lambda x: x, 'mean': torch.mean, 'sum':
            torch.sum}[reduction]

    def _apply_weight(self, x):
        """ Apply class_weights to calculate loss, ignore nan. """
        if self.weight is None:
            weight = torch.ones(x.shape[-1], device=x.device)
        else:
            weight = self.weight
        idx = np.ones(x.shape[-1], dtype=bool)
        idx[self.ignore_index] = False
        x, weight = x[:, idx], weight[idx]
        weight = ~torch.isnan(x) * weight
        return x * weight / weight.sum(-1, keepdim=True)

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