code: def set_final(sequence, sequence_length, values, time_major=False):
      """Sets the final values in a batch of sequences, and clears those after."""
      sequence_batch_major = (
          sequence
          if not time_major
          else tf.transpose(sequence, [1, 0, 2])
      )
      final_index = _get_final_index(sequence_length, time_major=False)
      mask = tf.sequence_mask(
          tf.maximum(0, sequence_length - 1),
          maxlen=sequence_batch_major.shape[1],
          dtype=tf.float32,
      )
      sequence_batch_major = tf.expand_dims(
          mask, axis=-1
      ) * sequence_batch_major + tf.scatter_nd(
          final_index, values, tf.shape(sequence_batch_major)
      )
      if time_major:
          return tf.transpose(sequence_batch_major, [1, 0, 2])
      return sequence_batch_major
  
imports: import collections
  from magenta.contrib import rnn
  from magenta.contrib import seq2seq
  import tensorflow.compat.v1
label_desc: Use image node_id 5 for calling a global function with example usage: set_final(sequence, sequence_length, values, time_major) and returns: sequence_batch_major, tf