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","]",")"]
code: def __init__(self, num_kernels, num_heads=1, max_node_type=100):
      super().__init__()
      self.num_kernels = num_kernels
      self.num_heads = num_heads
      self.max_node_type = max_node_type
      self.means = nn.Parameter(th.empty(num_kernels))
      self.stds = nn.Parameter(th.empty(num_kernels))
      self.linear_layer_1 = nn.Linear(num_kernels, num_kernels)
      self.linear_layer_2 = nn.Linear(num_kernels, num_heads)
      # There are 2 * max_node_type + 3 pairs of gamma and beta parameters:
      # 1. Parameters at position 0 are for default gamma/beta when no node
      #    type is given
      # 2. Parameters at position 1 to max_node_type+1 are for src node types.
      #    (position 1 is for padded unexisting nodes)
      # 3. Parameters at position max_node_type+2 to 2*max_node_type+2 are
      #    for tgt node types. (position max_node_type+2 is for padded)
      #    unexisting nodes)
      self.gamma = nn.Embedding(2 * max_node_type + 3, 1, padding_idx=0)
      self.beta = nn.Embedding(2 * max_node_type + 3, 1, padding_idx=0)
  
      nn.init.uniform_(self.means, 0, 3)
      nn.init.uniform_(self.stds, 0, 3)
      nn.init.constant_(self.gamma.weight, 1)
      nn.init.constant_(self.beta.weight, 0)
  
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 1 to create a new SpatialEncoder3d object from inherited base classes: nn with example: obj = SpatialEncoder3d(num_kernels, num_heads, max_node_type)