code: def _entmax_bisect_forward(
      X, alpha, dim, n_iter, ensure_sum_one, cls=EntmaxBisectFunction
  ):
      if not isinstance(alpha, torch.Tensor):
          alpha = torch.tensor(alpha, dtype=X.dtype, device=X.device)
  
      alpha_shape = list(X.shape)
      alpha_shape[dim] = 1
      alpha = alpha.expand(*alpha_shape)
  
      d = X.shape[dim]
  
      max_val, _ = X.max(dim=dim, keepdim=True)
      X = X * (alpha - 1)
      max_val = max_val * (alpha - 1)
  
      # Note: when alpha < 1, tau_lo > tau_hi. This still works since dm < 0.
      tau_lo = max_val - cls._gp(1, alpha)
      tau_hi = max_val - cls._gp(1 / d, alpha)
  
      f_lo = cls._p(X - tau_lo, alpha).sum(dim) - 1
  
      dm = tau_hi - tau_lo
  
      for it in range(n_iter):
          dm /= 2
          tau_m = tau_lo + dm
          p_m = cls._p(X - tau_m, alpha)
          f_m = p_m.sum(dim) - 1
  
          mask = (f_m * f_lo >= 0).unsqueeze(dim)
          tau_lo = torch.where(mask, tau_m, tau_lo)
  
      if ensure_sum_one:
          p_m /= p_m.sum(dim=dim).unsqueeze(dim=dim)
  
      return p_m, {"alpha": alpha, "dim": dim}
  
imports: import torch
  import torch.nn
  from torch.autograd import Function
label_desc: Use image node_id 1 for calling a global function with example usage: _entmax_bisect_forward(X, alpha, dim, n_iter, ensure_sum_one, cls) and returns: p_m,  dict