class_docstr: 3D Spatial Encoder, as introduced in
  `One Transformer Can Understand Both 2D & 3D Molecular Data
  <https://arxiv.org/pdf/2210.01765.pdf>`__
  
  This module encodes pair-wise relation between node pair :math:`(i,j)` in
  the 3D geometric space, according to the Gaussian Basis Kernel function:
  
  :math:`\psi _{(i,j)} ^k = \frac{1}{\sqrt{2\pi} \lvert \sigma^k \rvert}
  \exp{\left ( -\frac{1}{2} \left( \frac{\gamma_{(i,j)} \lvert \lvert r_i -
  r_j \rvert \rvert + \beta_{(i,j)} - \mu^k}{\lvert \sigma^k \rvert} \right)
  ^2 \right)}，k=1,...,K,`
  
  where :math:`K` is the number of Gaussian Basis kernels. :math:`r_i` is the
  Cartesian coordinate of node :math:`i`.
  :math:`\gamma_{(i,j)}, \beta_{(i,j)}` are learnable scaling factors and
  biases determined by node types. :math:`\mu^k, \sigma^k` are learnable
  centers and standard deviations of the Gaussian Basis kernels.
  
  Parameters
  ----------
  num_kernels : int
      Number of Gaussian Basis Kernels to be applied. Each Gaussian Basis
      Kernel contains a learnable kernel center and a learnable standard
      deviation.
  num_heads : int, optional
      Number of attention heads if multi-head attention mechanism is applied.
      Default : 1.
  max_node_type : int, optional
      Maximum number of node types. Each node type has a corresponding
      learnable scaling factor and a bias. Default : 100.
  
  Examples
  --------
  >>> import torch as th
  >>> import dgl
  >>> from dgl.nn import SpatialEncoder3d
  
  >>> coordinate = th.rand(1, 4, 3)
  >>> node_type = th.tensor([[1, 0, 2, 1]])
  >>> spatial_encoder = SpatialEncoder3d(num_kernels=4,
  ...                                    num_heads=8,
  ...                                    max_node_type=3)
  >>> out = spatial_encoder(coordinate, node_type=node_type)
  >>> print(out.shape)
  torch.Size([1, 4, 4, 8])
class_docstr_tok: ["3D","Spatial","Encoder",",","as","introduced","in","`","One","Transformer","Can","Understand","Both","2D","&","3D","Molecular","Data","<","https",":","\/\/arxiv.org\/pdf\/2210.01765.pdf",">","`","__","This","module","encodes","pair-wise","relation","between","node","pair",":","math",":","`","(","i",",","j",")","`","in","the","3D","geometric","space",",","according","to","the","Gaussian","Basis","Kernel","function",":",":","math",":","`","\\psi","_","{","(","i",",","j",")","}","^k","=","\\frac","{","1","}","{","\\sqrt","{","2\\pi","}","\\lvert","\\sigma^k","\\rvert","}","\\exp","{","\\left","(","-\\frac","{","1","}","{","2","}","\\left","(","\\frac","{","\\gamma_","{","(","i",",","j",")","}","\\lvert","\\lvert","r_i","-","r_j","\\rvert","\\rvert","+","\\beta_","{","(","i",",","j",")","}","-","\\mu^k","}","{","\\lvert","\\sigma^k","\\rvert","}","\\right",")","^2","\\right",")","}","\uff0ck=1",",","...",",","K",",","`","where",":","math",":","`","K","`","is","the","number","of","Gaussian","Basis","kernels",".",":","math",":","`","r_i","`","is","the","Cartesian","coordinate","of","node",":","math",":","`","i","`",".",":","math",":","`","\\gamma_","{","(","i",",","j",")","}",",","\\beta_","{","(","i",",","j",")","}","`","are","learnable","scaling","factors","and","biases","determined","by","node","types",".",":","math",":","`","\\mu^k",",","\\sigma^k","`","are","learnable","centers","and","standard","deviations","of","the","Gaussian","Basis","kernels",".","Parameters","--","--","--","--","--","num_kernels",":","int","Number","of","Gaussian","Basis","Kernels","to","be","applied",".","Each","Gaussian","Basis","Kernel","contains","a","learnable","kernel","center","and","a","learnable","standard","deviation",".","num_heads",":","int",",","optional","Number","of","attention","heads","if","multi-head","attention","mechanism","is","applied",".","Default",":","1.","max_node_type",":","int",",","optional","Maximum","number","of","node","types",".","Each","node","type","has","a","corresponding","learnable","scaling","factor","and","a","bias",".","Default",":","100",".","Examples","--","--","--","--",">",">",">","import","torch","as","th",">",">",">","import","dgl",">",">",">","from","dgl.nn","import","SpatialEncoder3d",">",">",">","coordinate","=","th.rand","(","1",",","4",",","3",")",">",">",">","node_type","=","th.tensor","(","[","[","1",",","0",",","2",",","1","]","]",")",">",">",">","spatial_encoder","=","SpatialEncoder3d","(","num_kernels=4",",","...","num_heads=8",",","...","max_node_type=3",")",">",">",">","out","=","spatial_encoder","(","coordinate",",","node_type=node_type",")",">",">",">","print","(","out.shape",")","torch.Size","(","[","1",",","4",",","4",",","8","]",")"]
docstr: Parameters
  ----------
  coord : torch.Tensor
      3D coordinates of nodes in shape :math:`(B, N, 3)`, where :math:`B`
      is the batch size, :math:`N`: is the maximum number of nodes.
  node_type : torch.Tensor, optional
      Node type ids of nodes. Default : None.
  
      * If specified, :attr:`node_type` should be a tensor in shape
        :math:`(B, N,)`. The scaling factors in gaussian kernels of each
        pair of nodes are determined by their node types.
      * Otherwise, :attr:`node_type` will be set to zeros of the same
        shape by default.
  
  Returns
  -------
  torch.Tensor
      Return attention bias as 3D spatial encoding of shape
      :math:`(B, N, N, H)`, where :math:`H` is :attr:`num_heads`.
