# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/05_losses.ipynb (unless otherwise specified).

__all__ = ['WeightedSoftmaxCrossEntropy']

# Cell
import torch
import torch.nn.functional as F

# Cell
class WeightedSoftmaxCrossEntropy(torch.nn.Module):
    "Weighted Softmax Cross Entropy loss functions"
    def __init__(self, *args, axis=-1, reduction = 'mean'):
        super().__init__()
        self.reduction = reduction
        self.axis = axis

    def forward(self, inputs, targets, weights):

        # Weighted soft-max cross-entropy loss
        loss = F.cross_entropy(inputs, targets, reduction='none')
        loss = loss * weights
        if  self.reduction == 'mean':
            return loss.mean()

        elif self.reduction == 'sum':
            return loss.sum()

        else:
            return loss

    def decodes(self, x): return x.argmax(dim=self.axis)
    def activation(self, x): return F.softmax(x, dim=self.axis)