def get_optimizer(
    learning_rate_fn: Callable[[int], float],
    config: ml_collections.ConfigDict) -> optax.GradientTransformation:
  """Constructs the optimizer from the given HParams."""
  weight_decay = config.optimizer.get('weight_decay', 0.0)
  decay_list = ['kernel']
  no_decay_list = ['bias', 'scale', 'embedding', 'embed_pos']

  decay_param_mask = train_utils.flattened_traversal(
      lambda path, _: any(pn in path for pn in decay_list))
  no_decay_param_mask = train_utils.flattened_traversal(
      lambda path, _: any(pn in path for pn in no_decay_list))
  # using inject_hyperparams gets learning_rate but needs manual lr management.
  tx1 = optax.adamw(  # pytype: disable=wrong-arg-types  # numpy-scalars
      learning_rate=learning_rate_fn,
      b1=config.optimizer.beta1,
      b2=config.optimizer.beta2,
      weight_decay=weight_decay)

  tx2 = optax.adamw(  # pytype: disable=wrong-arg-types  # numpy-scalars
      learning_rate=learning_rate_fn,
      b1=config.optimizer.beta1,
      b2=config.optimizer.beta2,
      weight_decay=0.0)

  tx = optax.chain(
      optax.masked(tx1, decay_param_mask),
      optax.masked(tx2, no_decay_param_mask),
  )
  return tx
Constructs the optimizer from the given HParams.
videogvt/train_lib/maskgvt_train_lib.py