text: 
  ## practice adoption from https://github.com/gordicaleksa/pytorch-GAT
  
  import torch
  import torch.nn as nn
  
  class GATLayer(nn.Module):
  
      head_dim = 1
  
      def __init__(self,
                  num_in_features,
                  num_out_features,
                  num_of_heads,
                  concat=True,
                  activation=nn.ELU(),
                  dropout_prob=0.6,
                  add_skip_connection=True,
                  bias=True,
                  log_attention_weights=False):
  
          super().__init__()
  
          self.num_of_heads = num_of_heads
          self.num_out_features = num_out_features
          self.concat = concat
          self.add_skip_connection = add_skip_connection
  
          self.linear_proj = nn.Linear(num_in_features, num_of_heads * num_out_features, bias=False)
  
          self.scoring_fn_target = nn.Parameter(torch.tensor(1, num_of_heads, num_out_features))
          self.scoring_fn_source = nn.Parameter(torch.tensor(1, num_of_heads, num_out_features))
  
          if bias and concat:
              self.bias = nn.Parameter(torch.Tensor(num_of_heads * num_out_features))
          elif bias and not cat:
              self.bias = nn.Parameter(torch.Tensor(num_out_features))
          else:
              self.register_parameter('bias', None)
  
          if add_skip_connection:
              self.skip_proj = nn.Linear(num_in_features, num_of_heads * num_out_features, bias=False)
          else:
              self.register_parameter('skip_proj', None)
  
          self.leakyReLu = nn.leakyReLu(0.2)
          self.softmax = nn.Softmax(dim=-1)
          self.activation = activation
          self.dropout = nn.Dropout(p=dropout_prob)
  
          self.log_attention_weights = log_attention_weights
          self.attention_weights = None
  
          self.init_params()
  
  class GATLayer2(GATLayer):
  
      head_dim = 1
  
      def __init__(self,
                  num_in_features,
                  num_out_features,
                  num_of_heads,
                  concat=True,
                  activation=nn.ELU(),
                  dropout_prob=0.6,
                  add_skip_connection=True,
                  bias=True,
                  log_attention_weights=False):
  
          super().__init__(num_in_features, num_out_features, num_of_heads, concat, activation, dropout_prob,
                              add_skip_connection, bias, log_attention_weights)
  
      def forward(self, data):
  
          in_node_features, connectivity_mask = data
          num_of_nodes = in_node_features.shape[0]
  
          assert connectivity_mask == (num_of_nodes, num_of_nodes), \
              f"Expected shape=({num_of_nodes},{num_of_nodes}), got shape={connectivity_mask.shape}"
  
          ## shape = (N, FIN); N=number of nodes in the graph, FIN=number of input features per node
          in_nodes_features = self.dropout(in_nodes_features)
  
          ## shape = (N, FIN) * (FIN, NH*FOUT) -> (N, NH, FOUT); NH=number of heads, FOUT=number of output features
          nodes_features_proj = self.linear_proj(in_node_features).view(-1, self.num_of_heads, self.num_out_features)
          nodes_features_proj = self.dropout(nodes_features_proj)
  
          ## shape = (N, FIN, FOUT) * (1, NH, FOUT) -> (NH, 1, N)
          scores_source = torch.sum((nodes_features_proj * self.scoring_fn_source), dim=-1, keepdim=True)
          scores_target = torch.sum((nodes_features_proj * self.scoring_fn_target), dim=-1, keepdim=True)
  
          ## src shape = (NH, N, 1) and trg shape = (NH, 1, N)
          scores_source = scores_source.transpose(0, 1)
          scores_target = scores_target.permute(1, 2, 0)
  
  
      def init_params(self):
          nn.init.xavier_uniform_(self.linear_proj.weight)
          nn.init.xavier_uniform_(self.scoring_fn_target)
          nn.init.xavier_uniform_(self.scoring_fn_source)
  
          if self.bias is not None:
              torch.nn.init.zeros_(self.bias)
  