docstr_tok: ["Parameters","--","--","--","--","--","coord",":","torch.Tensor","3D","coordinates","of","nodes","in","shape",":","math",":","`","(","B",",","N",",","3",")","`",",","where",":","math",":","`","B","`","is","the","batch","size",",",":","math",":","`","N","`",":","is","the","maximum","number","of","nodes",".","node_type",":","torch.Tensor",",","optional","Node","type","ids","of","nodes",".","Default",":","None",".","*","If","specified",",",":","attr",":","`","node_type","`","should","be","a","tensor","in","shape",":","math",":","`","(","B",",","N",",",")","`",".","The","scaling","factors","in","gaussian","kernels","of","each","pair","of","nodes","are","determined","by","their","node","types",".","*","Otherwise",",",":","attr",":","`","node_type","`","will","be","set","to","zeros","of","the","same","shape","by","default",".","Returns","--","--","--","-","torch.Tensor","Return","attention","bias","as","3D","spatial","encoding","of","shape",":","math",":","`","(","B",",","N",",","N",",","H",")","`",",","where",":","math",":","`","H","`","is",":","attr",":","`","num_heads","`","."]
code: def forward(self, coord, node_type=None):
      """
      Parameters
      ----------
      coord : torch.Tensor
          3D coordinates of nodes in shape :math:`(B, N, 3)`, where :math:`B`
          is the batch size, :math:`N`: is the maximum number of nodes.
      node_type : torch.Tensor, optional
          Node type ids of nodes. Default : None.
  
          * If specified, :attr:`node_type` should be a tensor in shape
            :math:`(B, N,)`. The scaling factors in gaussian kernels of each
            pair of nodes are determined by their node types.
          * Otherwise, :attr:`node_type` will be set to zeros of the same
            shape by default.
  
      Returns
      -------
      torch.Tensor
          Return attention bias as 3D spatial encoding of shape
          :math:`(B, N, N, H)`, where :math:`H` is :attr:`num_heads`.
      """
      bsz, N = coord.shape[:2]
      euc_dist = th.cdist(coord, coord, p=2.0)  # shape: [B, n, n]
      if node_type is None:
          node_type = th.zeros(
              [bsz, N, N, 2], device=coord.device
          ).long()
      else:
          src_node_type = node_type.unsqueeze(-1).repeat(1, 1, N)
          tgt_node_type = node_type.unsqueeze(1).repeat(1, N, 1)
          node_type = th.stack(
              [
                  src_node_type + 2,
                  tgt_node_type + self.max_node_type + 3,
              ],
              dim=-1,
          )  # shape: [B, n, n, 2]
  
      # scaled euclidean distance
      gamma = self.gamma(node_type).sum(dim=-2)  # shape: [B, n, n, 1]
      beta = self.beta(node_type).sum(dim=-2)  # shape: [B, n, n, 1]
      euc_dist = (
          gamma * euc_dist.unsqueeze(-1) + beta
      )  # shape: [B, n, n, 1]
      # gaussian basis kernel
      euc_dist = euc_dist.expand(-1, -1, -1, self.num_kernels)
      gaussian_kernel = gaussian(
          euc_dist, self.means, self.stds.abs() + 1e-2
      )  # shape: [B, n, n, K]
      # linear projection
      encoding = self.linear_layer_1(gaussian_kernel)
      encoding = F.gelu(encoding)
      encoding = self.linear_layer_2(encoding)  # shape: [B, n, n, H]
  
      return 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 SpatialEncoder3d obj's underlying member method code with example usage: obj.forward(coord, node_type) and returns: encoding