code: def waveform_summary(audio, audio_gen, step, name=""):
      """Creates a waveform plot summary for a batch of audio."""
  
      def plot_waveform(i, length=None, prefix="waveform", name=""):
          """Plots a waveforms."""
          waveform = np.squeeze(audio[i])
          waveform = (
              waveform[:length] if length is not None else waveform
          )
          waveform_gen = np.squeeze(audio_gen[i])
          waveform_gen = (
              waveform_gen[:length]
              if length is not None
              else waveform_gen
          )
          # Manually specify exact size of fig for tensorboard
          fig, (ax0, ax1) = plt.subplots(2, 1, figsize=(2.5, 2.5))
          ax0.plot(waveform)
          ax1.plot(waveform_gen)
  
          # Format and save plot to image
          name = name + "_" if name else ""
          tag = f"waveform/{name}{prefix}_{i+1}"
          fig_summary(tag, fig, step)
  
      # Make plots at multiple lengths.
      batch_size = int(audio.shape[0])
      for i in range(batch_size):
          plot_waveform(i, length=None, prefix="full", name=name)
          plot_waveform(i, length=2000, prefix="125ms", name=name)
  
imports: import io
  import ddsp
  from ddsp.core import tf_float32
  from ddsp.training.plotting import pianoroll_plot_setup
  import matplotlib.patches
  import matplotlib.pyplot
  from matplotlib.ticker import MaxNLocator
  import note_seq
  from note_seq import sequences_lib
  import numpy
  import tensorflow.compat.v2
label_desc: Use image node_id 2 for calling a global function with example usage: waveform_summary(audio, audio_gen, step, name) without return types