code: def lstm_layer(
      inputs,
      num_units,
      bidirectional,
      is_training,
      lengths=None,
      stack_size=1,
      dropout_keep_prob=1,
  ):
      """Create a LSTM layer using the specified backend."""
      cells_fw = []
      for i in range(stack_size):
          del i
          cell = tf.nn.rnn_cell.BasicLSTMCell(num_units)
          cell = tf.nn.rnn_cell.DropoutWrapper(
              cell,
              output_keep_prob=dropout_keep_prob
              if is_training
              else 1.0,
          )
          cells_fw.append(cell)
  
      if bidirectional:
          cells_bw = []
          for i in range(stack_size):
              del i
              cell = tf.nn.rnn_cell.BasicLSTMCell(num_units)
              cell = tf.nn.rnn_cell.DropoutWrapper(
                  cell,
                  output_keep_prob=dropout_keep_prob
                  if is_training
                  else 1.0,
              )
              cells_bw.append(cell)
          with tf.variable_scope("lstm"):
              (
                  outputs,
                  unused_state_f,
                  unused_state_b,
              ) = contrib_rnn.stack_bidirectional_dynamic_rnn(
                  cells_fw,
                  cells_bw,
                  inputs,
                  dtype=tf.float32,
                  sequence_length=lengths,
                  parallel_iterations=1,
              )
  
          return outputs
      else:
          with tf.variable_scope("lstm"):
              outputs, unused_state = tf.nn.dynamic_rnn(
                  cell=tf.nn.rnn_cell.MultiRNNCell(cells_fw),
                  inputs=inputs,
                  dtype=tf.float32,
                  sequence_length=lengths,
                  parallel_iterations=1,
              )
              return outputs
  
file_path: magenta/magenta/models/onsets_frames_transcription/model_tpu.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from magenta.contrib import rnn
  from magenta.contrib import training
  from magenta.models.onsets_frames_transcription import constants
  from magenta.models.onsets_frames_transcription import estimator_spec_util
  import tensorflow.compat.v1
  from tensorflow.compat.v1 import estimator
  import tf_slim
label_desc: Use image node_id 2 for calling a global function with example usage: lstm_layer(inputs, num_units, bidirectional, is_training, lengths, stack_size, dropout_keep_prob) and returns: outputs, outputs