code: def set_lr_configs(
      config: ml_collections.ConfigDict,
      datasets_metadata: Dict[str, Dict[str, Any]],
  ) -> ml_collections.ConfigDict:
      """Sets learning rate configurations."""
  
      num_total_train_examples = 0
      for ds_metadata in datasets_metadata.values():
          num_total_train_examples += ds_metadata.get(
              "num_train_examples", 0
          )
  
      with config.unlocked():
          decay_events = {500: 400}
          steps_per_epoch = (
              num_total_train_examples // config.batch_size
          )
          config.lr_configs.learning_rate_schedule = "compound"
          config.lr_configs.factors = "constant*piecewise_constant"
          config.lr_configs.decay_events = [
              decay_events.get(
                  config.num_training_epochs,
                  config.num_training_epochs * 2 // 3,
              )
              * steps_per_epoch,
          ]
          # Note: this is absolute (not relative):
          config.lr_configs.decay_factors = [0.1]
          config.lr_configs.base_learning_rate = 1e-4
  
          # Also set backbone lr configs
          config.backbone_training.lr_configs = copy.deepcopy(
              config.lr_configs
          )
          config.backbone_training.lr_configs.base_learning_rate = 1e-5
      return config
  
imports: import collections
  import copy
  import functools
  from typing import Any, Dict, Mapping, Optional, Sequence, Tuple, Union
  from absl import logging
  import flax
  import flax.linen
  import jax
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.common_lib import debug_utils
  from scenic.train_lib_deprecated import optimizers
  from scenic.train_lib_deprecated import train_utils
label_desc: Use image node_id 3 for calling a global function with example usage: set_lr_configs(config, datasets_metadata) and returns: config