def get_model(model_name):
  """Get the corresponding model class based on the model string.

  API:
  model_builder, hps = get_model("fully_connected")
  ... modify/parse hparams
  model = model_builder(hps, num_classes)

  Args:
    model_name: (str) e.g. fully_connected.

  Returns:
    The model architecture (currently a flax Model) along with its
    default hparams.
  Raises:
    ValueError if model is unrecognized.
  """
  try:
    return _ALL_MODELS[model_name][0]
  except KeyError:
    raise ValueError('Unrecognized model: {}'.format(model_name)) from None