code: def get_model_cls(
      model_name,
      data,
      dataset_name,
      hidden=None,
      dropout=None,
      **unused_kwargs,
  ):
      """Picks the models depending on the provided configuration flags."""
      # Create model classification.
      if model_name == "mlp":
          hidden = (
              parse_layers_string(hidden) if hidden is not None else ()
          )
          return MLP(
              output_dim=data.num_classes,
              hidden_sizes=hidden,
              activation=tf.nn.leaky_relu,
              name="mlp_cls",
          )
      elif model_name == "cnn":
          if dataset_name in ("mnist", "fashion_mnist"):
              channels = 1
          elif dataset_name in (
              "cifar10",
              "cifar100",
              "svhn_cropped",
              "svhn",
          ):
              channels = 3
          else:
              raise ValueError(
                  "Dataset name `%s` unsupported." % dataset_name
              )
          return ImageCNNAgreement(
              output_dim=data.num_classes,
              channels=channels,
              activation=tf.nn.leaky_relu,
              name="cnn_cls",
          )
      elif model_name == "wide_resnet":
          return WideResnet(
              num_classes=data.num_classes,
              lrelu_leakiness=0.1,
              horizontal_flip=dataset_name in ("cifar10",),
              random_translation=False,
              gaussian_noise=dataset_name
              not in ("svhn", "svhn_cropped"),
              width=2,
              num_residual_units=4,
              name="wide_resnet_cls",
          )
      elif model_name == "gcn":
          hidden = (
              parse_layers_string(hidden) if hidden is not None else ()
          )
          assert (
              len(hidden) == 1
          ), "GCN implementation currently supports only one hidden layer."
          return GCN(
              input_dim=data.num_features,
              output_dim=data.num_classes,
              hidden=hidden[0],
              dropout=dropout,
              aggregation=None,
              hidden_aggregation=(),
              activation=tf.nn.leaky_relu,
              is_binary_classification=False,
              name="gcn_cls",
          )
      else:
          raise NotImplementedError()
  
file_path: neural-structured-learning/research/gam/gam/experiments/helper.py
imports: from ..models.cnn import ImageCNNAgreement
  from ..models.gcn import GCN
  from ..models.mlp import MLP
  from ..models.wide_resnet import WideResnet
  import tensorflow
label_desc: Use image node_id 2 for calling a global function with example usage: get_model_cls(model_name, data, dataset_name, hidden, dropout) and returns: MLP, ImageCNNAgreement, WideResnet, GCN