code: def gnn_map(
      gnn_name, in_dim, out_dim, concat=False, bias=True
  ) -> nn.Module:
      """
  
      :param gnn_name:
      :param in_dim:
      :param out_dim:
      :param concat: for gat, concat multi-head output or not
      :return: GNN model
      """
      if gnn_name == "gat_8":
          return GATConv(in_dim, out_dim, 8, concat=concat, bias=bias)
      elif gnn_name == "gat_6":
          return GATConv(in_dim, out_dim, 6, concat=concat, bias=bias)
      elif gnn_name == "gat_4":
          return GATConv(in_dim, out_dim, 4, concat=concat, bias=bias)
      elif gnn_name == "gat_2":
          return GATConv(in_dim, out_dim, 2, concat=concat, bias=bias)
      elif gnn_name in ["gat_1", "gat"]:
          return GATConv(in_dim, out_dim, 1, concat=concat, bias=bias)
      elif gnn_name == "gcn":
          return GCNConv(in_dim, out_dim)
      elif gnn_name == "cheb":
          return ChebConv(in_dim, out_dim, K=2, bias=bias)
      elif gnn_name == "sage":
          return SAGEConv(in_dim, out_dim, bias=bias)
      elif gnn_name == "gated":
          return GatedGraphConv(in_dim, out_dim, bias=bias)
      elif gnn_name == "arma":
          return ARMAConv(in_dim, out_dim, bias=bias)
      elif gnn_name == "sg":
          return SGConv(in_dim, out_dim, bias=bias)
      elif gnn_name == "linear":
          return LinearConv(in_dim, out_dim, bias=bias)
      elif gnn_name == "zero":
          return ZeroConv()
      elif gnn_name == "identity":
          return Identity()
      elif hasattr(torch_geometric.nn, gnn_name):
          cls = getattr(torch_geometric.nn, gnn_name)
          assert isinstance(
              cls, type
          ), "Only support modules, get %s" % (gnn_name)
          kwargs = {
              "in_channels": in_dim,
              "out_channels": out_dim,
              "concat": concat,
              "bias": bias,
          }
          kwargs = {
              key: kwargs[key]
              for key in cls.__init__.__code__.co_varnames
              if key in kwargs
          }
          return cls(**kwargs)
      raise KeyError("Cannot parse key %s" % (gnn_name))
  
imports: from torch_geometric.nn import GATConv, GCNConv, ChebConv, SAGEConv, GatedGraphConv, ARMAConv, SGConv
  import torch_geometric.nn
  import torch
  from torch import nn
  import torch.nn.functional
  import inspect
  import sys
  from torch_scatter import scatter_add
  import torch_scatter
  from torch.nn import Parameter
  from torch_geometric.nn.inits import glorot, zeros
  from torch_geometric.utils import remove_self_loops, add_self_loops, add_remaining_self_loops, softmax
label_desc: Use image node_id 1 for calling a global function with example usage: gnn_map(gnn_name, in_dim, out_dim, concat, bias) and returns: GATConv, GATConv, GATConv, GATConv, GATConv, GCNConv, ChebConv, SAGEConv, GatedGraphConv, ARMAConv, SGConv, LinearConv, ZeroConv, Identity, cls