code: def _entmax15_forward(X, dim, k):
      max_val, _ = X.max(dim=dim, keepdim=True)
      X = X - max_val  # same numerical stability trick as for softmax
      X = X / 2  # divide by 2 to solve actual Entmax
  
      tau_star, _ = _entmax_threshold_and_support(X, dim=dim, k=k)
  
      Y = torch.clamp(X - tau_star, min=0) ** 2
      return Y, {}
  
imports: import torch
  import torch.nn
  from torch.autograd import Function
label_desc: Use image node_id 6 for calling a global function with example usage: _entmax15_forward(X, dim, k) and returns: Y,  dict