code: def optimizer_with_layerwise_decay(
      config: ml_collections.ConfigDict,
      params: PyTree,
      use_frozen_params: bool = True,
  ):
      """Returns an optimizer with layerwise decay.
  
      Implementation of layerwise decay follows BEIT and MAE.
      Reference: https://github.com/facebookresearch/mae/blob/main/util/lr_decay.py
  
      This function can apply layerwise decay to any optimizer, although this is
      typically done with Adam.
  
      Args:
        config: The training config.
        params: The parameters of the model being trained.
        use_frozen_params: If True, the optimizer will always expect to receive
          a FrozenDict of parameters and gradients.
  
      Returns:
        An Optax optimizer.
      """
      if config.model_name not in {
          "vit_multilabel_classification",
          "vit_multilabel_classification_mae",
          "vit_classification_mae",
          "vivit_classification_mae",
          "vivit_multimodal_classification",
          "mbt_classification",
          "vivit_multimodal_multiclassification",
          "mbt_multilabel_classification",
      }:
          raise ValueError(f"Unsupported model: {config.model_name}.")
  
      optimizer_config = optimizer_lib.get_optax_optimizer_config(
          config
      )
      # Avoid modifying original config and allow alteration.
      optimizer_config = copy.deepcopy(optimizer_config).unlock()
      base_learning_rate = config.lr_configs.base_learning_rate
  
      if optimizer_config.get("layerwise_decay", 0) <= 0:
          logging.info("Not performing any layerwise decay.")
          if "layerwise_decay" in optimizer_config:
              del optimizer_config.layerwise_decay
          lr_fn = lr_schedules.get_learning_rate_fn(config)
          return optimizer_lib.get_optimizer(
              optimizer_config, lr_fn, params
          )
  
      num_transformer_layers = config.model.num_layers
      num_layers = num_transformer_layers + 1
      layer_decay = optimizer_config.layerwise_decay
      learning_rate_scales = [
          layer_decay ** (num_layers - i) for i in range(num_layers + 1)
      ]
      logging.info("Learning rate scales: %s", learning_rate_scales)
  
      layer_configs = [
          copy.deepcopy(config) for _ in range(num_layers + 1)
      ]
      for index in range(len(layer_configs)):
          learning_rate = (
              base_learning_rate * learning_rate_scales[index]
          )
          layer_configs[
              index
          ].lr_configs.base_learning_rate = learning_rate
  
      learning_rate_fns = [
          lr_schedules.get_learning_rate_fn(layer_config)
          for layer_config in layer_configs
      ]
  
      # Weight decay mask is applied within optimizer_lib.get_optimizer.
      # Note that we need to delete the layerwise_decay attribute, as Optax
      # optimizers do not accept this argument.
      del optimizer_config.layerwise_decay
      optimizers = {
          i: optimizer_lib.get_optimizer(
              optimizer_config, learning_rate_fns[i], params
          )
          for i in range(num_layers + 1)
      }
  
      def _get_layer_id(name: str, num_layers: int) -> int:
          if (
              name == "cls"
              or "posembed_input" in name
              or "embedding" in name
              or "modality_token" in name
          ):
              return 0
          elif (
              re.match(r"Transformer(.*)/encoderblock_", name)
              is not None
          ):
              substring = re.findall(r"encoderblock_\d+", name)[0]
              layer_id = int(substring.replace("encoderblock_", ""))
              if "Transformer_concat" in name:
                  return (
                      layer_id + 11
                  )  # TODO(lgeorgescu): compute it somehow.
              return layer_id + 1
          else:
              return num_layers
  
      flat_params = flax.traverse_util.flatten_dict(
          flax.core.unfreeze(params), keep_empty_nodes=True, sep="/"
      )
      flat_layer_map = {
          k: _get_layer_id(k, num_layers) for k in flat_params
      }
      layer_map = flax.traverse_util.unflatten_dict(
          flat_layer_map, sep="/"
      )
      if use_frozen_params:
          layer_map = flax.core.freeze(layer_map)
  
      logging.info(
          "Layer assignments:\n%s",
          flax.traverse_util.flatten_dict(layer_map, sep="/"),
      )
      tx = optax.multi_transform(optimizers, layer_map)
  
      return tx
  
file_path: scenic/scenic/projects/av_mae/optimizer_utils.py
imports: import copy
  import re
  from typing import Any, Callable, Optional, Union
  from absl import logging
  import flax
  import ml_collections
  import optax
  from scenic.train_lib import lr_schedules
  from scenic.train_lib import optimizers
label_desc: Use image node_id 1 for calling a global function with example usage: optimizer_with_layerwise_decay(config, params, use_frozen_params) and returns: tx, optimizer_lib, int, unknown, num_layers, unknown