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],
  ) -> torch.Tensor:
      message_targets = torch.cat(
          [adj_list[1] for adj_list in adjacency_lists]
      )  # num_messages
      assert len(adjacency_lists) == len(
          self.__edge_message_transformation_layers
      )
  
      all_messages = []
      for edge_type_idx, (
          adj_list,
          features,
          edge_transformation_layer,
      ) in enumerate(
          zip(
              adjacency_lists,
              edge_features,
              self.__edge_message_transformation_layers,
          )
      ):
          edge_sources_idxs = adj_list[0]
          edge_source_states = nn.functional.embedding(
              edge_sources_idxs, node_states
          )  # [num_edges_of_type_edge_type_idx, H]
          all_messages.append(
              edge_transformation_layer(
                  self.__dropout(
                      torch.cat([edge_source_states, features], -1)
                  )
              )
          )  # [num_edges_of_type_edge_type_idx, D]
  
      aggregated_messages = self._aggregate_messages(
          messages=torch.cat(all_messages, dim=0),
          message_targets=message_targets,
          num_nodes=node_states.shape[0],
          aggregation_fn=self.__aggregation_fn,
      )  # [num_nodes, D]
      return self.__state_update(aggregated_messages, node_states)
  
file_path: ptgnn/ptgnn/neuralmodels/gnn/messagepassing/gatedmessagepassing.py
imports: import torch
  from torch import nn
  from typing import Dict, List, Tuple
  from ptgnn.neuralmodels.gnn.messagepassing.abstractmessagepassing import AbstractMessagePassingLayer
label_desc: Use image node_id 2 for calling the GatedMessagePassingLayer 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: self