class_docstr: A generic message-passing graph neural network with discrete edge types.
class_docstr_tok: ["A","generic","message-passing","graph","neural","network","with","discrete","edge","types","."]
docstr: :param node_representations: A [num_nodes, hidden_dimension] matrix of node representations.
  :param adjacency_lists: a list of [num_edges_per_type, 2] adjacency lists per edge type.
          The order is fixed across runs. Backwards edges and self-edges are included if
          the appropriate hyperparameter is set.
  :param edge_feature_embeddings: a list of the edge features per edge-type.
  :param node_to_graph_idx: A mapping that tells us which graph the node belongs to
  :param reference_node_ids: A dictionary indicating the reference node index
  :param reference_node_graph_idx: A dictionary indicating the graph index for reference node
  :param return_all_states: Whether to return all states
  :return: a [num_nodes, output_hidden_dimension] matrix of the output representations
docstr_tok: [":","param","node_representations",":","A","[","num_nodes",",","hidden_dimension","]","matrix","of","node","representations",".",":","param","adjacency_lists",":","a","list","of","[","num_edges_per_type",",","2","]","adjacency","lists","per","edge","type",".","The","order","is","fixed","across","runs",".","Backwards","edges","and","self-edges","are","included","if","the","appropriate","hyperparameter","is","set",".",":","param","edge_feature_embeddings",":","a","list","of","the","edge","features","per","edge-type",".",":","param","node_to_graph_idx",":","A","mapping","that","tells","us","which","graph","the","node","belongs","to",":","param","reference_node_ids",":","A","dictionary","indicating","the","reference","node","index",":","param","reference_node_graph_idx",":","A","dictionary","indicating","the","graph","index","for","reference","node",":","param","return_all_states",":","Whether","to","return","all","states",":","return",":","a","[","num_nodes",",","output_hidden_dimension","]","matrix","of","the","output","representations"]
code: def gnn(
      self,
      node_representations: torch.Tensor,
      adjacency_lists: List[Tuple[torch.Tensor, torch.Tensor]],
      edge_feature_embeddings: List[torch.Tensor],
      node_to_graph_idx: torch.Tensor,
      reference_node_ids: Dict[str, torch.Tensor],
      reference_node_graph_idx: Dict[str, torch.Tensor],
      return_all_states: bool = False,
  ) -> torch.Tensor:
      """
      :param node_representations: A [num_nodes, hidden_dimension] matrix of node representations.
      :param adjacency_lists: a list of [num_edges_per_type, 2] adjacency lists per edge type.
              The order is fixed across runs. Backwards edges and self-edges are included if
              the appropriate hyperparameter is set.
      :param edge_feature_embeddings: a list of the edge features per edge-type.
      :param node_to_graph_idx: A mapping that tells us which graph the node belongs to
      :param reference_node_ids: A dictionary indicating the reference node index
      :param reference_node_graph_idx: A dictionary indicating the graph index for reference node
      :param return_all_states: Whether to return all states
      :return: a [num_nodes, output_hidden_dimension] matrix of the output representations
      """
      if self.__edge_dropout_rate > 0 and self.training:
          dropped_adj_list, dropped_edge_features = [], []
          for (
              edge_sources_idxs,
              edge_target_idxs,
          ), edge_features in zip(
              adjacency_lists, edge_feature_embeddings
          ):
              mask = (
                  torch.rand_like(
                      edge_sources_idxs, dtype=torch.float32
                  )
                  > self.__edge_dropout_rate
              )
              dropped_adj_list.append(
                  (
                      edge_sources_idxs.masked_select(mask),
                      edge_target_idxs.masked_select(mask),
                  )
              )
              dropped_edge_features.append(edge_features[mask])
          adjacency_lists = dropped_adj_list
          edge_feature_embeddings = dropped_edge_features
  
      all_states = [node_representations]
      for mp_layer in self.__message_passing_layers:
          node_representations = mp_layer(
              node_states=node_representations,
              adjacency_lists=adjacency_lists,
              node_to_graph_idx=node_to_graph_idx,
              reference_node_ids=reference_node_ids,
              reference_node_graph_idx=reference_node_graph_idx,
              edge_features=edge_feature_embeddings,
          )
          all_states.append(node_representations)
      if return_all_states:
          node_representations = torch.cat(all_states, dim=-1)
      return node_representations
  
imports: from typing_extensions import Final
  import logging
  import numpy
  import torch
  from collections import defaultdict
  from torch import nn
  from typing import Any, Callable, Dict, Iterable, Iterator, List, Optional, Set, Tuple, TypeVar, Union
  from ptgnn.baseneuralmodel import AbstractNeuralModel, ModuleWithMetrics
  from ptgnn.baseneuralmodel.utils.data import enforce_not_None
  from ptgnn.neuralmodels.gnn.messagepassing.abstractmessagepassing import AbstractMessagePassingLayer
  from ptgnn.neuralmodels.gnn.structs import GnnOutput, GraphData, TensorizedGraphData
label_desc: Use image node_id 7 for calling the GraphNeuralNetwork obj's underlying member method code with example usage: obj.gnn(node_representations, adjacency_lists, edge_feature_embeddings, node_to_graph_idx, reference_node_ids, reference_node_graph_idx, return_all_states) and returns: node_representations