code: def get_build_graph_fn(
      mode, config, sequence_example_file_paths=None
  ):
      """Returns a function that builds the TensorFlow graph.
  
      Args:
        mode: 'train', 'eval', or 'generate'. Only mode related ops are added to
            the graph.
        config: An EventSequenceRnnConfig containing the encoder/decoder and HParams
            to use.
        sequence_example_file_paths: A list of paths to TFRecord files containing
            tf.train.SequenceExample protos. Only needed for training and
            evaluation. May be a sharded file of the form.
  
      Returns:
        A function that builds the TF ops when called.
  
      Raises:
        ValueError: If mode is not 'train', 'eval', or 'generate'.
      """
      if mode not in ("train", "eval", "generate"):
          raise ValueError(
              "The mode parameter must be 'train', 'eval', "
              "or 'generate'. The mode parameter was: %s" % mode
          )
  
      hparams = config.hparams
      encoder_decoder = config.encoder_decoder
  
      tf.logging.info("hparams = %s", hparams.values())
  
      input_size = encoder_decoder.input_size
  
      def build():
          """Builds the Tensorflow graph."""
          inputs, lengths = None, None
  
          if mode in ("train", "eval"):
              inputs, _, lengths = magenta.common.get_padded_batch(
                  sequence_example_file_paths,
                  hparams.batch_size,
                  input_size,
                  shuffle=mode == "train",
              )
  
          elif mode == "generate":
              inputs = tf.placeholder(
                  tf.float32, [hparams.batch_size, None, input_size]
              )
  
          cell = events_rnn_graph.make_rnn_cell(
              hparams.rnn_layer_sizes,
              dropout_keep_prob=hparams.dropout_keep_prob
              if mode == "train"
              else 1.0,
              attn_length=hparams.attn_length,
              residual_connections=hparams.residual_connections,
          )
  
          rnn_nade = RnnNade(
              cell,
              num_dims=input_size,
              num_hidden=hparams.nade_hidden_units,
          )
  
          if mode in ("train", "eval"):
              log_probs, cond_probs = rnn_nade.log_prob(inputs, lengths)
  
              inputs_flat = tf.to_float(
                  magenta.common.flatten_maybe_padded_sequences(
                      inputs, lengths
                  )
              )
              predictions_flat = tf.to_float(
                  tf.greater_equal(cond_probs, 0.5)
              )
  
              if mode == "train":
                  loss = tf.reduce_mean(-log_probs)
                  perplexity = tf.reduce_mean(tf.exp(log_probs))
                  correct_predictions = tf.to_float(
                      tf.equal(inputs_flat, predictions_flat)
                  )
                  accuracy = tf.reduce_mean(correct_predictions)
                  precision = tf.reduce_sum(
                      inputs_flat * predictions_flat
                  ) / tf.reduce_sum(predictions_flat)
                  recall = tf.reduce_sum(
                      inputs_flat * predictions_flat
                  ) / tf.reduce_sum(inputs_flat)
  
                  optimizer = tf.train.AdamOptimizer(
                      learning_rate=hparams.learning_rate
                  )
  
                  train_op = tf_slim.learning.create_train_op(
                      loss,
                      optimizer,
                      clip_gradient_norm=hparams.clip_norm,
                  )
                  tf.add_to_collection("train_op", train_op)
  
                  vars_to_summarize = {
                      "loss": loss,
                      "metrics/perplexity": perplexity,
                      "metrics/accuracy": accuracy,
                      "metrics/precision": precision,
                      "metrics/recall": recall,
                  }
              elif mode == "eval":
                  (
                      vars_to_summarize,
                      update_ops,
                  ) = tf_slim.metrics.aggregate_metric_map(
                      {
                          "loss": tf.metrics.mean(-log_probs),
                          "metrics/perplexity": tf.metrics.mean(
                              tf.exp(log_probs)
                          ),
                          "metrics/accuracy": tf.metrics.accuracy(
                              inputs_flat, predictions_flat
                          ),
                          "metrics/precision": tf.metrics.precision(
                              inputs_flat, predictions_flat
                          ),
                          "metrics/recall": tf.metrics.recall(
                              inputs_flat, predictions_flat
                          ),
                      }
                  )
                  for updates_op in update_ops.values():
                      tf.add_to_collection("eval_ops", updates_op)
  
              precision = vars_to_summarize["metrics/precision"]
              recall = vars_to_summarize["metrics/precision"]
              f1_score = tf.where(
                  tf.greater(precision + recall, 0),
                  2 * ((precision * recall) / (precision + recall)),
                  0,
              )
              vars_to_summarize["metrics/f1_score"] = f1_score
              for var_name, var_value in vars_to_summarize.items():
                  tf.summary.scalar(var_name, var_value)
                  tf.add_to_collection(var_name, var_value)
  
          elif mode == "generate":
              initial_state = rnn_nade.zero_state(hparams.batch_size)
  
              final_state = rnn_nade.steps(inputs, initial_state)
              samples, log_prob = rnn_nade.sample_single(initial_state)
  
              tf.add_to_collection("inputs", inputs)
              tf.add_to_collection("sample", samples)
              tf.add_to_collection("log_prob", log_prob)
  
              # Flatten state tuples for metagraph compatibility.
              for state in tf.nest.flatten(initial_state):
                  tf.add_to_collection("initial_state", state)
              for state in tf.nest.flatten(final_state):
                  tf.add_to_collection("final_state", state)
  
      return build
  
file_path: magenta/magenta/models/pianoroll_rnn_nade/pianoroll_rnn_nade_graph.py
imports: import collections
  import magenta
  from magenta.common import Nade
  from magenta.contrib import seq2seq
  from magenta.models.shared import events_rnn_graph
  import tensorflow.compat.v1
  import tf_slim
label_desc: Use image node_id 1 for calling a global function with example usage: get_build_graph_fn(mode, config, sequence_example_file_paths) and returns: build