code: def get_optimizer(
      optimizer_config: ml_collections.ConfigDict,
      learning_rate_fn: ScalarOrSchedule,
      params: Optional[PyTree] = None,
  ) -> optax.GradientTransformation:
      """Constructs the optimizer from the given configuration.
  
      Args:
        optimizer_config: Configuration specific to the optimizer.
        learning_rate_fn: Learning rate schedule.
        params: Parameters pytree, used when we want to skip weight decay on bias
          and scale parameters.
  
      Returns:
        An optax GradientTransformation, this consists of a pair of pure functions
        implementing a gradient transformation.
      """
      # Avoid modifying original config and allow alteration.
      config = copy.deepcopy(optimizer_config).unlock()
  
      # Skip weight decay for BatchNorm scale or for the bias parameters.
      weight_decay_mask = None
      if "skip_scale_and_bias_regularization" in config:
          if config.skip_scale_and_bias_regularization and config.get(
              "weight_decay", 0
          ):
              if params is None:
                  raise ValueError(
                      "params must be given to obtain weight_decay_mask."
                  )
              weight_decay_mask = jax.tree_util.tree_map(
                  lambda x: x.ndim != 1, params
              )
          del config.skip_scale_and_bias_regularization
  
      optim_ops = []
      # Add weight decay for sgd (possibly with momentum and nesterov).
      if config.optimizer == "sgd" and "weight_decay" in config:
          if config.weight_decay:
              optim_ops.append(
                  optax.add_decayed_weights(
                      config.weight_decay, weight_decay_mask
                  )
              )
          del config.weight_decay
  
      if weight_decay_mask and config.optimizer in {
          "adamw",
          "lamb",
          "adamaxw",
          "radamw",
      }:
          config.mask = weight_decay_mask
      elif weight_decay_mask and config.optimizer in {
          "adafactor",
          "lars",
      }:
          config.weight_decay_mask = weight_decay_mask
  
          if config.optimizer == "adafactor":
              config.weight_decay_rate = config.weight_decay
              del config.weight_decay
  
      eps = compute_eps(params)
      if "ada" in config.optimizer or config.optimizer in {
          "lamb",
          "lars",
          "rmsprop",
          "yogi",
      }:
          config.eps = eps
  
      # Add gradient clipping before optimizer operations.
      if "grad_clip" in config:
          grad_clip_config = config.grad_clip
          clip_method = grad_clip_config.get("clip_method", None)
          clip_value = grad_clip_config.get("clip_value", None)
          if clip_method is not None and clip_value is not None:
              if clip_method == "clip_by_global_norm":
                  optim_ops.append(
                      optax.clip_by_global_norm(clip_value)
                  )
              elif clip_method == "adaptive_grad_clip":
                  optim_ops.append(
                      optax.adaptive_grad_clip(clip_value, eps=eps)
                  )
              elif clip_method == "clip":
                  optim_ops.append(optax.clip(clip_value))
              elif clip_method == "clip_by_block_rms":
                  optim_ops.append(optax.clip_by_block_rms(clip_value))
              else:
                  logging.info("%s is not supported", clip_method)
          del config.grad_clip
  
      if "gradient_centralization" in config:
          if config["gradient_centralization"]:
              optim_ops.append(optax.centralize())
          del config["gradient_centralization"]
  
      if "lookahead" in config:
          lookahead = config["lookahead"]
          lookahead_sync_period = config.get("lookahead_sync_period")
          lookahead_slow_step_size = config.get(
              "lookahead_slow_step_size"
          )
  
          del config["lookahead"]
  
          if "lookahead_sync_period" in config:
              del config["lookahead_sync_period"]
  
          if "lookahead_slow_step_size" in config:
              del config["lookahead_slow_step_size"]
      else:
          lookahead = False
          lookahead_sync_period = None
          lookahead_slow_step_size = None
  
      if "params_to_freeze" in config:
          params_to_freeze = config.params_to_freeze
          del config["params_to_freeze"]
      else:
          params_to_freeze = []
  
      # Call the optax optimizer with exact arguments as in the config.
      optimizer_fn = getattr(optax, config.optimizer)
      del config.optimizer
      optimizer = optimizer_fn(learning_rate=learning_rate_fn, **config)
  
      if lookahead:
          # FIXME: it messes with the masks, because the params are gonna be diff.
          optimizer = optax.lookahead(
              optimizer,
              sync_period=lookahead_sync_period,
              slow_step_size=lookahead_slow_step_size,
          )
  
      optim_ops.append(optimizer)
  
      if params_to_freeze:
          params_to_freeze = [re.compile(p) for p in params_to_freeze]
  
          def _freeze_param(name: str) -> bool:
              return any(bool(p.search(name)) for p in params_to_freeze)
  
          zero_mask = tree_map_with_names_values(
              lambda _, n: _freeze_param(n), params
          )
          logging.info("Params that are frozen: %s", zero_mask)
          optim_ops.append(optax.masked(optax.set_to_zero(), zero_mask))
  
      return optax.chain(*optim_ops)
  
imports: import copy
  import dataclasses
  import re
  from typing import Any, Callable, Generator, List, Optional, Tuple, Union
  from absl import logging
  import flax
  import jax
  import jax.numpy
  import ml_collections
  import numpy
  import optax
label_desc: Use image node_id 3 for calling a global function with example usage: get_optimizer(optimizer_config, learning_rate_fn, params) and returns: optax, any