def create_train_state(model, config, rng, inputs):
  """Create and initialize the model.

  Args:
    model: Flax nn module.
    config: Configuration for model.
    rng: JAX PRNG Key.
    inputs: The init inputs fed into the model.

  Returns:
    The initialized TrainState with the optimizer.
  """
  rng, params_rng, dropout_rng = jax.random.split(rng, 3)
  variables = model.init({'params': params_rng, 'dropout': dropout_rng}, inputs)
  params = variables['params']
  parameter_overview.log_parameter_overview(params)  # pytype: disable=wrong-arg-types

  if config.optimizer == 'AdamW':

    def w_decay_fn(path):
      return all(pn not in path for pn in config.no_weight_decay)

    def wo_decay_fn(path):
      return any(pn in path for pn in config.no_weight_decay)

    optimizer_w_decay = flax.optim.Adam(
        learning_rate=config.learning_rate,
        weight_decay=config.learning_rate_weight_decay)
    optimizer_wo_decay = flax.optim.Adam(
        learning_rate=config.learning_rate, weight_decay=0)
    params_w_decay = flax.optim.ModelParamTraversal(
        lambda path, _: w_decay_fn(path))
    params_wo_decay = flax.optim.ModelParamTraversal(
        lambda path, _: wo_decay_fn(path))
    optimizer = flax.optim.MultiOptimizer(
        (params_w_decay, optimizer_w_decay),
        (params_wo_decay, optimizer_wo_decay)).create(params)
  else:
    raise NotImplementedError
  return TrainState(step=0, optimizer=optimizer)
Create and initialize the model.

  Args:
    model: Flax nn module.
    config: Configuration for model.
    rng: JAX PRNG Key.
    inputs: The init inputs fed into the model.

  Returns:
    The initialized TrainState with the optimizer.
  
language/gscan/xattn_model/train_utils.py