text: # 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
  from fastai.torch_core import TensorBase
  
  # 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 _contiguous(self,x): return TensorBase(x.contiguous())
      def forward(self, inp, targ, weights):
  
          inp, targ  = map(self._contiguous, (inp, targ))
          # Weighted soft-max cross-entropy loss
          loss = F.cross_entropy(inp, targ, 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)