code: def create_optimizer(
      config: ml_collections.ConfigDict,
  ) -> optax.GradientTransformation:
      """Creates an optimizer, as specified by the config."""
      if config.optimizer == "adam":
          return optax.adam(learning_rate=config.learning_rate)
      if config.optimizer == "sgd":
          return optax.sgd(
              learning_rate=config.learning_rate,
              momentum=config.momentum,
          )
      raise ValueError(f"Unsupported optimizer: {config.optimizer}.")
  
imports: import os
  from typing import Any, Dict, Iterable, Tuple, Optional
  from absl import logging
  from clu import checkpoint
  from clu import metric_writers
  from clu import metrics
  from clu import parameter_overview
  from clu import periodic_actions
  import flax
  import flax.core
  import flax.linen
  from flax.training import train_state
  import jax
  import jax.numpy
  import jraph
  import ml_collections
  import numpy
  import optax
  import sklearn.metrics
  import tensorflow
  import input_pipeline
  import models
label_desc: Use image node_id 2 for calling a global function with example usage: create_optimizer(config) and returns: optax, optax