def create_train_state(config: ml_collections.ConfigDict, rng: np.ndarray,
                       input_shape: Sequence[int],
                       num_classes: int) -> Tuple[Any, TrainState]:
  """Create and initialize the model.

  Args:
    config: Configuration for model.
    rng: JAX PRNG Key.
    input_shape: Shape of the inputs fed into the model.
    num_classes: Number of classes in the output layer.

  Returns:
    The initialized TrainState with the optimizer.
  """
  # Create model function.
  if config.model_name.startswith("resnet"):
    model_cls = resnet_v1.create_model(config.model_name, config)
  elif config.model_name.startswith("nest"):
    model_cls = nest_net.create_model(config.model_name, config)
  else:
    raise ValueError(f"Model {config.model_name} not supported.")
  model = functools.partial(model_cls, num_classes=num_classes)
  variables = model(train=False).init(rng, jnp.ones(input_shape))
  model_state = dict(variables)
  params = model_state.pop("params")
  parameter_overview.log_parameter_overview(params)
  if config.get("log_model_profile"):  # Be True or [1, 2]
    message_1 = utils.log_throughput(model, variables, input_shape)
    message_2 = utils.compute_flops(model, variables,
                                    [1] + list(input_shape[1:]))
    count = parameter_overview.count_parameters(params)
    message_3 = "Params: {:,}".format(count)
    message = ", ".join([message_1, message_2, message_3])
    logging.info("Profile results %s", message)
    if (isinstance(config.log_model_profile, (int,)) and
        config.log_model_profile >= 2):
      sys.exit(0)
  # Create optimizer.
  if config.optim in ("adamw", "adam"):
    if config.get("optim_wd_ignore"):
      # Allow zero weight decay for certain parameters listed in optim_wd_ignore
      igns = config.optim_wd_ignore
      p = flax.optim.ModelParamTraversal(
          lambda path, _: not any([i in path for i in igns]))
      p_nowd = flax.optim.ModelParamTraversal(
          lambda path, _: any([i in path for i in igns]))
      p_opt = flax.optim.Adam(weight_decay=config.weight_decay)
      p_nowd_opt = flax.optim.Adam(weight_decay=0)
      optimizer = flax.optim.MultiOptimizer((p, p_opt),
                                            (p_nowd, p_nowd_opt)).create(params)
    else:
      optimizer = flax.optim.Adam(
          weight_decay=config.weight_decay).create(params)
  elif config.optim == "sgd":
    optimizer = flax.optim.Momentum(
        beta=config.sgd_momentum, nesterov=True).create(params)
  else:
    raise NotImplementedError(f"{config.optim} does not exist.")

  return model, TrainState(step=0, optimizer=optimizer, model_state=model_state)