text: from torch import nn
  from torch.autograd import Function
  import torch
  
  
  # Other possible implementations:
  #  https://github.com/KrisKorrel/sparsemax-pytorch/blob/master/sparsemax.py
  #  https://github.com/msobroza/SparsemaxPytorch/blob/master/mnist/sparsemax.py
  # https://github.com/vene/sparse-structured-attention/blob/master/pytorch/torchsparseattn/sparsemax.py
  
  
  
  # credits to Yandex https://github.com/Qwicen/node/blob/master/lib/nn_utils.py
  def _make_ix_like(input, dim=0):
      d = input.size(dim)
      rho = torch.arange(1, d + 1, device=input.device, dtype=input.dtype)
      view = [1] * input.dim()
      view[0] = -1
      return rho.view(view).transpose(0, dim)
  
  
  class SparsemaxFunction(Function):
      """
      An implementation of sparsemax (Martins & Astudillo, 2016). See
      :cite:`DBLP:journals/corr/MartinsA16` for detailed description.
      By Ben Peters and Vlad Niculae
      """
  
      @staticmethod
      def forward(ctx, input, dim=-1):
          """sparsemax: normalizing sparse transform (a la softmax)
          Parameters:
              input (Tensor): any shape
              dim: dimension along which to apply sparsemax
          Returns:
              output (Tensor): same shape as input
          """
          ctx.dim = dim
          max_val, _ = input.max(dim=dim, keepdim=True)
          input -= max_val  # same numerical stability trick as for softmax
          tau, supp_size = SparsemaxFunction._threshold_and_support(input, dim=dim)
          output = torch.clamp(input - tau, min=0)
          ctx.save_for_backward(supp_size, output)
          return output
  
      @staticmethod
      def backward(ctx, grad_output):
          supp_size, output = ctx.saved_tensors
          dim = ctx.dim
          grad_input = grad_output.clone()
          grad_input[output == 0] = 0
  
          v_hat = grad_input.sum(dim=dim) / supp_size.to(output.dtype).squeeze()
          v_hat = v_hat.unsqueeze(dim)
          grad_input = torch.where(output != 0, grad_input - v_hat, grad_input)
          return grad_input, None
  
  
      @staticmethod
      def _threshold_and_support(input, dim=-1):
          """Sparsemax building block: compute the threshold
          Args:
              input: any dimension
              dim: dimension along which to apply the sparsemax
          Returns:
              the threshold value
          """
  
          input_srt, _ = torch.sort(input, descending=True, dim=dim)
          input_cumsum = input_srt.cumsum(dim) - 1
          rhos = _make_ix_like(input, dim)
          support = rhos * input_srt > input_cumsum
  
          support_size = support.sum(dim=dim).unsqueeze(dim)
          tau = input_cumsum.gather(dim, support_size - 1)
          tau /= support_size.to(input.dtype)
          return tau, support_size
  
  
  #sparsemax = lambda input, dim=-1: SparsemaxFunction.apply(input, dim)
  
  
  sparsemax = SparsemaxFunction.apply
  
  
  class Sparsemax(nn.Module):
  
      def __init__(self, dim=-1):
          self.dim = dim
          super(Sparsemax, self).__init__()
  
      def forward(self, input):
          return sparsemax(input, self.dim)
  
  
  class Entmax15Function(Function):
      """
      An implementation of exact Entmax with alpha=1.5 (B. Peters, V. Niculae, A. Martins). See
      :cite:`https://arxiv.org/abs/1905.05702 for detailed description.
      Source: https://github.com/deep-spin/entmax
      """
  
      @staticmethod
      def forward(ctx, input, dim=-1):
          ctx.dim = dim
  
          max_val, _ = input.max(dim=dim, keepdim=True)
          input = input - max_val  # same numerical stability trick as for softmax
          input = input / 2  # divide by 2 to solve actual Entmax
  
          tau_star, _ = Entmax15Function._threshold_and_support(input, dim)
          output = torch.clamp(input - tau_star, min=0) ** 2
          ctx.save_for_backward(output)
          return output
  
      @staticmethod
      def backward(ctx, grad_output):
          Y, = ctx.saved_tensors
          gppr = Y.sqrt()  # = 1 / g'' (Y)
          dX = grad_output * gppr
          q = dX.sum(ctx.dim) / gppr.sum(ctx.dim)
          q = q.unsqueeze(ctx.dim)
          dX -= q * gppr
          return dX, None
  
      @staticmethod
      def _threshold_and_support(input, dim=-1):
          Xsrt, _ = torch.sort(input, descending=True, dim=dim)
  
          rho = _make_ix_like(input, dim)
          mean = Xsrt.cumsum(dim) / rho
          mean_sq = (Xsrt ** 2).cumsum(dim) / rho
          ss = rho * (mean_sq - mean ** 2)
          delta = (1 - ss) / rho
  
          # NOTE this is not exactly the same as in reference algo
          # Fortunately it seems the clamped values never wrongly
          # get selected by tau <= sorted_z. Prove this!
          delta_nz = torch.clamp(delta, 0)
          tau = mean - torch.sqrt(delta_nz)
  
          support_size = (tau <= Xsrt).sum(dim).unsqueeze(dim)
          tau_star = tau.gather(dim, support_size - 1)
          return tau_star, support_size
  
  
  class Entmoid15(Function):
      """ A highly optimized equivalent of labda x: Entmax15([x, 0]) """
  
      @staticmethod
      def forward(ctx, input):
          output = Entmoid15._forward(input)
          ctx.save_for_backward(output)
          return output
  
      @staticmethod
      def _forward(input):
          input, is_pos = abs(input), input >= 0
          tau = (input + torch.sqrt(F.relu(8 - input ** 2))) / 2
          tau.masked_fill_(tau <= input, 2.0)
          y_neg = 0.25 * F.relu(tau - input, inplace=True) ** 2
          return torch.where(is_pos, 1 - y_neg, y_neg)
  
      @staticmethod
      def backward(ctx, grad_output):
          return Entmoid15._backward(ctx.saved_tensors[0], grad_output)
  
      @staticmethod
      def _backward(output, grad_output):
          gppr0, gppr1 = output.sqrt(), (1 - output).sqrt()
          grad_input = grad_output * gppr0
          q = grad_input / (gppr0 + gppr1)
          grad_input -= q * gppr0
          return grad_input
  
  
  
  entmax15 = Entmax15Function.apply
  entmoid15 = Entmoid15.apply
  
  class Entmax15(nn.Module):
  
      def __init__(self, dim=-1):
          self.dim = dim
          super(Entmax15, self).__init__()
  
      def forward(self, input):
          return entmax15(input, self.dim)
  
  
  
  
  # Credits were lost...
  # def _make_ix_like(input, dim=0):
  #     d = input.size(dim)
  #     rho = torch.arange(1, d + 1, device=input.device, dtype=input.dtype)
  #     view = [1] * input.dim()
  #     view[0] = -1
  #     return rho.view(view).transpose(0, dim)
  #
  #
  # def _threshold_and_support(input, dim=0):
  #     """Sparsemax building block: compute the threshold
  #     Args:
  #         input: any dimension
  #         dim: dimension along which to apply the sparsemax
  #     Returns:
  #         the threshold value
  #     """
  #
  #     input_srt, _ = torch.sort(input, descending=True, dim=dim)
  #     input_cumsum = input_srt.cumsum(dim) - 1
  #     rhos = _make_ix_like(input, dim)
  #     support = rhos * input_srt > input_cumsum
  #
  #     support_size = support.sum(dim=dim).unsqueeze(dim)
  #     tau = input_cumsum.gather(dim, support_size - 1)
  #     tau /= support_size.to(input.dtype)
  #     return tau, support_size
  #
  #
  # class SparsemaxFunction(Function):
  #
  #     @staticmethod
  #     def forward(ctx, input, dim=0):
  #         """sparsemax: normalizing sparse transform (a la softmax)
  #         Parameters:
  #             input (Tensor): any shape
  #             dim: dimension along which to apply sparsemax
  #         Returns:
  #             output (Tensor): same shape as input
  #         """
  #         ctx.dim = dim
  #         max_val, _ = input.max(dim=dim, keepdim=True)
  #         input -= max_val  # same numerical stability trick as for softmax
  #         tau, supp_size = _threshold_and_support(input, dim=dim)
  #         output = torch.clamp(input - tau, min=0)
  #         ctx.save_for_backward(supp_size, output)
  #         return output
  #
  #     @staticmethod
  #     def backward(ctx, grad_output):
  #         supp_size, output = ctx.saved_tensors
  #         dim = ctx.dim
  #         grad_input = grad_output.clone()
  #         grad_input[output == 0] = 0
  #
  #         v_hat = grad_input.sum(dim=dim) / supp_size.to(output.dtype).squeeze()
  #         v_hat = v_hat.unsqueeze(dim)
  #         grad_input = torch.where(output != 0, grad_input - v_hat, grad_input)
  #         return grad_input, None
  #
  #
  # sparsemax = SparsemaxFunction.apply
  #
  #
  # class Sparsemax(nn.Module):
  #
  #     def __init__(self, dim=0):
  #         self.dim = dim
  #         super(Sparsemax, self).__init__()
  #
  #     def forward(self, input):
  #         return sparsemax(input, self.dim)
  