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:
      assert len(adjacency_lists) == len(
          self.__edge_message_transformation_layers
      ), "The number of adjacency lists must be equal to the number of edge types."
  
      all_message_targets, 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, edge_target_idxs = adj_list
          all_message_targets.append(edge_target_idxs)
  
          edge_source_states = nn.functional.embedding(
              edge_sources_idxs, node_states
          )
  
          if self.__use_target_state_as_message_input:
              edge_target_states = nn.functional.embedding(
                  edge_target_idxs, node_states
              )
              message_input = torch.cat(
                  [edge_source_states, edge_target_states], dim=-1
              )
          else:
              message_input = edge_source_states
  
          all_messages.append(
              edge_transformation_layer(
                  torch.cat([message_input, features], dim=-1)
              )
          )
  
      if isinstance(self.__aggregation_fn, AbstractMessageAggregation):
          aggregated_messages = self.__aggregation_fn(
              messages=torch.cat(all_messages, dim=0),
              message_targets=torch.cat(all_message_targets, dim=0),
              num_nodes=node_states.shape[0],
          )
      else:
          aggregated_messages = self._aggregate_messages(
              messages=torch.cat(all_messages, dim=0),
              message_targets=torch.cat(all_message_targets, dim=0),
              num_nodes=node_states.shape[0],
              aggregation_fn=self.__aggregation_fn,
          )
  
      if self.__message_activation is not None:
          aggregated_messages = self.__message_activation(
              aggregated_messages
          )
  
      return self.__state_update(aggregated_messages)
  
file_path: ptgnn/ptgnn/neuralmodels/gnn/messagepassing/mlpmessagepassing.py
imports: import torch
  from torch import nn
  from typing import Callable, Dict, List, Optional, Tuple, Union
  from ptgnn.neuralmodels.gnn.messagepassing.abstractmessagepassing import AbstractMessageAggregation, AbstractMessagePassingLayer
  from ptgnn.neuralmodels.mlp import MLP
label_desc: Use image node_id 2 for calling the MlpMessagePassingLayer 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