class_docstr: Head of GNN, edge prediction
class_docstr_tok: ["Head","of","GNN",",","edge","prediction"]
docstr: Head of Edge and link prediction models.
  
  Args:
      dim_out: output dimension. For binary prediction, dim_out=1.
docstr_tok: ["Head","of","Edge","and","link","prediction","models",".","Args",":","dim_out",":","output","dimension",".","For","binary","prediction",",","dim_out=1","."]
code: def __init__(self, dim_in, dim_out):
      """Head of Edge and link prediction models.
  
      Args:
          dim_out: output dimension. For binary prediction, dim_out=1.
      """
      # Use dim_in for graph conv, since link prediction dim_out could be
      # binary
      # E.g. if decoder='dot', link probability is dot product between
      # node embeddings, of dimension dim_in
      super(GNNEdgeHead, self).__init__()
      # module to decode edges from node embeddings
  
      if cfg.model.edge_decoding == "concat":
          self.layer_post_mp = MLP(
              dim_in * 2,
              dim_out,
              num_layers=cfg.gnn.layers_post_mp,
              bias=True,
          )
          # requires parameter
          self.decode_module = lambda v1, v2: self.layer_post_mp(
              torch.cat((v1, v2), dim=-1)
          )
      else:
          if dim_out > 1:
              raise ValueError(
                  "Binary edge decoding ({})is used for multi-class "
                  "edge/link prediction.".format(
                      cfg.model.edge_decoding
                  )
              )
          self.layer_post_mp = MLP(
              dim_in,
              dim_in,
              num_layers=cfg.gnn.layers_post_mp,
              bias=True,
          )
          if cfg.model.edge_decoding == "dot":
              self.decode_module = lambda v1, v2: torch.sum(
                  v1 * v2, dim=-1
              )
          elif cfg.model.edge_decoding == "cosine_similarity":
              self.decode_module = nn.CosineSimilarity(dim=-1)
          else:
              raise ValueError(
                  "Unknown edge decoding {}.".format(
                      cfg.model.edge_decoding
                  )
              )
  
imports: import torch
  import torch.nn
  import graphgym.register
  from graphgym.config import cfg
  from graphgym.models.layer_pyg import MLP
  from graphgym.models.pooling import pooling_dict
label_desc: Use image node_id 1 to create a new GNNEdgeHead object from inherited base classes: nn with example: obj = GNNEdgeHead(dim_in, dim_out)