def METHOD_NAME(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