code: def train_op(batch, hparams, config_name):
      """Define a training op, including summaries and optimization.
  
      Args:
        batch: Dictionary produced by NSynthDataset.
        hparams: Hyperparameters dictionary.
        config_name: Name of config module.
  
      Returns:
        train_op: A complete iteration of training with summaries.
      """
      config = utils.get_module(
          "baseline.models.ae_configs.%s" % config_name
      )
  
      if hparams.raw_audio:
          x = batch["audio"]
          # Add height and channel dims
          x = tf.expand_dims(tf.expand_dims(x, 1), -1)
      else:
          x = batch["spectrogram"]
  
      # Define the model
      with tf.name_scope("Model"):
          z = config.encode(x, hparams)
          xhat = config.decode(z, batch, hparams)
  
      # For interpolation
      tf.add_to_collection("x", x)
      tf.add_to_collection("pitch", batch["pitch"])
      tf.add_to_collection("z", z)
      tf.add_to_collection("xhat", xhat)
  
      # Compute losses
      total_loss = compute_mse_loss(x, xhat, hparams)
  
      # Apply optimizer
      with tf.name_scope("Optimizer"):
          global_step = tf.get_variable(
              "global_step",
              [],
              tf.int64,
              initializer=tf.constant_initializer(0),
              trainable=False,
          )
          optimizer = tf.train.AdamOptimizer(
              hparams.learning_rate, hparams.adam_beta
          )
          train_step = slim.learning.create_train_op(
              total_loss, optimizer, global_step=global_step
          )
  
      return train_step
  
file_path: magenta/magenta/models/nsynth/baseline/models/ae.py
imports: from magenta.contrib import training
  from magenta.models.nsynth import utils
  import numpy
  import tensorflow.compat.v1
  import tf_slim
label_desc: Use image node_id 3 for calling a global function with example usage: train_op(batch, hparams, config_name) and returns: train_step