code: def optimize(
      loss, learning_rate, hparams, use_tpu=False, variables=None
  ):
      """Minimize loss."""
      loss = weight_decay_and_noise(loss, hparams, learning_rate)
      loss = tf.identity(loss, name="total_loss")
      if variables is None:
          variables = tf.trainable_variables()
      # Print trainable variables.
      log_variable_sizes(variables, verbose=hparams.summarize_vars)
      # Print non-trainable variables.
      non_trainable_variables = list(
          set(tf.global_variables()) - set(variables)
      )
      log_variable_sizes(
          non_trainable_variables,
          tag="Non-trainable variables",
          verbose=hparams.summarize_vars,
      )
      if hparams.summarize_vars:
          summarize_variables(variables)
          # Summarize non-trainable variables as well
          summarize_variables(
              non_trainable_variables, tag="Non-trainable variables"
          )
      diet_vars = [
          v
          for v in tf.global_variables()
          if v.dtype == dtypes.float16_ref
      ]
      log_variable_sizes(
          diet_vars, "Diet Variables", verbose=hparams.summarize_vars
      )
      opt = ConditionalOptimizer(
          hparams.optimizer, learning_rate, hparams, use_tpu
      )
      if use_tpu:
          opt = contrib.tpu().CrossShardOptimizer(opt)
      if getattr(hparams, "gpu_automatic_mixed_precision", False):
          if use_tpu:
              raise RuntimeError(
                  "GPU auto mixed precision cannot be used with TPU"
              )
          elif _mixed_precision_is_enabled(hparams):
              raise RuntimeError(
                  "GPU auto mixed precision cannot be used with manual mixed precision"
              )
          else:
              setattr(opt, "_use_locking", "True")
              setattr(opt, "_name", "ConditionalOptimizer")
              opt = tf.train.experimental.enable_mixed_precision_graph_rewrite(
                  opt
              )
  
      opt_summaries = []
      if common_layers.should_generate_summaries():
          tf.summary.scalar("learning_rate", learning_rate)
          opt_summaries.append("loss")
          if hparams.summarize_grads:
              tf.logging.info("Summarizing gradients")
              opt_summaries.extend(
                  ["gradients", "gradient_norm", "global_gradient_norm"]
              )
  
      if hparams.clip_grad_norm:
          tf.logging.info(
              "Clipping gradients, norm: %0.5f", hparams.clip_grad_norm
          )
      if hparams.grad_noise_scale:
          tf.logging.info(
              "Adding noise to gradients, noise scale: %0.5f",
              hparams.grad_noise_scale,
          )
  
      train_op = contrib.layers().optimize_loss(
          name="training",
          loss=loss,
          global_step=tf.train.get_or_create_global_step(),
          learning_rate=learning_rate,
          clip_gradients=hparams.clip_grad_norm or None,
          gradient_noise_scale=hparams.grad_noise_scale or None,
          optimizer=opt,
          summaries=opt_summaries,
          colocate_gradients_with_ops=True,
          variables=variables,
      )
      return train_op
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import numpy
  from tensor2tensor.layers import common_layers
  from tensor2tensor.utils import adafactor
  from tensor2tensor.utils import contrib
  from tensor2tensor.utils import misc_utils
  from tensor2tensor.utils import mlperf_log
  from tensor2tensor.utils import multistep_optimizer
  from tensor2tensor.utils import registry
  from tensor2tensor.utils import yellowfin
  import tensorflow.compat.v1
  from tensorflow.python.framework import dtypes
label_desc: Use image node_id 2 for calling a global function with example usage: optimize(loss, learning_rate, hparams, use_tpu, variables) and returns: train_op