code: def _testStackBidirectionalDynamicRNNStates(self, use_gpu):
      # Check that the states are correctly initialized.
      # - Create a net and iterate for 3 states. Keep the state (state_3).
      # - Reset states, and iterate for 5 steps. Last state is state_5.
      # - Reset the sets to state_3 and iterate for 2 more steps,
      #   last state will be state_5'.
      # - Check that the state_5 and state_5' (forward and backward) are the
      #   same for the first layer (it does not apply for the second layer since
      #   it has forward-backward dependencies).
      with self.session(use_gpu=use_gpu, graph=tf.Graph()) as sess:
          batch_size = 2
          # Create states placeholders.
          initial_states_fw = [
              tf.placeholder(tf.float32, shape=(batch_size, layer * 2))
              for layer in self.layers
          ]
          initial_states_bw = [
              tf.placeholder(tf.float32, shape=(batch_size, layer * 2))
              for layer in self.layers
          ]
          # Create the net
          (
              input_value,
              inputs,
              outputs,
              state_fw,
              state_bw,
              sequence_length,
          ) = self._createStackBidirectionalDynamicRNN(
              use_gpu,
              use_shape=True,
              use_state_tuple=False,
              initial_states_fw=initial_states_fw,
              initial_states_bw=initial_states_bw,
          )
          tf.global_variables_initializer().run()
  
          # Run 3 steps.
          feed_dict = {inputs[0]: input_value, sequence_length: [3, 2]}
          # Initialize to empty state.
          for i, layer in enumerate(self.layers):
              feed_dict[initial_states_fw[i]] = np.zeros(
                  (batch_size, layer * 2), dtype=np.float32
              )
              feed_dict[initial_states_bw[i]] = np.zeros(
                  (batch_size, layer * 2), dtype=np.float32
              )
          _, st_3_fw, st_3_bw = sess.run(
              [outputs, state_fw, state_bw], feed_dict=feed_dict
          )
  
          # Reset the net and run 5 steps.
          feed_dict = {inputs[0]: input_value, sequence_length: [5, 3]}
          for i, layer in enumerate(self.layers):
              feed_dict[initial_states_fw[i]] = np.zeros(
                  (batch_size, layer * 2), dtype=np.float32
              )
              feed_dict[initial_states_bw[i]] = np.zeros(
                  (batch_size, layer * 2), dtype=np.float32
              )
          _, st_5_fw, st_5_bw = sess.run(
              [outputs, state_fw, state_bw], feed_dict=feed_dict
          )
  
          # Reset the net to state_3 and run 2 more steps.
          feed_dict = {inputs[0]: input_value, sequence_length: [2, 1]}
          for i, _ in enumerate(self.layers):
              feed_dict[initial_states_fw[i]] = st_3_fw[i]
              feed_dict[initial_states_bw[i]] = st_3_bw[i]
          _, st_5p_fw, st_5p_bw = sess.run(
              [outputs, state_fw, state_bw], feed_dict=feed_dict
          )
  
          # Check that the 3+2 and 5 first layer states.
          self.assertAllEqual(st_5_fw[0], st_5p_fw[0])
          self.assertAllEqual(st_5_bw[0], st_5p_bw[0])
  
imports: import itertools
  from absl.testing import parameterized
  from magenta.contrib import rnn
  import numpy
  import tensorflow.compat.v1
label_desc: Use image node_id 4 for calling the StackBidirectionalRNNTest obj's underlying member method code with example usage: obj._testStackBidirectionalDynamicRNNStates(use_gpu) without return types