import tensorflow as tf


        batch_size = tf.shape(attention_weights)[0]
        src_len = tf.shape(attention_weights)[2]
        trg_len = tf.shape(attention_weights)[1]

        src_indices = tf.tile(tf.reshape(tf.range(src_len), shape=[1, 1, src_len]), [batch_size, trg_len, 1])
        trg_indices = tf.tile(tf.reshape(tf.range(trg_len), shape=[1, trg_len, 1]), [batch_size, 1, src_len])

        source_length = encoder_input_length[0]
        target_length = tf.to_int32(tf.reduce_sum(trg_mask, axis=1))
        true_src_len = tf.reshape(source_length, shape=[batch_size, 1, 1]) - 1
        true_trg_len = tf.reshape(target_length, shape=[batch_size, 1, 1]) - 1
