code: def _get_input_tensors(dataset, config):
      """Get input tensors from dataset."""
      batch_size = config.hparams.batch_size
      iterator = tf.data.make_one_shot_iterator(dataset)
      (
          input_sequence,
          output_sequence,
          control_sequence,
          sequence_length,
      ) = iterator.get_next()
      input_sequence.set_shape(
          [batch_size, None, config.data_converter.input_depth]
      )
      output_sequence.set_shape(
          [batch_size, None, config.data_converter.output_depth]
      )
      if not config.data_converter.control_depth:
          control_sequence = None
      else:
          control_sequence.set_shape(
              [batch_size, None, config.data_converter.control_depth]
          )
      sequence_length.set_shape(
          [batch_size] + sequence_length.shape[1:].as_list()
      )
  
      return {
          "input_sequence": input_sequence,
          "output_sequence": output_sequence,
          "control_sequence": control_sequence,
          "sequence_length": sequence_length,
      }
  
file_path: magenta/magenta/models/music_vae/music_vae_train.py
imports: import os
  from magenta.models.music_vae import configs
  from magenta.models.music_vae import data
  import tensorflow.compat.v1
  import tf_slim
label_desc: Use image node_id 2 for calling a global function with example usage: _get_input_tensors(dataset, config) and returns: dict