class_docstr: Linear transformation according to types.
  
  For each sample of the input batch :math:`x \in X`, apply linear transformation
  :math:`xW_t`, where :math:`t` is the type of :math:`x`.
  
  The module supports two regularization methods (basis-decomposition and
  block-diagonal-decomposition) proposed by "`Modeling Relational Data
  with Graph Convolutional Networks <https://arxiv.org/abs/1703.06103>`__"
  
  The basis regularization decomposes :math:`W_t` by:
  
  .. math::
  
     W_t^{(l)} = \sum_{b=1}^B a_{tb}^{(l)}V_b^{(l)}
  
  where :math:`B` is the number of bases, :math:`V_b^{(l)}` are linearly combined
  with coefficients :math:`a_{tb}^{(l)}`.
  
  The block-diagonal-decomposition regularization decomposes :math:`W_t` into :math:`B`
  block-diagonal matrices. We refer to :math:`B` as the number of bases:
  
  .. math::
  
     W_t^{(l)} = \oplus_{b=1}^B Q_{tb}^{(l)}
  
  where :math:`B` is the number of bases, :math:`Q_{tb}^{(l)}` are block
  bases with shape :math:`R^{(d^{(l+1)}/B)\times(d^{l}/B)}`.
  
  Parameters
  ----------
  in_size : int
      Input feature size.
  out_size : int
      Output feature size.
  num_types : int
      Total number of types.
  regularizer : str, optional
      Which weight regularizer to use "basis" or "bdd":
  
       - "basis" is short for basis-decomposition.
       - "bdd" is short for block-diagonal-decomposition.
  
      Default applies no regularization.
  num_bases : int, optional
      Number of bases. Needed when ``regularizer`` is specified. Typically smaller
      than ``num_types``.
      Default: ``None``.
  
  Examples
  --------
  
  No regularization.
  
  >>> from dgl.nn import TypedLinear
  >>> import torch
  >>>
  >>> x = torch.randn(100, 32)
  >>> x_type = torch.randint(0, 5, (100,))
  >>> m = TypedLinear(32, 64, 5)
  >>> y = m(x, x_type)
  >>> print(y.shape)
  torch.Size([100, 64])
  
  With basis regularization
  
  >>> x = torch.randn(100, 32)
  >>> x_type = torch.randint(0, 5, (100,))
  >>> m = TypedLinear(32, 64, 5, regularizer='basis', num_bases=4)
  >>> y = m(x, x_type)
  >>> print(y.shape)
  torch.Size([100, 64])
