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
import functools
import numpy as np
import torch.nn.functional as F
import torch.nn as nn
import torch._C
import torch.serialization
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__log_softmax__softmax_0(in_ptr0, out_ptr0, out_ptr1,
    xnumel, XBLOCK: tl.constexpr):
    xnumel = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x2 = xindex
    x1 = xindex // 4
    tmp0 = tl.load(in_ptr0 + x2, xmask)
    tmp1 = tl.load(in_ptr0 + 4 * x1, xmask, eviction_policy='evict_last')
    tmp2 = tl.load(in_ptr0 + (1 + 4 * x1), xmask, eviction_policy='evict_last')
    tmp4 = tl.load(in_ptr0 + (2 + 4 * x1), xmask, eviction_policy='evict_last')
    tmp6 = tl.load(in_ptr0 + (3 + 4 * x1), xmask, eviction_policy='evict_last')
    tmp3 = triton_helpers.maximum(tmp1, tmp2)
    tmp5 = triton_helpers.maximum(tmp3, tmp4)
    tmp7 = triton_helpers.maximum(tmp5, tmp6)
    tmp8 = tmp0 - tmp7
    tmp9 = tl_math.exp(tmp8)
    tl.store(out_ptr0 + x2, tmp8, xmask)
    tl.store(out_ptr1 + x2, tmp9, xmask)


@triton.jit
def triton_poi_fused__softmax_1(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr
    ):
    xnumel = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x2 = xindex
    x1 = xindex // 4
    tmp0 = tl.load(in_ptr0 + x2, xmask)
    tmp1 = tl.load(in_ptr0 + 4 * x1, xmask, eviction_policy='evict_last')
    tmp2 = tl.load(in_ptr0 + (1 + 4 * x1), xmask, eviction_policy='evict_last')
    tmp4 = tl.load(in_ptr0 + (2 + 4 * x1), xmask, eviction_policy='evict_last')
    tmp6 = tl.load(in_ptr0 + (3 + 4 * x1), xmask, eviction_policy='evict_last')
    tmp3 = tmp1 + tmp2
    tmp5 = tmp3 + tmp4
    tmp7 = tmp5 + tmp6
    tmp8 = tmp0 / tmp7
    tl.store(out_ptr0 + x2, tmp8, xmask)


