code: def _entmax_threshold_and_support(X, dim=-1, k=None):
      """Core computation for 1.5-entmax: optimal threshold and support size.
  
      Parameters
      ----------
      X : torch.Tensor
          The input tensor to compute thresholds over.
  
      dim : int
          The dimension along which to apply 1.5-entmax.
  
      k : int or None
          number of largest elements to partial-sort over. For optimal
          performance, should be slightly bigger than the expected number of
          nonzeros in the solution. If the solution is more than k-sparse,
          this function is recursively called with a 2*k schedule.
          If `None`, full sorting is performed from the beginning.
  
      Returns
      -------
      tau : torch.Tensor like `X`, with all but the `dim` dimension intact
          the threshold value for each vector
      support_size : torch LongTensor, shape like `tau`
          the number of nonzeros in each vector.
      """
  
      if k is None or k >= X.shape[dim]:  # do full sort
          Xsrt, _ = torch.sort(X, dim=dim, descending=True)
      else:
          Xsrt, _ = torch.topk(X, k=k, dim=dim)
  
      rho = _make_ix_like(Xsrt, 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)
  
      if k is not None and k < X.shape[dim]:
          unsolved = (support_size == k).squeeze(dim)
  
          if torch.any(unsolved):
              X_ = _roll_last(X, dim)[unsolved]
              tau_, ss_ = _entmax_threshold_and_support(
                  X_, dim=-1, k=2 * k
              )
              _roll_last(tau_star, dim)[unsolved] = tau_
              _roll_last(support_size, dim)[unsolved] = ss_
  
      return tau_star, support_size
  
imports: import torch
  import torch.nn
  import torch.nn.functional
  from torch import Tensor
  from torch.autograd import Function
  from torch.jit import script
label_desc: Use image node_id 6 for calling a global function with example usage: _entmax_threshold_and_support(X, dim, k) and returns: tau_star,  support_size