def get_att(self, s, prob):
        '''
        :param s: [src_sequence_length, batch_size, src_dim]
        :param prob: [src_sequence_length, batch_size]\
            or [tgt_sequence_length, src_sequence_length, batch_size]
        :return: [batch_size, src_dim] or [tgt_sequence_length, batch_size, src_dim]
        '''
        buf = s * tf.expand_dims(prob, axis=-1)
        att = tf.reduce_sum(buf, axis=-3)
        return att