@triton.jit
def triton_per_fused__log_softmax__to_copy_add_div_mean_mul_ne_neg_pow_rsub_sum_view_2(
    in_out_ptr0, in_ptr0, in_ptr1, in_ptr2, 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 + 4 * r0, None, eviction_policy='evict_last')
    tmp2 = tl.load(in_ptr0 + (1 + 4 * r0), None, eviction_policy='evict_last')
    tmp5 = tl.load(in_ptr0 + (2 + 4 * r0), None, eviction_policy='evict_last')
    tmp8 = tl.load(in_ptr0 + (3 + 4 * r0), None, eviction_policy='evict_last')
    tmp13 = tl.load(in_ptr1 + 4 * r0, None, eviction_policy='evict_last')
    tmp16 = tl.load(in_ptr1 + (1 + 4 * r0), None, eviction_policy='evict_last')
    tmp20 = tl.load(in_ptr1 + (2 + 4 * r0), None, eviction_policy='evict_last')
    tmp24 = tl.load(in_ptr1 + (3 + 4 * r0), None, eviction_policy='evict_last')
    tmp31 = tl.load(in_ptr2 + 4 * r0, None, eviction_policy='evict_last')
    tmp98 = tl.load(in_ptr2 + (1 + 4 * r0), None, eviction_policy='evict_last')
    tmp139 = tl.load(in_ptr2 + (2 + 4 * r0), None, eviction_policy='evict_last'
        )
    tmp180 = tl.load(in_ptr2 + (3 + 4 * r0), None, eviction_policy='evict_last'
        )
    tmp1 = tl_math.exp(tmp0)
    tmp3 = tl_math.exp(tmp2)
    tmp4 = tmp1 + tmp3
    tmp6 = tl_math.exp(tmp5)
    tmp7 = tmp4 + tmp6
    tmp9 = tl_math.exp(tmp8)
    tmp10 = tmp7 + tmp9
    tmp11 = tl_math.log(tmp10)
    tmp12 = tmp0 - tmp11
    tmp14 = tmp12 * tmp13
    tmp15 = tmp2 - tmp11
    tmp17 = tmp15 * tmp16
    tmp18 = tmp14 + tmp17
    tmp19 = tmp5 - tmp11
    tmp21 = tmp19 * tmp20
    tmp22 = tmp18 + tmp21
    tmp23 = tmp8 - tmp11
    tmp25 = tmp23 * tmp24
    tmp26 = tmp22 + tmp25
    tmp27 = -tmp26
    tmp28 = tl.broadcast_to(tmp27, [XBLOCK, RBLOCK])
    tmp30 = tl.sum(tmp28, 1)[:, None]
    tmp32 = tmp13.to(tl.int64)
    tmp33 = tl.full([1, 1], 0, tl.int64)
    tmp34 = triton_helpers.maximum(tmp32, tmp33)
    tmp35 = tl.full([1, 1], 3, tl.int64)
    tmp36 = triton_helpers.minimum(tmp34, tmp35)
    tmp37 = tmp36 == tmp33
    tmp38 = tmp37.to(tl.int64)
    tmp39 = tmp38.to(tl.float32)
    tmp40 = tmp31 * tmp39
    tmp41 = 255.0
    tmp42 = tmp13 != tmp41
    tmp43 = tmp42.to(tl.int64)
    tmp44 = tmp43.to(tl.float32)
    tmp45 = tmp40 * tmp44
    tmp46 = tmp16.to(tl.int64)
    tmp47 = triton_helpers.maximum(tmp46, tmp33)
    tmp48 = triton_helpers.minimum(tmp47, tmp35)
    tmp49 = tmp48 == tmp33
    tmp50 = tmp49.to(tl.int64)
    tmp51 = tmp50.to(tl.float32)
    tmp52 = tmp31 * tmp51
    tmp53 = tmp16 != tmp41
    tmp54 = tmp53.to(tl.int64)
    tmp55 = tmp54.to(tl.float32)
    tmp56 = tmp52 * tmp55
    tmp57 = tmp45 + tmp56
    tmp58 = tmp20.to(tl.int64)
    tmp59 = triton_helpers.maximum(tmp58, tmp33)
    tmp60 = triton_helpers.minimum(tmp59, tmp35)
    tmp61 = tmp60 == tmp33
    tmp62 = tmp61.to(tl.int64)
    tmp63 = tmp62.to(tl.float32)
    tmp64 = tmp31 * tmp63
    tmp65 = tmp20 != tmp41
    tmp66 = tmp65.to(tl.int64)
    tmp67 = tmp66.to(tl.float32)
    tmp68 = tmp64 * tmp67
    tmp69 = tmp57 + tmp68
    tmp70 = tmp24.to(tl.int64)
    tmp71 = triton_helpers.maximum(tmp70, tmp33)
    tmp72 = triton_helpers.minimum(tmp71, tmp35)
    tmp73 = tmp72 == tmp33
    tmp74 = tmp73.to(tl.int64)
    tmp75 = tmp74.to(tl.float32)
    tmp76 = tmp31 * tmp75
    tmp77 = tmp24 != tmp41
    tmp78 = tmp77.to(tl.int64)
    tmp79 = tmp78.to(tl.float32)
    tmp80 = tmp76 * tmp79
    tmp81 = tmp69 + tmp80
    tmp82 = tmp31 * tmp31
    tmp83 = tmp38 * tmp38
    tmp84 = tmp83.to(tl.float32)
    tmp85 = tmp82 + tmp84
    tmp86 = tmp50 * tmp50
    tmp87 = tmp86.to(tl.float32)
    tmp88 = tmp82 + tmp87
    tmp89 = tmp85 + tmp88
    tmp90 = tmp62 * tmp62
    tmp91 = tmp90.to(tl.float32)
    tmp92 = tmp82 + tmp91
    tmp93 = tmp89 + tmp92
    tmp94 = tmp74 * tmp74
    tmp95 = tmp94.to(tl.float32)
    tmp96 = tmp82 + tmp95
    tmp97 = tmp93 + tmp96
    tmp99 = tl.full([1, 1], 1, tl.int64)
    tmp100 = tmp36 == tmp99
    tmp101 = tmp100.to(tl.int64)
    tmp102 = tmp101.to(tl.float32)
    tmp103 = tmp98 * tmp102
    tmp104 = tmp103 * tmp44
    tmp105 = tmp48 == tmp99
    tmp106 = tmp105.to(tl.int64)
    tmp107 = tmp106.to(tl.float32)
    tmp108 = tmp98 * tmp107
    tmp109 = tmp108 * tmp55
    tmp110 = tmp104 + tmp109
    tmp111 = tmp60 == tmp99
    tmp112 = tmp111.to(tl.int64)
    tmp113 = tmp112.to(tl.float32)
    tmp114 = tmp98 * tmp113
    tmp115 = tmp114 * tmp67
    tmp116 = tmp110 + tmp115
    tmp117 = tmp72 == tmp99
    tmp118 = tmp117.to(tl.int64)
    tmp119 = tmp118.to(tl.float32)
    tmp120 = tmp98 * tmp119
    tmp121 = tmp120 * tmp79
    tmp122 = tmp116 + tmp121
    tmp123 = tmp98 * tmp98
    tmp124 = tmp101 * tmp101
    tmp125 = tmp124.to(tl.float32)
    tmp126 = tmp123 + tmp125
    tmp127 = tmp106 * tmp106
    tmp128 = tmp127.to(tl.float32)
    tmp129 = tmp123 + tmp128
    tmp130 = tmp126 + tmp129
    tmp131 = tmp112 * tmp112
    tmp132 = tmp131.to(tl.float32)
    tmp133 = tmp123 + tmp132
    tmp134 = tmp130 + tmp133
    tmp135 = tmp118 * tmp118
    tmp136 = tmp135.to(tl.float32)
    tmp137 = tmp123 + tmp136
    tmp138 = tmp134 + tmp137
    tmp140 = tl.full([1, 1], 2, tl.int64)
    tmp141 = tmp36 == tmp140
    tmp142 = tmp141.to(tl.int64)
    tmp143 = tmp142.to(tl.float32)
    tmp144 = tmp139 * tmp143
    tmp145 = tmp144 * tmp44
    tmp146 = tmp48 == tmp140
    tmp147 = tmp146.to(tl.int64)
    tmp148 = tmp147.to(tl.float32)
    tmp149 = tmp139 * tmp148
    tmp150 = tmp149 * tmp55
    tmp151 = tmp145 + tmp150
    tmp152 = tmp60 == tmp140
    tmp153 = tmp152.to(tl.int64)
    tmp154 = tmp153.to(tl.float32)
    tmp155 = tmp139 * tmp154
    tmp156 = tmp155 * tmp67
    tmp157 = tmp151 + tmp156
    tmp158 = tmp72 == tmp140
    tmp159 = tmp158.to(tl.int64)
    tmp160 = tmp159.to(tl.float32)
    tmp161 = tmp139 * tmp160
    tmp162 = tmp161 * tmp79
    tmp163 = tmp157 + tmp162
    tmp164 = tmp139 * tmp139
    tmp165 = tmp142 * tmp142
    tmp166 = tmp165.to(tl.float32)
    tmp167 = tmp164 + tmp166
    tmp168 = tmp147 * tmp147
    tmp169 = tmp168.to(tl.float32)
    tmp170 = tmp164 + tmp169
    tmp171 = tmp167 + tmp170
    tmp172 = tmp153 * tmp153
    tmp173 = tmp172.to(tl.float32)
    tmp174 = tmp164 + tmp173
    tmp175 = tmp171 + tmp174
    tmp176 = tmp159 * tmp159
    tmp177 = tmp176.to(tl.float32)
    tmp178 = tmp164 + tmp177
    tmp179 = tmp175 + tmp178
    tmp181 = tmp36 == tmp35
    tmp182 = tmp181.to(tl.int64)
    tmp183 = tmp182.to(tl.float32)
    tmp184 = tmp180 * tmp183
    tmp185 = tmp184 * tmp44
    tmp186 = tmp48 == tmp35
    tmp187 = tmp186.to(tl.int64)
    tmp188 = tmp187.to(tl.float32)
    tmp189 = tmp180 * tmp188
    tmp190 = tmp189 * tmp55
    tmp191 = tmp185 + tmp190
    tmp192 = tmp60 == tmp35
    tmp193 = tmp192.to(tl.int64)
    tmp194 = tmp193.to(tl.float32)
    tmp195 = tmp180 * tmp194
    tmp196 = tmp195 * tmp67
    tmp197 = tmp191 + tmp196
    tmp198 = tmp72 == tmp35
    tmp199 = tmp198.to(tl.int64)
    tmp200 = tmp199.to(tl.float32)
    tmp201 = tmp180 * tmp200
    tmp202 = tmp201 * tmp79
    tmp203 = tmp197 + tmp202
    tmp204 = tmp180 * tmp180
    tmp205 = tmp182 * tmp182
    tmp206 = tmp205.to(tl.float32)
    tmp207 = tmp204 + tmp206
    tmp208 = tmp187 * tmp187
    tmp209 = tmp208.to(tl.float32)
    tmp210 = tmp204 + tmp209
    tmp211 = tmp207 + tmp210
    tmp212 = tmp193 * tmp193
    tmp213 = tmp212.to(tl.float32)
    tmp214 = tmp204 + tmp213
    tmp215 = tmp211 + tmp214
    tmp216 = tmp199 * tmp199
    tmp217 = tmp216.to(tl.float32)
    tmp218 = tmp204 + tmp217
    tmp219 = tmp215 + tmp218
    tmp220 = 2.0
    tmp221 = tmp81 * tmp220
    tmp222 = 1.0
    tmp223 = tmp221 + tmp222
    tmp224 = tmp97 + tmp222
    tmp225 = tmp223 / tmp224
    tmp226 = tmp222 - tmp225
    tmp227 = tl.broadcast_to(tmp226, [XBLOCK, RBLOCK])
    tmp229 = tl.sum(tmp227, 1)[:, None]
    tmp230 = tmp122 * tmp220
    tmp231 = tmp230 + tmp222
    tmp232 = tmp138 + tmp222
    tmp233 = tmp231 / tmp232
    tmp234 = tmp222 - tmp233
    tmp235 = tl.broadcast_to(tmp234, [XBLOCK, RBLOCK])
    tmp237 = tl.sum(tmp235, 1)[:, None]
    tmp238 = tmp163 * tmp220
    tmp239 = tmp238 + tmp222
    tmp240 = tmp179 + tmp222
    tmp241 = tmp239 / tmp240
    tmp242 = tmp222 - tmp241
    tmp243 = tl.broadcast_to(tmp242, [XBLOCK, RBLOCK])
    tmp245 = tl.sum(tmp243, 1)[:, None]
    tmp246 = tmp203 * tmp220
    tmp247 = tmp246 + tmp222
    tmp248 = tmp219 + tmp222
    tmp249 = tmp247 / tmp248
    tmp250 = tmp222 - tmp249
    tmp251 = tl.broadcast_to(tmp250, [XBLOCK, RBLOCK])
    tmp253 = tl.sum(tmp251, 1)[:, None]
    tmp254 = 4.0
    tmp255 = tmp30 / tmp254
    tmp256 = tmp255 * tmp222
    tmp257 = 0.5
    tmp258 = tmp256 * tmp257
    tmp259 = tmp229 / tmp254
    tmp260 = 0.0
    tmp261 = tmp259 + tmp260
    tmp262 = tmp237 / tmp254
    tmp263 = tmp261 + tmp262
    tmp264 = tmp245 / tmp254
    tmp265 = tmp263 + tmp264
    tmp266 = tmp253 / tmp254
    tmp267 = tmp265 + tmp266
    tmp268 = 0.25
    tmp269 = tmp267 * tmp268
    tmp270 = tmp269 / tmp222
    tmp271 = tmp270 * tmp222
    tmp272 = tmp271 * tmp257
    tmp273 = tmp258 + tmp272
    tl.debug_barrier()
    tl.store(in_out_ptr0 + tl.full([XBLOCK, 1], 0, tl.int32), tmp273, None)


