code: def get_optim(
      optimizer_name: str,
      learning_rate_fn: Callable[[int], float],
      wd: float,
  ):
      """Returns list of operation for optax optimizer."""
      optim_ops = []
      if optimizer_name in ["sgd", "momentum"]:
          if wd:
              optim_ops.append(optax.add_decayed_weights(wd))
          if optimizer_name == "sgd":
              optim_ops.append(
                  optax.sgd(learning_rate=learning_rate_fn, momentum=0)
              )  # pytype: disable=wrong-arg-types  # numpy-scalars
          else:
              optim_ops.append(
                  optax.sgd(
                      learning_rate=learning_rate_fn, momentum=0.9
                  )
              )  # pytype: disable=wrong-arg-types  # numpy-scalars
      elif optimizer_name == "adamw":
          optim_ops.append(
              optax.adamw(
                  learning_rate=learning_rate_fn,  # pytype: disable=wrong-arg-types  # numpy-scalars
                  weight_decay=wd,
              )
          )
      elif optimizer_name == "adafactor":
          optim_ops.append(
              optax.adafactor(
                  learning_rate=learning_rate_fn,  # pytype: disable=wrong-arg-types  # numpy-scalars
                  multiply_by_parameter_scale=False,
                  momentum=0.9,
                  decay_rate=0.999,
                  weight_decay_rate=wd,
              )
          )
      else:
          logging.info('Unknown optimizer "%s"', optimizer_name)
      return optim_ops
  
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 6 for calling a global function with example usage: get_optim(optimizer_name, learning_rate_fn, wd) and returns: optim_ops