class_docstr: Gumbel Sinkhorn Layer turns the input matrix into a bi-stochastic matrix.
  See details in `"Mena et al. Learning Latent Permutations with Gumbel-Sinkhorn Networks. ICLR 2018"
  <https://arxiv.org/abs/1802.08665>`_
  :param max_iter: maximum iterations (default: ``10``)
  :param tau: the hyper parameter :math:`     au` controlling the temperature (default: ``1``)
  :param epsilon: a small number for numerical stability (default: ``1e-4``)
  :param batched_operation: apply batched_operation for better efficiency (but may cause issues for back-propagation,
   default: ``False``)
  .. note::
      This module only supports log-scale Sinkhorn operation.
class_docstr_tok: ["Gumbel","Sinkhorn","Layer","turns","the","input","matrix","into","a","bi-stochastic","matrix",".","See","details","in","`","``","Mena","et","al",".","Learning","Latent","Permutations","with","Gumbel-Sinkhorn","Networks",".","ICLR","2018","''","<","https",":","\/\/arxiv.org\/abs\/1802.08665",">","`","_",":","param","max_iter",":","maximum","iterations","(","default",":","``","10","``",")",":","param","tau",":","the","hyper","parameter",":","math",":","`","au","`","controlling","the","temperature","(","default",":","``","1","``",")",":","param","epsilon",":","a","small","number","for","numerical","stability","(","default",":","``","1e-4","``",")",":","param","batched_operation",":","apply","batched_operation","for","better","efficiency","(","but","may","cause","issues","for","back-propagation",",","default",":","``","False","``",")","..","note",":",":","This","module","only","supports","log-scale","Sinkhorn","operation","."]
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 sample_num: number of samples
  :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 m\times n_1 \times n_2)` the computed doubly-stochastic matrix. :math:`m`: number of samples
   (``sample_num``)
  The samples are stacked at the fist dimension of the output tensor. You may reshape the output tensor ``s`` as:
  ::
      s = torch.reshape(s, (-1, sample_num, s.shape[1], s.shape[2]))
  .. 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","sample_num",":","number","of","samples",":","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","m\\times","n_1","\\times","n_2",")","`","the","computed","doubly-stochastic","matrix",".",":","math",":","`","m","`",":","number","of","samples","(","``","sample_num","``",")","The","samples","are","stacked","at","the","fist","dimension","of","the","output","tensor",".","You","may","reshape","the","output","tensor","``","s","``","as",":",":",":","s","=","torch.reshape","(","s",",","(","-1",",","sample_num",",","s.shape","[","1","]",",","s.shape","[","2","]",")",")","..","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,
      sample_num=5,
      dummy_row=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 sample_num: number of samples
      :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 m\times n_1 \times n_2)` the computed doubly-stochastic matrix. :math:`m`: number of samples
       (``sample_num``)
      The samples are stacked at the fist dimension of the output tensor. You may reshape the output tensor ``s`` as:
      ::
          s = torch.reshape(s, (-1, sample_num, s.shape[1], s.shape[2]))
      .. 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.
      """
  
      def sample_gumbel(t_like, eps=1e-20):
          """
          randomly sample standard gumbel variables
          """
          u = torch.empty_like(t_like).uniform_()
          return -torch.log(-torch.log(u + eps) + eps)
  
      s_rep = torch.repeat_interleave(s, sample_num, dim=0)
      s_rep = s_rep + sample_gumbel(s_rep)
      nrows_rep = torch.repeat_interleave(nrows, sample_num, dim=0)
      ncols_rep = torch.repeat_interleave(ncols, sample_num, dim=0)
      s_rep = self.sinkhorn(s_rep, nrows_rep, ncols_rep, dummy_row)
      # s_rep = torch.reshape(s_rep, (-1, sample_num, s_rep.shape[1], s_rep.shape[2]))
      return s_rep
  
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 GumbelSinkhorn obj's underlying member method code with example usage: obj.forward(s, nrows, ncols, sample_num, dummy_row) and returns: s_rep, unknown