class_docstr: A transformer layer among all nodes in a graph.
class_docstr_tok: ["A","transformer","layer","among","all","nodes","in","a","graph","."]
code: def forward(
      self,
      node_states: torch.Tensor,
      adjacency_lists: List[Tuple[torch.Tensor, torch.Tensor]],
      node_to_graph_idx: torch.Tensor,
      reference_node_ids: Dict[str, torch.Tensor],
      reference_node_graph_idx: Dict[str, torch.Tensor],
      edge_features: List[torch.Tensor],  # Not used
  ) -> torch.Tensor:
      if self.__target_reference == "all":
          relevant_node_states = node_states
      else:
          relevant_node_states = node_states[
              reference_node_ids[self.__target_reference]
          ]
          node_to_graph_idx = reference_node_graph_idx[
              self.__target_reference
          ]
  
      keys_queries_values = self.__selfatt_head_transforms(
          relevant_node_states
      ).reshape(relevant_node_states.shape[0], self.__num_heads, -1)
  
      keys = keys_queries_values[
          :, :, : self.__key_query_dim
      ]  # [num_nodes, num_heads, key_dim]
      queries = keys_queries_values[
          :, :, self.__key_query_dim : 2 * self.__key_query_dim
      ]  # [num_nodes, num_heads, key_dim]
      values = keys_queries_values[
          :, :, 2 * self.__key_query_dim :
      ]  # [num_nodes, num_heads, value_dim]
  
      all_out_values = []
      for graph_nodes in self.__iter_idxs_per_graph(node_to_graph_idx):
          # Loop is necessary due to limited memory and ease of coding.
          graph_keys = keys[graph_nodes]
          graph_queries = queries[graph_nodes]
          scores = torch.einsum(
              "khd,vhd->khv", graph_keys, graph_queries
          ) / (self.__key_query_dim**0.5)
          attention_probs = nn.functional.softmax(scores, dim=-1)
          attention_probs = self.__dropout_layer(attention_probs)
          out_values = torch.einsum(
              "khv,vhd->khd", attention_probs, values[graph_nodes]
          )
          all_out_values.append(out_values)
  
      values = torch.cat(all_out_values, dim=0)
      output = self.__summarization_layer(
          values.reshape(values.shape[0], -1)
      )
      attention_output = self.__layer_norm1(
          self.__dropout_layer(output) + node_states
      )
  
      intermediate = nn.functional.relu(
          self.__intermediate_layer(attention_output)
      )
      output = self.__dropout_layer(self.__output_layer(intermediate))
      output_node_states = self.__layer_norm2(output + attention_output)
      if self.__target_reference == "all":
          return output_node_states
      else:
          node_states[
              reference_node_ids[self.__target_reference]
          ] = output_node_states
          return node_states
  
file_path: ptgnn/ptgnn/neuralmodels/gnn/messagepassing/selfattmessagepassing.py
imports: import torch
  from torch import nn
  from torch_scatter import scatter_sum
  from typing import Dict, List, Tuple
  from ptgnn.neuralmodels.gnn.messagepassing.abstractmessagepassing import AbstractMessagePassingLayer
label_desc: Use image node_id 3 for calling the MultiHeadSelfAttentionMessagePassing obj's underlying member method code with example usage: obj.forward(node_states, adjacency_lists, node_to_graph_idx, reference_node_ids, reference_node_graph_idx, edge_features) and returns: output_node_states, node_states