def bahdanau_attention(key, context, hidden_size, projected_align=False):
if hidden_size % <NUM_LIT:2> != <NUM_LIT:0>:<EOL><INDENT>raise ValueError("<STR_LIT>")<EOL><DEDENT>batch_size = tf.shape(context)[<NUM_LIT:0>]<EOL>max_num_tokens, token_size = context.get_shape().as_list()[-<NUM_LIT:2>:]<EOL>r_context = tf.reshape(context, shape=[-<NUM_LIT:1>, max_num_tokens, token_size])<EOL>projected_key = tf.layers.dense(key, hidden_size, kernel_initializer=xav())<EOL>r_projected_key =tf.tile(tf.reshape(projected_key, shape=[-<NUM_LIT:1>, <NUM_LIT:1>, hidden_size]),<EOL>[<NUM_LIT:1>, max_num_tokens, <NUM_LIT:1>])<EOL>lstm_fw_cell = tf.nn.rnn_cell.LSTMCell(hidden_size//<NUM_LIT:2>)<EOL>lstm_bw_cell = tf.nn.rnn_cell.LSTMCell(hidden_size//<NUM_LIT:2>)<EOL>(output_fw, output_bw), states =tf.nn.bidirectional_dynamic_rnn(cell_fw=lstm_fw_cell,<EOL>cell_bw=lstm_bw_cell,<EOL>inputs=r_context,<EOL>dtype=tf.float32)<EOL>bilstm_output = tf.concat([output_fw, output_bw], -<NUM_LIT:1>)<EOL>concat_h_state = tf.concat([r_projected_key, output_fw, output_bw], -<NUM_LIT:1>)<EOL>projected_state =tf.layers.dense(concat_h_state, hidden_size, use_bias=False,<EOL>kernel_initializer=xav())<EOL>score =tf.layers.dense(tf.tanh(projected_state), units=<NUM_LIT:1>, use_bias=False,<EOL>kernel_initializer=xav())<EOL>attn = tf.nn.softmax(score, dim=<NUM_LIT:1>)<EOL>if projected_align:<EOL><INDENT>log.info("<STR_LIT>")<EOL>t_context = tf.transpose(bilstm_output, [<NUM_LIT:0>, <NUM_LIT:2>, <NUM_LIT:1>])<EOL>output = tf.reshape(tf.matmul(t_context, attn),<EOL>shape=[batch_size, -<NUM_LIT:1>, hidden_size])<EOL><DEDENT>else:<EOL><INDENT>log.info("<STR_LIT>")<EOL>t_context = tf.transpose(r_context, [<NUM_LIT:0>, <NUM_LIT:2>, <NUM_LIT:1>])<EOL>output = tf.reshape(tf.matmul(t_context, attn),<EOL>shape=[batch_size, -<NUM_LIT:1>, token_size])<EOL><DEDENT>return output<EOL>
It is a implementation of the Bahdanau et al. attention mechanism. Based on the paper:
        https://arxiv.org/abs/1409.0473 "Neural Machine Translation by Jointly Learning to Align and Translate"
    Args:
        key: A tensorflow tensor with dimensionality [None, None, key_size]
        context: A tensorflow tensor with dimensionality [None, None, max_num_tokens, token_size]
        hidden_size: Number of units in hidden representation
        projected_align: Using bidirectional lstm for hidden representation of context.
        If true, beetween input and attention mechanism insert layer of bidirectional lstm with dimensionality [hidden_size].
        If false, bidirectional lstm is not used.
    Returns:
        output: Tensor at the output with dimensionality [None, None, hidden_size]