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
      """
      # now gat,sage,chebconv may be different from pyg
      if gnn_name == "gat_8":
          return GATConvC(in_dim, out_dim, 8, concat=concat, bias=bias)
      elif gnn_name == "gat_6":
          return GATConvC(in_dim, out_dim, 6, concat=concat, bias=bias)
      elif gnn_name == "gat_4":
          return GATConvC(in_dim, out_dim, 4, concat=concat, bias=bias)
      elif gnn_name == "gat_2":
          return GATConvC(in_dim, out_dim, 2, concat=concat, bias=bias)
      elif gnn_name in ["gat_1", "gat"]:
          return GATConvC(in_dim, out_dim, 1, concat=concat, bias=bias)
      elif gnn_name == "gcn":
          return GraphConv(in_dim, out_dim)
      elif gnn_name == "cheb":
          return MChebConv(in_dim, out_dim, k=2, bias=bias)
      elif gnn_name == "sage":
          return SAGEConv(
              in_dim, out_dim, bias=bias, aggregator_type="pool"
          )
      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(dgl.nn.conv, gnn_name):
          cls = getattr(dgl.nn.conv, 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: import torch
  from torch import nn
  import torch.nn.functional
  import dgl
  from dgl.nn.pytorch.conv import *
  import dgl.function
  import math
  from dgl.utils import expand_as_pair
  from dgl.nn.functional import edge_softmax
  from torch.nn import Parameter
  import torch
  from torch import Tensor
  from ..scatter_utils import *
  from typing import Optional, Tuple, Union
label_desc: Use image node_id 3 for calling a global function with example usage: gnn_map(gnn_name, in_dim, out_dim, concat, bias) and returns: GATConvC, GATConvC, GATConvC, GATConvC, GATConvC, GraphConv, MChebConv, SAGEConv, GatedGraphConv, ARMAConv, SGConv, LinearConv, ZeroConv, Identity, cls