code: def _testAttentionCellWrapperCorrectResult(self):
      num_units = 4
      attn_length = 6
      batch_size = 2
      expected_output = np.array(
          [
              [1.068372, 0.45496, -0.678277, 0.340538],
              [1.018088, 0.378983, -0.572179, 0.268591],
          ],
          dtype=np.float32,
      )
      expected_state = np.array(
          [
              [
                  0.74946702,
                  0.34681597,
                  0.26474735,
                  1.06485605,
                  0.38465962,
                  0.11420801,
                  0.10272158,
                  0.30925757,
                  0.63899988,
                  0.7181077,
                  0.47534478,
                  0.33715725,
                  0.58086717,
                  0.49446869,
                  0.7641536,
                  0.12814975,
                  0.92231739,
                  0.89857256,
                  0.21889746,
                  0.38442063,
                  0.53481543,
                  0.8876909,
                  0.45823169,
                  0.5905602,
                  0.78038228,
                  0.56501579,
                  0.03971386,
                  0.09870267,
                  0.8074435,
                  0.66821432,
                  0.99211812,
                  0.12295902,
                  1.14606023,
                  0.34370938,
                  -0.79251152,
                  0.51843399,
              ],
              [
                  0.5179342,
                  0.48682183,
                  -0.25426468,
                  0.96810579,
                  0.28809637,
                  0.13607743,
                  -0.11446252,
                  0.26792109,
                  0.78047138,
                  0.63460857,
                  0.49122369,
                  0.52007174,
                  0.73000264,
                  0.66986895,
                  0.73576689,
                  0.86301267,
                  0.87887371,
                  0.35185754,
                  0.93417215,
                  0.64732957,
                  0.63173044,
                  0.66627824,
                  0.53644657,
                  0.20477486,
                  0.98458421,
                  0.38277245,
                  0.03746676,
                  0.92510188,
                  0.57714164,
                  0.84932971,
                  0.36127412,
                  0.12125921,
                  1.1362772,
                  0.34361625,
                  -0.78150457,
                  0.70582712,
              ],
          ],
          dtype=np.float32,
      )
      seed = 12345
      tf.set_random_seed(seed)
      rnn_scope = None
      for state_is_tuple in [False, True]:
          with tf.Session() as sess:
              with tf.variable_scope(
                  "state_is_tuple",
                  reuse=state_is_tuple,
                  initializer=tf.glorot_uniform_initializer(),
              ):
                  lstm_cell = rnn_cell.BasicLSTMCell(
                      num_units, state_is_tuple=state_is_tuple
                  )
                  cell = contrib_rnn.AttentionCellWrapper(
                      lstm_cell,
                      attn_length,
                      state_is_tuple=state_is_tuple,
                  )
                  # This is legacy behavior to preserve the test.  Weight
                  # sharing no longer works by creating a new RNNCell in the
                  # same variable scope; so here we restore the scope of the
                  # RNNCells after the first use below.
                  if rnn_scope is not None:
                      (
                          cell._scope,
                          lstm_cell._scope,
                      ) = rnn_scope  # pylint: disable=protected-access,unpacking-non-sequence
                  zeros1 = tf.random_uniform(
                      (batch_size, num_units), 0.0, 1.0, seed=seed + 1
                  )
                  zeros2 = tf.random_uniform(
                      (batch_size, num_units), 0.0, 1.0, seed=seed + 2
                  )
                  zeros3 = tf.random_uniform(
                      (batch_size, num_units), 0.0, 1.0, seed=seed + 3
                  )
                  attn_state_zeros = tf.random_uniform(
                      (batch_size, attn_length * num_units),
                      0.0,
                      1.0,
                      seed=seed + 4,
                  )
                  zero_state = (
                      (zeros1, zeros2),
                      zeros3,
                      attn_state_zeros,
                  )
                  if not state_is_tuple:
                      zero_state = tf.concat(
                          [
                              zero_state[0][0],
                              zero_state[0][1],
                              zero_state[1],
                              zero_state[2],
                          ],
                          1,
                      )
                  inputs = tf.random_uniform(
                      (batch_size, num_units), 0.0, 1.0, seed=seed + 5
                  )
                  output, state = cell(inputs, zero_state)
                  # This is legacy behavior to preserve the test.  Weight
                  # sharing no longer works by creating a new RNNCell in the
                  # same variable scope; so here we store the scope of the
                  # first RNNCell for reuse above.
                  if rnn_scope is None:
                      rnn_scope = (
                          cell._scope,
                          lstm_cell._scope,
                      )  # pylint: disable=protected-access
                  if state_is_tuple:
                      state = tf.concat(
                          [
                              state[0][0],
                              state[0][1],
                              state[1],
                              state[2],
                          ],
                          1,
                      )
                  sess.run(tf.global_variables_initializer())
                  self.assertAllClose(sess.run(output), expected_output)
                  self.assertAllClose(sess.run(state), expected_state)
  
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 5 for calling the RNNCellTest obj's underlying member method code with example usage: obj._testAttentionCellWrapperCorrectResult() without return types