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: :param s: :math:`(b\times n_1 \times n_2)` input 3d tensor. :math:`b`: batch size
  :param nrows: :math:`(b)` number of objects in dim1
  :param ncols: :math:`(b)` number of objects in dim2
  :param dummy_row: whether to add dummy rows (rows whose elements are all 0) to pad the matrix to square matrix.
   default: ``False``
  :return: :math:`(b\times n_1 \times n_2)` the computed doubly-stochastic matrix
  .. note::
      We support batched instances with different number of nodes, therefore ``nrows`` and ``ncols`` are
      required to specify the exact number of objects of each dimension in the batch. If not specified, we assume
      the batched matrices are not padded.
  .. note::
      The original Sinkhorn algorithm only works for square matrices. To handle cases where the graphs to be
      matched have different number of nodes, it is a common practice to add dummy rows to construct a square
      matrix. After the row and column normalizations, the padded rows are discarded.
  .. note::
      We assume row number <= column number. If not, the input matrix will be transposed.
docstr_tok: [":","param","s",":",":","math",":","`","(","b\\times","n_1","\\times","n_2",")","`","input","3d","tensor",".",":","math",":","`","b","`",":","batch","size",":","param","nrows",":",":","math",":","`","(","b",")","`","number","of","objects","in","dim1",":","param","ncols",":",":","math",":","`","(","b",")","`","number","of","objects","in","dim2",":","param","dummy_row",":","whether","to","add","dummy","rows","(","rows","whose","elements","are","all","0",")","to","pad","the","matrix","to","square","matrix",".","default",":","``","False","``",":","return",":",":","math",":","`","(","b\\times","n_1","\\times","n_2",")","`","the","computed","doubly-stochastic","matrix","..","note",":",":","We","support","batched","instances","with","different","number","of","nodes",",","therefore","``","nrows","``","and","``","ncols","``","are","required","to","specify","the","exact","number","of","objects","of","each","dimension","in","the","batch",".","If","not","specified",",","we","assume","the","batched","matrices","are","not","padded",".","..","note",":",":","The","original","Sinkhorn","algorithm","only","works","for","square","matrices",".","To","handle","cases","where","the","graphs","to","be","matched","have","different","number","of","nodes",",","it","is","a","common","practice","to","add","dummy","rows","to","construct","a","square","matrix",".","After","the","row","and","column","normalizations",",","the","padded","rows","are","discarded",".","..","note",":",":","We","assume","row","number","<","=","column","number",".","If","not",",","the","input","matrix","will","be","transposed","."]
code: def forward(
      self,
      s: Tensor,
      nrows: Tensor = None,
      ncols: Tensor = None,
      dummy_row: bool = False,
  ) -> Tensor:
      r"""
      :param s: :math:`(b\times n_1 \times n_2)` input 3d tensor. :math:`b`: batch size
      :param nrows: :math:`(b)` number of objects in dim1
      :param ncols: :math:`(b)` number of objects in dim2
      :param dummy_row: whether to add dummy rows (rows whose elements are all 0) to pad the matrix to square matrix.
       default: ``False``
      :return: :math:`(b\times n_1 \times n_2)` the computed doubly-stochastic matrix
      .. note::
          We support batched instances with different number of nodes, therefore ``nrows`` and ``ncols`` are
          required to specify the exact number of objects of each dimension in the batch. If not specified, we assume
          the batched matrices are not padded.
      .. note::
          The original Sinkhorn algorithm only works for square matrices. To handle cases where the graphs to be
          matched have different number of nodes, it is a common practice to add dummy rows to construct a square
          matrix. After the row and column normalizations, the padded rows are discarded.
      .. note::
          We assume row number <= column number. If not, the input matrix will be transposed.
      """
      if self.log_forward:
          return self.forward_log(s, nrows, ncols, dummy_row)
      else:
          return self.forward_ori(s, nrows, ncols, dummy_row)
  
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 2 for calling the Sinkhorn obj's underlying member method code with example usage: obj.forward(s, nrows, ncols, dummy_row) and returns: self, self