def call(args):
    arg0_1, arg1_1 = args
    args.clear()
    assert_size_stride(arg0_1, (4, 4), (4, 1))
    assert_size_stride(arg1_1, (4, 4), (4, 1))
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        buf0 = empty_strided_cuda((4, 4), (4, 1), torch.float32)
        buf2 = empty_strided_cuda((4, 4), (4, 1), torch.float32)
        get_raw_stream(0)
        triton_poi_fused__log_softmax__softmax_0[grid(16)](arg0_1, buf0,
            buf2, 16, XBLOCK=16, num_warps=1, num_stages=1)
        del arg0_1
        buf3 = empty_strided_cuda((4, 4), (4, 1), torch.float32)
        triton_poi_fused__softmax_1[grid(16)](buf2, buf3, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        del buf2
        buf1 = empty_strided_cuda((), (), torch.float32)
        buf16 = buf1
        del buf1
        triton_per_fused__log_softmax__to_copy_add_div_mean_mul_ne_neg_pow_rsub_sum_view_2[
            grid(1)](buf16, buf0, arg1_1, buf3, 1, 4, XBLOCK=1, num_warps=2,
            num_stages=1)
        del arg1_1
        del buf0
        del buf3
    return buf16,


def _expand_onehot_labels(labels, label_weights, target_shape, ignore_index):
    """Expand onehot labels to match the size of prediction."""
    bin_labels = labels.new_zeros(target_shape)
    valid_mask = (labels >= 0) & (labels != ignore_index)
    inds = torch.nonzero(valid_mask, as_tuple=True)
    if inds[0].numel() > 0:
        if labels.dim() == 3:
            bin_labels[inds[0], labels[valid_mask], inds[1], inds[2]] = 1
        else:
            bin_labels[inds[0], labels[valid_mask]] = 1
    valid_mask = valid_mask.unsqueeze(1).expand(target_shape).float()
    if label_weights is None:
        bin_label_weights = valid_mask
    else:
        bin_label_weights = label_weights.unsqueeze(1).expand(target_shape)
        bin_label_weights *= valid_mask
    return bin_labels, bin_label_weights


def reduce_loss(loss, reduction):
    """Reduce loss as specified.

    Args:
        loss (Tensor): Elementwise loss tensor.
        reduction (str): Options are "none", "mean" and "sum".

    Return:
        Tensor: Reduced loss tensor.
    """
    reduction_enum = F._Reduction.get_enum(reduction)
    if reduction_enum == 0:
        return loss
    elif reduction_enum == 1:
        return loss.mean()
    elif reduction_enum == 2:
        return loss.sum()


def weight_reduce_loss(loss, weight=None, reduction='mean', avg_factor=None):
    """Apply element-wise weight and reduce loss.

    Args:
        loss (Tensor): Element-wise loss.
        weight (Tensor): Element-wise weights.
        reduction (str): Same as built-in losses of PyTorch.
        avg_factor (float): Avarage factor when computing the mean of losses.

    Returns:
        Tensor: Processed loss values.
    """
    if weight is not None:
        assert weight.dim() == loss.dim()
        if weight.dim() > 1:
            assert weight.size(1) == 1 or weight.size(1) == loss.size(1)
        loss = loss * weight
    if avg_factor is None:
        loss = reduce_loss(loss, reduction)
    elif reduction == 'mean':
        loss = loss.sum() / avg_factor
    elif reduction != 'none':
        raise ValueError('avg_factor can not be used with reduction="sum"')
    return loss


def binary_cross_entropy(pred, label, weight=None, reduction='mean',
    avg_factor=None, class_weight=None, ignore_index=255):
    """Calculate the binary CrossEntropy loss.

    Args:
        pred (torch.Tensor): The prediction with shape (N, 1).
        label (torch.Tensor): The learning label of the prediction.
        weight (torch.Tensor, optional): Sample-wise loss weight.
        reduction (str, optional): The method used to reduce the loss.
            Options are "none", "mean" and "sum".
        avg_factor (int, optional): Average factor that is used to average
            the loss. Defaults to None.
        class_weight (list[float], optional): The weight for each class.
        ignore_index (int | None): The label index to be ignored. Default: 255

    Returns:
        torch.Tensor: The calculated loss
    """
    if pred.dim() != label.dim():
        assert pred.dim() == 2 and label.dim() == 1 or pred.dim(
            ) == 4 and label.dim(
            ) == 3, 'Only pred shape [N, C], label shape [N] or pred shape [N, C, H, W], label shape [N, H, W] are supported'
        label, weight = _expand_onehot_labels(label, weight, pred.shape,
            ignore_index)
    if weight is not None:
        weight = weight.float()
    loss = F.binary_cross_entropy_with_logits(pred, label.float(),
        pos_weight=class_weight, reduction='none')
    loss = weight_reduce_loss(loss, weight, reduction=reduction, avg_factor
        =avg_factor)
    return loss


def cross_entropy(pred, label, weight=None, class_weight=None, reduction=
    'mean', avg_factor=None, ignore_index=-100):
    """The wrapper function for :func:`F.cross_entropy`"""
    loss = F.cross_entropy(pred, label, weight=class_weight, reduction=
        'none', ignore_index=ignore_index)
    if weight is not None:
        weight = weight.float()
    loss = weight_reduce_loss(loss, weight=weight, reduction=reduction,
        avg_factor=avg_factor)
    return loss


def weighted_loss(loss_func):
    """Create a weighted version of a given loss function.

    To use this decorator, the loss function must have the signature like
    `loss_func(pred, target, **kwargs)`. The function only needs to compute
    element-wise loss without any reduction. This decorator will add weight
    and reduction arguments to the function. The decorated function will have
    the signature like `loss_func(pred, target, weight=None, reduction='mean',
    avg_factor=None, **kwargs)`.

    :Example:

    >>> import torch
    >>> @weighted_loss
    >>> def l1_loss(pred, target):
    >>>     return (pred - target).abs()

    >>> pred = torch.Tensor([0, 2, 3])
    >>> target = torch.Tensor([1, 1, 1])
    >>> weight = torch.Tensor([1, 0, 1])

    >>> l1_loss(pred, target)
    tensor(1.3333)
    >>> l1_loss(pred, target, weight)
    tensor(1.)
    >>> l1_loss(pred, target, reduction='none')
    tensor([1., 1., 2.])
    >>> l1_loss(pred, target, weight, avg_factor=2)
    tensor(1.5000)
    """

    @functools.wraps(loss_func)
    def wrapper(pred, target, weight=None, reduction='mean', avg_factor=
        None, **kwargs):
        loss = loss_func(pred, target, **kwargs)
        loss = weight_reduce_loss(loss, weight, reduction, avg_factor)
        return loss
    return wrapper


@weighted_loss
def binary_dice_loss(pred, target, valid_mask, smooth=1, exponent=2, **kwards):
    assert pred.shape[0] == target.shape[0]
    pred = pred.reshape(pred.shape[0], -1)
    target = target.reshape(target.shape[0], -1)
    valid_mask = valid_mask.reshape(valid_mask.shape[0], -1)
    num = torch.sum(torch.mul(pred, target) * valid_mask, dim=1) * 2 + smooth
    den = torch.sum(pred.pow(exponent) + target.pow(exponent), dim=1) + smooth
    return 1 - num / den


@weighted_loss
def dice_loss(pred, target, valid_mask, smooth=1, exponent=2, class_weight=
    None, ignore_index=255):
    assert pred.shape[0] == target.shape[0]
    total_loss = 0
    num_classes = pred.shape[1]
    for i in range(num_classes):
        if i != ignore_index:
            dice_loss = binary_dice_loss(pred[:, i], target[..., i],
                valid_mask=valid_mask, smooth=smooth, exponent=exponent)
            if class_weight is not None:
                dice_loss *= class_weight[i]
            total_loss += dice_loss
    return total_loss / num_classes


def get_class_weight(class_weight):
    """Get class weight for loss function.

    Args:
        class_weight (list[float] | str | None): If class_weight is a str,
            take it as a file name and read from it.
    """
    if isinstance(class_weight, str):
        if class_weight.endswith('.npy'):
            class_weight = np.load(class_weight)
        else:
            class_weight = mmcv.load(class_weight)
    return class_weight


def mask_cross_entropy(pred, target, label, reduction='mean', avg_factor=
    None, class_weight=None, ignore_index=None):
    """Calculate the CrossEntropy loss for masks.

    Args:
        pred (torch.Tensor): The prediction with shape (N, C), C is the number
            of classes.
        target (torch.Tensor): The learning label of the prediction.
        label (torch.Tensor): ``label`` indicates the class label of the mask'
            corresponding object. This will be used to select the mask in the
            of the class which the object belongs to when the mask prediction
            if not class-agnostic.
        reduction (str, optional): The method used to reduce the loss.
            Options are "none", "mean" and "sum".
        avg_factor (int, optional): Average factor that is used to average
            the loss. Defaults to None.
        class_weight (list[float], optional): The weight for each class.
        ignore_index (None): Placeholder, to be consistent with other loss.
            Default: None.

    Returns:
        torch.Tensor: The calculated loss
    """
    assert ignore_index is None, 'BCE loss does not support ignore_index'
    assert reduction == 'mean' and avg_factor is None
    num_rois = pred.size()[0]
    inds = torch.arange(0, num_rois, dtype=torch.long, device=pred.device)
    pred_slice = pred[inds, label].squeeze(1)
    return F.binary_cross_entropy_with_logits(pred_slice, target, weight=
        class_weight, reduction='mean')[None]


class BCEDiceLossNew(nn.Module):
    """CrossEntropyLoss.

    Args:
        use_sigmoid (bool, optional): Whether the prediction uses sigmoid
            of softmax. Defaults to False.
        use_mask (bool, optional): Whether to use mask cross entropy loss.
            Defaults to False.
        reduction (str, optional): . Defaults to 'mean'.
            Options are "none", "mean" and "sum".
        class_weight (list[float] | str, optional): Weight of each class. If in
            str format, read them from a file. Defaults to None.
        loss_weight (float, optional): Weight of the loss. Defaults to 1.0.
    """

    def __init__(self, use_sigmoid=False, use_mask=False, smooth=1,
        exponent=2, reduction='mean', class_weight=None, loss_weight=1.0,
        ignore_index=255, bce_weight=0.5, **kwards):
        super(BCEDiceLossNew, self).__init__()
        assert use_sigmoid is False or use_mask is False
        self.use_sigmoid = use_sigmoid
        self.use_mask = use_mask
        self.smooth = smooth
        self.exponent = exponent
        self.reduction = reduction
        self.class_weight = get_class_weight(class_weight)
        self.loss_weight = loss_weight
        self.ignore_index = ignore_index
        self.bce_weight = bce_weight
        if self.use_sigmoid:
            self.cls_criterion = binary_cross_entropy
        elif self.use_mask:
            self.cls_criterion = mask_cross_entropy
        else:
            self.cls_criterion = cross_entropy

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