code: def build_optimizer(
      config: ml_collections.ConfigDict, params: PyTree
  ):
      """Builds optimizer."""
      # Defaults
      wd = config.get("weight_decay", 0)
      optimizer_name = config.get("optimizer", "sgd")
      lr_config = config.get(
          "lr_configs", ml_collections.ConfigDict({"factors": ""})
      )
      optim_ops = get_optim(
          optimizer_name,
          lr_schedules.compound_lr_scheduler(lr_config),
          wd,
      )
  
      if config.get("multi_optim", False):
          # Define the multiple optimizers.
          transforms = {"default": optax.chain(*optim_ops)}
          for sub_optim_name in config.multi_optim_configs.strings:
              sub_optim_ops = get_optim(
                  config.get(
                      "optimizer_" + sub_optim_name, optimizer_name
                  ),
                  lr_schedules.compound_lr_scheduler(
                      config.get("lr_" + sub_optim_name, lr_config)
                  ),
                  config.get("weight_decay_" + sub_optim_name, wd),
              )
              transforms[sub_optim_name] = optax.chain(*sub_optim_ops)
  
          # Combine optimizers.
          def label_param(name):
              for sub_optim_name in config.multi_optim_configs.strings:
                  if sub_optim_name in name:
                      return sub_optim_name
              return "default"
  
          optim_ops = [
              optax.multi_transform(
                  transforms,
                  optimizers.tree_map_with_names_values(
                      lambda _, name: label_param(name), params
                  ),
              )
          ]
  
      # Explicit freezing of some parts of the model. Note that freezing can also
      # be done by setting learning rate to 0 with config.multi_optim set to True.
      if config.get("freeze_video_encoder", False):
          # Zero out updates for video encoder parameters.
          mask = optimizers.tree_map_with_names_values(
              lambda _, n: True if "video_encoder" in n else False,
              params,
          )
          optim_ops.append(optax.masked(optax.set_to_zero(), mask))
  
      if config.get("freeze_text_encoder", False):
          # Zero out updates for text encoder parameters.
          mask = optimizers.tree_map_with_names_values(
              lambda _, n: True if "text_encoder" in n else False,
              params,
          )
          optim_ops.append(optax.masked(optax.set_to_zero(), mask))
  
      tx = optax.chain(*optim_ops)
      opt_state = jax.jit(tx.init, backend="cpu")(params)
      return tx, opt_state
  
imports: import os
  import re
  from typing import Any, Callable, Dict, List, Mapping, Optional, Tuple, Union
  from absl import logging
  import flax
  from flax import struct
  from flax.core import frozen_dict
  import flax.linen
  from flax.training import checkpoints
  import jax
  import jax.numpy
  import jax.profiler
  import ml_collections
  import numpy
  import optax
  from scenic.dataset_lib import dataset_utils
  from scenic.model_lib.base_models import base_model
  from scenic.projects.baselines.clip import model
  from scenic.train_lib import lr_schedules
  from scenic.train_lib import optimizers
  from scenic.train_lib import train_utils
  from tensorflow.io import gfile
label_desc: Use image node_id 7 for calling a global function with example usage: build_optimizer(config, params) and returns: tx,  opt_state, str, sub_optim_name