import tensorflow as tf

        dec, mem = tf.nn.seq2seq.attention_decoder(
            dec_inp, enc_state,
            attn_states, cell, output_size=4,
            num_heads=2)
        sess.run([tf.global_variables_initializer()])
        res = sess.run(dec)
        self.assertEqual(3, len(res))
        self.assertEqual((2, 4), res[0].shape)
