code: def get_optax_optimizer_config(
      config: ml_collections.ConfigDict,
  ) -> ml_collections.ConfigDict:
      """Obtain optimizer from main config."""
      optimizer_config = config.get(
          "optimizer_configs", ml_collections.ConfigDict()
      )
  
      # New-style config: all optimizer-related fields are in optimizer_configs.
      if "optimizer" in optimizer_config:
          if "optimizer" in config:
              raise ValueError(
                  "Both config.optimizer and config.optimizer_configs.optimizer are "
                  "defined. Define it only once to avoid possible contradictions. "
                  "The preferred location is in config.optimizer_configs.optimizer"
              )
          return optimizer_config
  
      # Backwards compatibility: copy optimizer field into the optimizer config.
      optimizer_config = copy.deepcopy(optimizer_config).unlock()
      if "optimizer" in config:
          optimizer_config.optimizer = config.optimizer
  
          # The old optimizers have adam with weight decay. However, in optax this is
          # done using the adamw optimizer.
          if (
              config.optimizer == "adam"
              and "weight_decay" in optimizer_config
          ):
              optimizer_config.optimizer = "adamw"
  
          if config.optimizer == "momentum":
              optimizer_config.optimizer = "sgd"
  
          if config.optimizer == "nesterov":
              optimizer_config.optimizer = "sgd"
              optimizer_config.nesterov = True
  
      if "skip_scale_and_bias_regularization" in config:
          optimizer_config.skip_scale_and_bias_regularization = (
              config.skip_scale_and_bias_regularization
          )
  
      optimizer_config = _scenic_optimizer_args_to_optax_args(
          optimizer_config
      )
  
      if "grad_clip_configs" in config:
          optimizer_config.grad_clip = config.grad_clip_configs
  
      optimizer_config.lock()
      logging.info(
          "Optimizer config after backwards compatibility operations:\n%s",
          optimizer_config,
      )
      return optimizer_config
  
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 4 for calling a global function with example usage: get_optax_optimizer_config(config) and returns: optimizer_config, optimizer_config