class_docstr: EGTLayer for Edge-augmented Graph Transformer (EGT), as introduced in
  `Global Self-Attention as a Replacement for Graph Convolution
  Reference `<https://arxiv.org/pdf/2108.03348.pdf>`_
  
  Parameters
  ----------
  feat_size : int
      Node feature size.
  edge_feat_size : int
      Edge feature size.
  num_heads : int
      Number of attention heads, by which :attr: `feat_size` is divisible.
  num_virtual_nodes : int
      Number of virtual nodes.
  dropout : float, optional
      Dropout probability. Default: 0.0.
  attn_dropout : float, optional
      Attention dropout probability. Default: 0.0.
  activation : callable activation layer, optional
      Activation function. Default: nn.ELU().
  edge_update : bool, optional
      Whether to update the edge embedding. Default: True.
  
  Examples
  --------
  >>> import torch as th
  >>> from dgl.nn import EGTLayer
  
  >>> batch_size = 16
  >>> num_nodes = 100
  >>> feat_size, edge_feat_size = 128, 32
  >>> nfeat = th.rand(batch_size, num_nodes, feat_size)
  >>> efeat = th.rand(batch_size, num_nodes, num_nodes, edge_feat_size)
  >>> net = EGTLayer(
          feat_size=feat_size,
          edge_feat_size=edge_feat_size,
          num_heads=8,
          num_virtual_nodes=4,
      )
  >>> out = net(nfeat, efeat)
code: def __init__(
      self,
      feat_size,
      edge_feat_size,
      num_heads,
      num_virtual_nodes,
      dropout=0,
      attn_dropout=0,
      activation=nn.ELU(),
      edge_update=True,
  ):
      super().__init__()
      self.num_heads = num_heads
      self.num_virtual_nodes = num_virtual_nodes
      self.edge_update = edge_update
  
      assert (
          feat_size % num_heads == 0
      ), "feat_size must be divisible by num_heads"
      self.dot_dim = feat_size // num_heads
      self.mha_ln_h = nn.LayerNorm(feat_size)
      self.mha_ln_e = nn.LayerNorm(edge_feat_size)
      self.edge_input = nn.Linear(edge_feat_size, num_heads)
      self.qkv_proj = nn.Linear(feat_size, feat_size * 3)
      self.gate = nn.Linear(edge_feat_size, num_heads)
      self.attn_dropout = nn.Dropout(attn_dropout)
      self.node_output = nn.Linear(feat_size, feat_size)
      self.mha_dropout_h = nn.Dropout(dropout)
  
      self.node_ffn = nn.Sequential(
          nn.LayerNorm(feat_size),
          nn.Linear(feat_size, feat_size),
          activation,
          nn.Linear(feat_size, feat_size),
          nn.Dropout(dropout),
      )
  
      if self.edge_update:
          self.edge_output = nn.Linear(num_heads, edge_feat_size)
          self.mha_dropout_e = nn.Dropout(dropout)
          self.edge_ffn = nn.Sequential(
              nn.LayerNorm(edge_feat_size),
              nn.Linear(edge_feat_size, edge_feat_size),
              activation,
              nn.Linear(edge_feat_size, edge_feat_size),
              nn.Dropout(dropout),
          )
  
imports: import torch
  import torch.nn
  import torch.nn.functional
label_desc: Use image node_id 1 to create a new EGTLayer object from inherited base classes: nn with example: obj = EGTLayer(feat_size, edge_feat_size, num_heads, num_virtual_nodes, dropout, attn_dropout, activation, edge_update)