class_docstr_tok: ["Linear","transformation","according","to","types",".","For","each","sample","of","the","input","batch",":","math",":","`","x","\\in","X","`",",","apply","linear","transformation",":","math",":","`","xW_t","`",",","where",":","math",":","`","t","`","is","the","type","of",":","math",":","`","x","`",".","The","module","supports","two","regularization","methods","(","basis-decomposition","and","block-diagonal-decomposition",")","proposed","by","``","`","Modeling","Relational","Data","with","Graph","Convolutional","Networks","<","https",":","\/\/arxiv.org\/abs\/1703.06103",">","`","__","''","The","basis","regularization","decomposes",":","math",":","`","W_t","`","by",":","..","math",":",":","W_t^","{","(","l",")","}","=","\\sum_","{","b=1","}","^B","a_","{","tb","}","^","{","(","l",")","}","V_b^","{","(","l",")","}","where",":","math",":","`","B","`","is","the","number","of","bases",",",":","math",":","`","V_b^","{","(","l",")","}","`","are","linearly","combined","with","coefficients",":","math",":","`","a_","{","tb","}","^","{","(","l",")","}","`",".","The","block-diagonal-decomposition","regularization","decomposes",":","math",":","`","W_t","`","into",":","math",":","`","B","`","block-diagonal","matrices",".","We","refer","to",":","math",":","`","B","`","as","the","number","of","bases",":","..","math",":",":","W_t^","{","(","l",")","}","=","\\oplus_","{","b=1","}","^B","Q_","{","tb","}","^","{","(","l",")","}","where",":","math",":","`","B","`","is","the","number","of","bases",",",":","math",":","`","Q_","{","tb","}","^","{","(","l",")","}","`","are","block","bases","with","shape",":","math",":","`","R^","{","(","d^","{","(","l+1",")","}","\/B",")","\\times","(","d^","{","l","}","\/B",")","}","`",".","Parameters","--","--","--","--","--","in_size",":","int","Input","feature","size",".","out_size",":","int","Output","feature","size",".","num_types",":","int","Total","number","of","types",".","regularizer",":","str",",","optional","Which","weight","regularizer","to","use","``","basis","''","or","``","bdd","''",":","-","``","basis","''","is","short","for","basis-decomposition",".","-","``","bdd","''","is","short","for","block-diagonal-decomposition",".","Default","applies","no","regularization",".","num_bases",":","int",",","optional","Number","of","bases",".","Needed","when","``","regularizer","``","is","specified",".","Typically","smaller","than","``","num_types","``",".","Default",":","``","None","``",".","Examples","--","--","--","--","No","regularization",".",">",">",">","from","dgl.nn","import","TypedLinear",">",">",">","import","torch",">",">",">",">",">",">","x","=","torch.randn","(","100",",","32",")",">",">",">","x_type","=","torch.randint","(","0",",","5",",","(","100",",",")",")",">",">",">","m","=","TypedLinear","(","32",",","64",",","5",")",">",">",">","y","=","m","(","x",",","x_type",")",">",">",">","print","(","y.shape",")","torch.Size","(","[","100",",","64","]",")","With","basis","regularization",">",">",">","x","=","torch.randn","(","100",",","32",")",">",">",">","x_type","=","torch.randint","(","0",",","5",",","(","100",",",")",")",">",">",">","m","=","TypedLinear","(","32",",","64",",","5",",","regularizer='basis","'",",","num_bases=4",")",">",">",">","y","=","m","(","x",",","x_type",")",">",">",">","print","(","y.shape",")","torch.Size","(","[","100",",","64","]",")"]
docstr: Forward computation.
  
  Parameters
  ----------
  x : torch.Tensor
      A 2D input tensor. Shape: (N, D1)
  x_type : torch.Tensor
      A 1D integer tensor storing the type of the elements in ``x`` with one-to-one
      correspondenc. Shape: (N,)
  sorted_by_type : bool, optional
      Whether the inputs have been sorted by the types. Forward on pre-sorted inputs may
      be faster.
  
  Returns
  -------
  y : torch.Tensor
      The transformed output tensor. Shape: (N, D2)
docstr_tok: ["Forward","computation",".","Parameters","--","--","--","--","--","x",":","torch.Tensor","A","2D","input","tensor",".","Shape",":","(","N",",","D1",")","x_type",":","torch.Tensor","A","1D","integer","tensor","storing","the","type","of","the","elements","in","``","x","``","with","one-to-one","correspondenc",".","Shape",":","(","N",",",")","sorted_by_type",":","bool",",","optional","Whether","the","inputs","have","been","sorted","by","the","types",".","Forward","on","pre-sorted","inputs","may","be","faster",".","Returns","--","--","--","-","y",":","torch.Tensor","The","transformed","output","tensor",".","Shape",":","(","N",",","D2",")"]
code: def forward(self, x, x_type, sorted_by_type=False):
      """Forward computation.
  
      Parameters
      ----------
      x : torch.Tensor
          A 2D input tensor. Shape: (N, D1)
      x_type : torch.Tensor
          A 1D integer tensor storing the type of the elements in ``x`` with one-to-one
          correspondenc. Shape: (N,)
      sorted_by_type : bool, optional
          Whether the inputs have been sorted by the types. Forward on pre-sorted inputs may
          be faster.
  
      Returns
      -------
      y : torch.Tensor
          The transformed output tensor. Shape: (N, D2)
      """
      w = self.get_weight()
      if self.regularizer == "bdd":
          w = w.index_select(0, x_type).view(
              -1, self.submat_in, self.submat_out
          )
          x = x.view(-1, 1, self.submat_in)
          return torch.bmm(x, w).view(-1, self.out_size)
      elif sorted_by_type:
          pos_l = torch.searchsorted(
              x_type, torch.arange(self.num_types, device=x.device)
          )
          pos_r = torch.cat(
              [pos_l[1:], torch.tensor([len(x_type)], device=x.device)]
          )
          seglen = (
              pos_r - pos_l
          ).cpu()  # XXX(minjie): cause device synchronize
          return segment_mm(x, w, seglen_a=seglen)
      else:
          return gather_mm(x, w, idx_b=x_type)
  
file_path: dgl/python/dgl/nn/pytorch/linear.py
imports: import math
  import torch
  import torch.nn
  from ...ops import gather_mm, segment_mm
label_desc: Use image node_id 4 for calling the TypedLinear obj's underlying member method code with example usage: obj.forward(x, x_type, sorted_by_type) and returns: torch, segment_mm, gather_mm