class_docstr: Sinkhorn algorithm turns the input matrix into a bi-stochastic matrix.
  Sinkhorn algorithm firstly applies an ``exp`` function with temperature :math:`\tau`:
  .. math::
      \mathbf{S}_{i,j} = \exp \left(\frac{\mathbf{s}_{i,j}}{\tau}\right)
  And then turns the matrix into doubly-stochastic matrix by iterative row- and column-wise normalization:
  .. math::
      \mathbf{S} &= \mathbf{S} \oslash (\mathbf{1}_{n_2} \mathbf{1}_{n_2}^\top \cdot \mathbf{S}) \\
      \mathbf{S} &= \mathbf{S} \oslash (\mathbf{S} \cdot \mathbf{1}_{n_2} \mathbf{1}_{n_2}^\top)
  where :math:`\oslash` means element-wise division, :math:`\mathbf{1}_n` means a column-vector with length :math:`n`
  whose elements are all :math:`1`\ s.
  :param max_iter: maximum iterations (default: ``10``)
  :param tau: the hyper parameter :math:`\tau` controlling the temperature (default: ``1``)
  :param epsilon: a small number for numerical stability (default: ``1e-4``)
  :param log_forward: apply log-scale computation for better numerical stability (default: ``True``)
  :param batched_operation: apply batched_operation for better efficiency (but may cause issues for back-propagation,
   default: ``False``)
  .. note::
      ``tau`` is an important hyper parameter to be set for Sinkhorn algorithm. ``tau`` controls the distance between
      the predicted doubly-stochastic matrix, and the discrete permutation matrix computed by Hungarian algorithm (see
      :func:`~src.lap_solvers.hungarian.hungarian`). Given a small ``tau``, Sinkhorn performs more closely to
      Hungarian, at the cost of slower convergence speed and reduced numerical stability.
  .. note::
      We recommend setting ``log_forward=True`` because it is more numerically stable. It provides more precise
      gradient in back propagation and helps the model to converge better and faster.
  .. note::
      Setting ``batched_operation=True`` may be preferred when you are doing inference with this module and do not
      need the gradient.
class_docstr_tok: ["Sinkhorn","algorithm","turns","the","input","matrix","into","a","bi-stochastic","matrix",".","Sinkhorn","algorithm","firstly","applies","an","``","exp","``","function","with","temperature",":","math",":","`","\\tau","`",":","..","math",":",":","\\mathbf","{","S","}","_","{","i",",","j","}","=","\\exp","\\left","(","\\frac","{","\\mathbf","{","s","}","_","{","i",",","j","}","}","{","\\tau","}","\\right",")","And","then","turns","the","matrix","into","doubly-stochastic","matrix","by","iterative","row-","and","column-wise","normalization",":","..","math",":",":","\\mathbf","{","S","}","&","=","\\mathbf","{","S","}","\\oslash","(","\\mathbf","{","1","}","_","{","n_2","}","\\mathbf","{","1","}","_","{","n_2","}","^\\top","\\cdot","\\mathbf","{","S","}",")","\\\\","\\mathbf","{","S","}","&","=","\\mathbf","{","S","}","\\oslash","(","\\mathbf","{","S","}","\\cdot","\\mathbf","{","1","}","_","{","n_2","}","\\mathbf","{","1","}","_","{","n_2","}","^\\top",")","where",":","math",":","`","\\oslash","`","means","element-wise","division",",",":","math",":","`","\\mathbf","{","1","}","_n","`","means","a","column-vector","with","length",":","math",":","`","n","`","whose","elements","are","all",":","math",":","`","1","`","\\","s.",":","param","max_iter",":","maximum","iterations","(","default",":","``","10","``",")",":","param","tau",":","the","hyper","parameter",":","math",":","`","\\tau","`","controlling","the","temperature","(","default",":","``","1","``",")",":","param","epsilon",":","a","small","number","for","numerical","stability","(","default",":","``","1e-4","``",")",":","param","log_forward",":","apply","log-scale","computation","for","better","numerical","stability","(","default",":","``","True","``",")",":","param","batched_operation",":","apply","batched_operation","for","better","efficiency","(","but","may","cause","issues","for","back-propagation",",","default",":","``","False","``",")","..","note",":",":","``","tau","``","is","an","important","hyper","parameter","to","be","set","for","Sinkhorn","algorithm",".","``","tau","``","controls","the","distance","between","the","predicted","doubly-stochastic","matrix",",","and","the","discrete","permutation","matrix","computed","by","Hungarian","algorithm","(","see",":","func",":","`","~src.lap_solvers.hungarian.hungarian","`",")",".","Given","a","small","``","tau","``",",","Sinkhorn","performs","more","closely","to","Hungarian",",","at","the","cost","of","slower","convergence","speed","and","reduced","numerical","stability",".","..","note",":",":","We","recommend","setting","``","log_forward=True","``","because","it","is","more","numerically","stable",".","It","provides","more","precise","gradient","in","back","propagation","and","helps","the","model","to","converge","better","and","faster",".","..","note",":",":","Setting","``","batched_operation=True","``","may","be","preferred","when","you","are","doing","inference","with","this","module","and","do","not","need","the","gradient","."]
docstr: Computing sinkhorn with row/column normalization.
  .. warning::
      This function is deprecated because :meth:`~src.lap_solvers.sinkhorn.Sinkhorn.forward_log` is more
      numerically stable.
