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 libdevice, math as tl_math
from typing import Callable
from functools import partial
from torch import nn
import torch.distributed
from torch.nn.modules.loss import *
from torch.nn.modules import *
from torch.optim import *
from torch.optim.lr_scheduler import *
import torch.backends
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_binary_cross_entropy_with_logits_div_eq_mul_rsub_sigmoid_sub_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)
    tmp5 = tl.load(in_ptr1 + r0, None)
    tmp1 = tl.sigmoid(tmp0)
    tmp2 = tl.broadcast_to(tmp1, [RBLOCK])
    tmp4 = triton_helpers.promote_to_tensor(tl.sum(tmp2, 0))
    tmp6 = tmp5 * tmp1
    tmp7 = tl.broadcast_to(tmp6, [RBLOCK])
    tmp9 = triton_helpers.promote_to_tensor(tl.sum(tmp7, 0))
    tmp10 = 1.0
    tmp11 = tmp10 - tmp5
    tmp12 = tmp11 * tmp0
    tmp13 = 0.0
    tmp14 = triton_helpers.minimum(tmp13, tmp0)
    tmp15 = tl_math.abs(tmp0)
    tmp16 = -tmp15
    tmp17 = tl_math.exp(tmp16)
    tmp18 = libdevice.log1p(tmp17)
    tmp19 = tmp14 - tmp18
    tmp20 = tmp12 - tmp19
    tmp21 = tl.broadcast_to(tmp20, [RBLOCK])
    tmp23 = triton_helpers.promote_to_tensor(tl.sum(tmp21, 0))
    tmp24 = tl.broadcast_to(tmp5, [RBLOCK])
    tmp26 = triton_helpers.promote_to_tensor(tl.sum(tmp24, 0))
    tmp27 = tmp26 + tmp4
    tmp28 = tmp27 == tmp13
    tmp29 = tmp28.to(tl.float32)
    tmp30 = 1e-07
    tmp31 = tmp29 * tmp30
    tmp32 = tmp9 + tmp31
    tmp33 = tmp27 - tmp9
    tmp34 = tmp33 + tmp30
    tmp35 = tmp32 / tmp34
    tmp36 = tmp10 - tmp35
    tmp37 = 256.0
    tmp38 = tmp23 / tmp37
    tmp39 = tmp36 + tmp38
    tl.debug_barrier()
    tl.store(in_out_ptr0 + tl.full([1], 0, tl.int32), tmp39, 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)
        buf4 = buf0
        del buf0
        get_raw_stream(0)
        triton_per_fused_add_binary_cross_entropy_with_logits_div_eq_mul_rsub_sigmoid_sub_sum_0[
            grid(1)](buf4, arg0_1, arg1_1, 1, 256, num_warps=2, num_stages=1)
        del arg0_1
        del arg1_1
    return buf4,


def get_activation_fn(activation: 'str'=None):
    """Returns the activation function from ``torch.nn`` by its name."""
    if activation is None or activation.lower() == 'none':

        def activation_fn(x):
            return x
    else:
        activation_fn = torch.nn.__dict__[activation]()
    return activation_fn


def wrap_metric_fn_with_activation(metric_fn: 'Callable', activation: 'str'
    =None):
    """Wraps model outputs for ``metric_fn` with specified ``activation``.

    Args:
        metric_fn: metric function to compute
        activation: activation name to use

    Returns:
        wrapped metric function with wrapped model outputs

    .. note::
        Works only with ``metric_fn`` like
        ``metric_fn(outputs, targets, *args, **kwargs)``.
    """
    activation_fn = get_activation_fn(activation)

    def wrapped_metric_fn(outputs: 'torch.Tensor', targets: 'torch.Tensor',
        *args, **kwargs):
        outputs = activation_fn(outputs)
        output = metric_fn(outputs, targets, *args, **kwargs)
        return output
    return wrapped_metric_fn


def iou(outputs: 'torch.Tensor', targets: 'torch.Tensor', eps: 'float'=
    1e-07, threshold: 'float'=None) ->torch.Tensor:
    """Computes the dice score.

    Args:
        outputs: A list of predicted elements
        targets:  A list of elements that are to be predicted
        eps: epsilon to avoid zero division
        threshold: threshold for outputs binarization

    Returns:
        IoU (Jaccard) score

    Examples:
        >>> iou(
        >>>     outputs=torch.tensor([
        >>>         [1, 0, 0],
        >>>         [0, 1, 0],
        >>>         [0, 0, 1],
        >>>     ]),
        >>>     targets=torch.tensor([
        >>>         [1, 0, 0],
        >>>         [0, 1, 0],
        >>>         [0, 0, 1],
        >>>     ]),
        >>>     threshold=0.5,
        >>> )
        tensor(1.0)
        >>> iou(
        >>>     outputs=torch.tensor([
        >>>         [1, 0, 0],
        >>>         [0, 1, 0],
        >>>         [0, 0, 1],
        >>>     ]),
        >>>     targets=torch.tensor([
        >>>         [1, 0, 0],
        >>>         [0, 1, 0],
        >>>         [0, 0, 0],
        >>>     ]),
        >>>     threshold=0.5,
        >>> )
        tensor(0.6667)
    """
    if threshold is not None:
        outputs = (outputs > threshold).float()
    intersection = torch.sum(targets * outputs)
    union = torch.sum(targets) + torch.sum(outputs)
    iou_score = (intersection + eps * (union == 0)) / (union - intersection +
        eps)
    return iou_score


class IoULoss(nn.Module):
    """The intersection over union (Jaccard) loss.

    @TODO: Docs. Contribution is welcome.
    """

    def __init__(self, eps: 'float'=1e-07, threshold: 'float'=None,
        activation: 'str'='Sigmoid'):
        """
        Args:
            eps: epsilon to avoid zero division
            threshold: threshold for outputs binarization
            activation: An torch.nn activation applied to the outputs.
                Must be one of ``'none'``, ``'Sigmoid'``, ``'Softmax'``
        """
        super().__init__()
        metric_fn = wrap_metric_fn_with_activation(metric_fn=iou,
            activation=activation)
        self.loss_fn = partial(metric_fn, eps=eps, threshold=threshold)

    def forward(self, outputs, targets):
        """@TODO: Docs. Contribution is welcome."""
        iou = self.loss_fn(outputs, targets)
        return 1 - iou


class BCEIoULossNew(nn.Module):
    """The Intersection over union (Jaccard) with BCE loss.

    @TODO: Docs. Contribution is welcome.
    """

    def __init__(self, eps: 'float'=1e-07, threshold: 'float'=None,
        activation: 'str'='Sigmoid', reduction: 'str'='mean'):
        """
        Args:
            eps: epsilon to avoid zero division
            threshold: threshold for outputs binarization
            activation: An torch.nn activation applied to the outputs.
                Must be one of ``'none'``, ``'Sigmoid'``, ``'Softmax'``
            reduction: Specifies the reduction to apply
                to the output of BCE
        """
        super().__init__()
        self.bce_loss = nn.BCEWithLogitsLoss(reduction=reduction)
        self.iou_loss = IoULoss(eps, threshold, activation)

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