code: def _get_optimizer_fn(
      optimizer,
      learning_rate,
      use_tpu,
      exponential_decay_steps=-1,
      exponential_decay_rate=-1,
      lr_warmup_steps=0,
  ):
      """Returns a function that gives the optimizer."""
  
      def optimizer_fn():
          global_step = tf.compat.v1.train.get_or_create_global_step()
          new_learning_rate = learning_rate
          if exponential_decay_steps > 0 and exponential_decay_rate > 0:
              new_learning_rate = tf.compat.v1.train.exponential_decay(
                  learning_rate=learning_rate,
                  global_step=global_step,
                  decay_steps=exponential_decay_steps,
                  decay_rate=exponential_decay_rate,
                  staircase=True,
              )
  
          if lr_warmup_steps > 0:
              warmup_lr = (
                  learning_rate
                  * tf.cast(global_step, tf.float32)
                  / tf.cast(lr_warmup_steps, tf.float32)
              )
  
              warmup_fn = lambda: tf.minimum(warmup_lr, learning_rate)
  
              new_learning_rate = tf.cond(
                  pred=global_step < lr_warmup_steps,
                  true_fn=warmup_fn,
                  false_fn=lambda: new_learning_rate,
              )
  
          if use_tpu:
              opt = tf.compat.v1.tpu.CrossShardOptimizer(
                  _OPTIMIZERS[optimizer](
                      learning_rate=new_learning_rate
                  )
              )
          else:
              opt = _OPTIMIZERS[optimizer](
                  learning_rate=new_learning_rate
              )
          return opt
  
      return optimizer_fn
  
imports: import collections
  import functools
  import inspect
  from absl import logging
  from model_search import block_builder
  from model_search import ensembler
  from model_search.architecture import architecture_utils
  from model_search.architecture import tower
  from model_search.ensembler import EnsembleLogits
  from model_search.generators import base_tower_generator
  from model_search.generators import trial_utils
  from model_search.meta import distillation
  import numpy
  import tensorflow.compat.v2
label_desc: Use image node_id 3 for calling a global function with example usage: _get_optimizer_fn(optimizer, learning_rate, use_tpu, exponential_decay_steps, exponential_decay_rate, lr_warmup_steps) and returns: optimizer_fn, opt