docstr_tok: ["Computing","sinkhorn","with","row\/column","normalization",".","..","warning",":",":","This","function","is","deprecated","because",":","meth",":","`","~src.lap_solvers.sinkhorn.Sinkhorn.forward_log","`","is","more","numerically","stable","."]
code: def forward_ori(self, s, nrows=None, ncols=None, dummy_row=False):
      r"""
      Computing sinkhorn with row/column normalization.
      .. warning::
          This function is deprecated because :meth:`~src.lap_solvers.sinkhorn.Sinkhorn.forward_log` is more
          numerically stable.
      """
      if len(s.shape) == 2:
          s = s.unsqueeze(0)
          matrix_input = True
      elif len(s.shape) == 3:
          matrix_input = False
      else:
          raise ValueError("input data shape not understood.")
  
      batch_size = s.shape[0]
  
      # s = s.to(dtype=dtype)
  
      if nrows is None:
          nrows = [s.shape[1] for _ in range(batch_size)]
      if ncols is None:
          ncols = [s.shape[2] for _ in range(batch_size)]
  
      # tau scaling
      ret_s = torch.zeros_like(s)
      for b, n in enumerate(nrows):
          ret_s[b, 0:n, 0 : ncols[b]] = nn.functional.softmax(
              s[b, 0:n, 0 : ncols[b]] / self.tau, dim=-1
          )
      s = ret_s
  
      # add dummy elements
      if dummy_row:
          dummy_shape = list(s.shape)
          dummy_shape[1] = s.shape[2] - s.shape[1]
          # s = torch.cat((s, torch.full(dummy_shape, self.epsilon * 10).to(s.device)), dim=1)
          # nrows = nrows + dummy_shape[1] # non in-place
          s = torch.cat(
              (s, torch.full(dummy_shape, 0.0).to(s.device)), dim=1
          )
          ori_nrows = nrows
          nrows = ncols
          for b in range(batch_size):
              s[b, ori_nrows[b] : nrows[b], : ncols[b]] = self.epsilon
  
      row_norm_ones = torch.zeros(
          batch_size,
          s.shape[1],
          s.shape[1],
          device=s.device,
          dtype=s.dtype,
      )  # size: row x row
      col_norm_ones = torch.zeros(
          batch_size,
          s.shape[2],
          s.shape[2],
          device=s.device,
          dtype=s.dtype,
      )  # size: col x col
      for b in range(batch_size):
          row_slice = slice(0, nrows[b])
          col_slice = slice(0, ncols[b])
          row_norm_ones[b, row_slice, row_slice] = 1
          col_norm_ones[b, col_slice, col_slice] = 1
  
      s += self.epsilon
  
      for i in range(self.max_iter):
          if i % 2 == 0:
              # column norm
              # ones = torch.ones(batch_size, s.shape[1], s.shape[1], device=s.device)
              sum = torch.sum(
                  torch.mul(s.unsqueeze(3), col_norm_ones.unsqueeze(1)),
                  dim=2,
              )
          else:
              # row norm
              # ones = torch.ones(batch_size, s.shape[2], s.shape[2], device=s.device)
              sum = torch.sum(
                  torch.mul(row_norm_ones.unsqueeze(3), s.unsqueeze(1)),
                  dim=2,
              )
  
          tmp = torch.zeros_like(s)
          for b in range(batch_size):
              row_slice = slice(
                  0, nrows[b] if nrows is not None else s.shape[2]
              )
              col_slice = slice(
                  0, ncols[b] if ncols is not None else s.shape[1]
              )
              tmp[b, row_slice, col_slice] = (
                  1 / sum[b, row_slice, col_slice]
              )
          s = s * tmp
  
      if dummy_row:
          if dummy_shape[1] > 0:
              s = s[:, : -dummy_shape[1]]
          for b in range(batch_size):
              s[b, ori_nrows[b] : nrows[b], : ncols[b]] = 0
  
      if matrix_input:
          s.squeeze_(0)
  
      return s
  
file_path: DIG/dig/auggraph/method/SMixup/utils/sinkhorn.py
imports: import torch
  import torch.nn
  from torch import Tensor
  import pygmtools
label_desc: Use image node_id 4 for calling the Sinkhorn obj's underlying member method code with example usage: obj.forward_ori(s, nrows, ncols, dummy_row) and returns: s