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_data: The data for the node embedder to compute the initial node representations.
  :param adjacency_lists: A list of [num_edges, 2] matrices for each edge type.
  :param edge_feature_data: A list of the same size as `adjacency_lists` with the data
      to compute the edge features.
  :param node_to_graph_idx: A [num_nodes] vector that contains the index of the graph it belongs to.
  :param reference_node_ids: A dictionary with values the indices of the reference nodes.
  :param reference_node_graph_idx: A dictionary with values the index of the graph that each
      node in the `reference_node_ids` belongs to.
  :param num_graphs: the number of graphs.
docstr_tok: [":","param","node_data",":","The","data","for","the","node","embedder","to","compute","the","initial","node","representations",".",":","param","adjacency_lists",":","A","list","of","[","num_edges",",","2","]","matrices","for","each","edge","type",".",":","param","edge_feature_data",":","A","list","of","the","same","size","as","`","adjacency_lists","`","with","the","data","to","compute","the","edge","features",".",":","param","node_to_graph_idx",":","A","[","num_nodes","]","vector","that","contains","the","index","of","the","graph","it","belongs","to",".",":","param","reference_node_ids",":","A","dictionary","with","values","the","indices","of","the","reference","nodes",".",":","param","reference_node_graph_idx",":","A","dictionary","with","values","the","index","of","the","graph","that","each","node","in","the","`","reference_node_ids","`","belongs","to",".",":","param","num_graphs",":","the","number","of","graphs","."]
code: def forward(
      self,
      *,
      node_data,
      adjacency_lists: List[Tuple[torch.Tensor, torch.Tensor]],
      edge_feature_data: List,
      node_to_graph_idx: torch.Tensor,
      reference_node_ids: Dict[str, torch.Tensor],
      reference_node_graph_idx: Dict[str, torch.Tensor],
      num_graphs,
      **kwargs,
  ) -> GnnOutput:
      """
  
      :param node_data: The data for the node embedder to compute the initial node representations.
      :param adjacency_lists: A list of [num_edges, 2] matrices for each edge type.
      :param edge_feature_data: A list of the same size as `adjacency_lists` with the data
          to compute the edge features.
      :param node_to_graph_idx: A [num_nodes] vector that contains the index of the graph it belongs to.
      :param reference_node_ids: A dictionary with values the indices of the reference nodes.
      :param reference_node_graph_idx: A dictionary with values the index of the graph that each
          node in the `reference_node_ids` belongs to.
      :param num_graphs: the number of graphs.
      """
      initial_node_representations = self.__node_embedder(
          **node_data
      )  # [num_nodes, D]
  
      if self.__edge_feature_embedder is None:
          edge_feature_embeddings = [
              torch.empty(
                  f.shape[0], 0, device=node_to_graph_idx.device
              )
              for f, _ in adjacency_lists
          ]
      else:
          edge_feature_embeddings = [
              self.__edge_feature_embedder(**edge_data)
              for edge_data in edge_feature_data
          ]
  
      if self.__introduce_backwards_edges:
          adjacency_lists += [(t, f) for f, t in adjacency_lists]
          edge_feature_embeddings += [
              e for e in edge_feature_embeddings
          ]
  
      if self.__add_self_edges:
          num_nodes = node_to_graph_idx.shape[0]
          idents = torch.arange(
              num_nodes,
              dtype=torch.int64,
              device=node_to_graph_idx.device,
          )
          adjacency_lists.append((idents, idents))
          edge_feature_embeddings.append(
              torch.zeros(
                  num_nodes,
                  edge_feature_embeddings[-1].shape[-1],
                  device=node_to_graph_idx.device,
              )
          )
  
      output_representations = self.gnn(
          initial_node_representations,
          adjacency_lists,
          edge_feature_embeddings,
          node_to_graph_idx,
          reference_node_ids,
          reference_node_graph_idx,
          **kwargs,
      )  # [num_nodes, H]
  
      with torch.no_grad():
          self.__num_edges += sum(
              adj[0].shape[0] for adj in adjacency_lists
          )
          self.__num_graphs += num_graphs
          self.__num_nodes += node_to_graph_idx.shape[0]
      return GnnOutput(
          input_node_representations=initial_node_representations,
          output_node_representations=output_representations,
          node_to_graph_idx=node_to_graph_idx,
          node_idx_references=reference_node_ids,
          node_graph_idx_reference=reference_node_graph_idx,
          num_graphs=num_graphs,
      )
  
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 8 for calling the GraphNeuralNetwork obj's underlying member method code with example usage: obj.forward() and returns: GnnOutput