import tensorflow as tf

def generate_seq2seq_mask(attention_mask, mask_sequence, seq_type, **kargs):
    if seq_type == 'seq2seq':
        if mask_sequence is not None:
            seq_shape = get_shape_list(mask_sequence, expected_rank=2)
            seq_len = seq_shape[1]
            ones = tf.ones((1, seq_len, seq_len))
            a_mask = tf.matrix_band_part(ones, -1, 0)
            s_ex12 = tf.expand_dims(tf.expand_dims(mask_sequence, 1), 2)
            s_ex13 = tf.expand_dims(tf.expand_dims(mask_sequence, 1), 3)
            a_mask = (1 - s_ex13) * (1 - s_ex12) + s_ex13 * a_mask
            # generate mask of batch x seq_len x seq_len
            a_mask = tf.reshape(a_mask, (-1, seq_len, seq_len))
            out_mask = attention_mask * a_mask
        else:
            ones = tf.ones_like(attention_mask[:1])
            mask = (tf.matrix_band_part(ones, -1, 0))
