class_docstr: Spatial Encoder, as introduced in
  `Do Transformers Really Perform Bad for Graph Representation?
  <https://proceedings.neurips.cc/paper/2021/file/f1c1592588411002af340cbaedd6fc33-Paper.pdf>`__
  
  This module is a learnable spatial embedding module, which encodes
  the shortest distance between each node pair for attention bias.
  
  Parameters
  ----------
  max_dist : int
      Upper bound of the shortest path distance
      between each node pair to be encoded.
      All distance will be clamped into the range `[0, max_dist]`.
  num_heads : int, optional
      Number of attention heads if multi-head attention mechanism is applied.
      Default : 1.
  
  Examples
  --------
  >>> import torch as th
  >>> import dgl
  >>> from dgl.nn import SpatialEncoder
  >>> from dgl import shortest_dist
  
  >>> g1 = dgl.graph(([0,0,0,1,1,2,3,3], [1,2,3,0,3,0,0,1]))
  >>> g2 = dgl.graph(([0,1], [1,0]))
  >>> n1, n2 = g1.num_nodes(), g2.num_nodes()
  >>> # use -1 padding since shortest_dist returns -1 for unreachable node pairs
  >>> dist = -th.ones((2, 4, 4), dtype=th.long)
  >>> dist[0, :n1, :n1] = shortest_dist(g1, root=None, return_paths=False)
  >>> dist[1, :n2, :n2] = shortest_dist(g2, root=None, return_paths=False)
  >>> spatial_encoder = SpatialEncoder(max_dist=2, num_heads=8)
  >>> out = spatial_encoder(dist)
  >>> print(out.shape)
  torch.Size([2, 4, 4, 8])
class_docstr_tok: ["Spatial","Encoder",",","as","introduced","in","`","Do","Transformers","Really","Perform","Bad","for","Graph","Representation","?","<","https",":","\/\/proceedings.neurips.cc\/paper\/2021\/file\/f1c1592588411002af340cbaedd6fc33-Paper.pdf",">","`","__","This","module","is","a","learnable","spatial","embedding","module",",","which","encodes","the","shortest","distance","between","each","node","pair","for","attention","bias",".","Parameters","--","--","--","--","--","max_dist",":","int","Upper","bound","of","the","shortest","path","distance","between","each","node","pair","to","be","encoded",".","All","distance","will","be","clamped","into","the","range","`","[","0",",","max_dist","]","`",".","num_heads",":","int",",","optional","Number","of","attention","heads","if","multi-head","attention","mechanism","is","applied",".","Default",":","1",".","Examples","--","--","--","--",">",">",">","import","torch","as","th",">",">",">","import","dgl",">",">",">","from","dgl.nn","import","SpatialEncoder",">",">",">","from","dgl","import","shortest_dist",">",">",">","g1","=","dgl.graph","(","(","[","0,0,0,1,1,2,3,3","]",",","[","1,2,3,0,3,0,0,1","]",")",")",">",">",">","g2","=","dgl.graph","(","(","[","0,1","]",",","[","1,0","]",")",")",">",">",">","n1",",","n2","=","g1.num_nodes","(",")",",","g2.num_nodes","(",")",">",">",">","#","use","-1","padding","since","shortest_dist","returns","-1","for","unreachable","node","pairs",">",">",">","dist","=","-th.ones","(","(","2",",","4",",","4",")",",","dtype=th.long",")",">",">",">","dist","[","0",",",":","n1",",",":","n1","]","=","shortest_dist","(","g1",",","root=None",",","return_paths=False",")",">",">",">","dist","[","1",",",":","n2",",",":","n2","]","=","shortest_dist","(","g2",",","root=None",",","return_paths=False",")",">",">",">","spatial_encoder","=","SpatialEncoder","(","max_dist=2",",","num_heads=8",")",">",">",">","out","=","spatial_encoder","(","dist",")",">",">",">","print","(","out.shape",")","torch.Size","(","[","2",",","4",",","4",",","8","]",")"]
docstr: Parameters
  ----------
  dist : Tensor
      Shortest path distance of the batched graph with -1 padding, a tensor
      of shape :math:`(B, N, N)`, where :math:`B` is the batch size of
      the batched graph, and :math:`N` is the maximum number of nodes.
  
  Returns
  -------
  torch.Tensor
      Return attention bias as spatial encoding of shape
      :math:`(B, N, N, H)`, where :math:`H` is :attr:`num_heads`.
docstr_tok: ["Parameters","--","--","--","--","--","dist",":","Tensor","Shortest","path","distance","of","the","batched","graph","with","-1","padding",",","a","tensor","of","shape",":","math",":","`","(","B",",","N",",","N",")","`",",","where",":","math",":","`","B","`","is","the","batch","size","of","the","batched","graph",",","and",":","math",":","`","N","`","is","the","maximum","number","of","nodes",".","Returns","--","--","--","-","torch.Tensor","Return","attention","bias","as","spatial","encoding","of","shape",":","math",":","`","(","B",",","N",",","N",",","H",")","`",",","where",":","math",":","`","H","`","is",":","attr",":","`","num_heads","`","."]
code: def forward(self, dist):
      """
      Parameters
      ----------
      dist : Tensor
          Shortest path distance of the batched graph with -1 padding, a tensor
          of shape :math:`(B, N, N)`, where :math:`B` is the batch size of
          the batched graph, and :math:`N` is the maximum number of nodes.
  
      Returns
      -------
      torch.Tensor
          Return attention bias as spatial encoding of shape
          :math:`(B, N, N, H)`, where :math:`H` is :attr:`num_heads`.
      """
      spatial_encoding = self.embedding_table(
          th.clamp(
              dist,
              min=-1,
              max=self.max_dist,
          )
          + 1
      )
      return spatial_encoding
  
file_path: dgl/python/dgl/nn/pytorch/gt/spatial_encoder.py
imports: import torch
  import torch.nn
  import torch.nn.functional
label_desc: Use image node_id 2 for calling the SpatialEncoder obj's underlying member method code with example usage: obj.forward(dist) and returns: spatial_encoding