:param s: [src_sequence_length, batch_size, src_dim]
        :param h: [batch_size, tgt_dim] or [tgt_sequence_length, batch_size, tgt_dim]
        :param mask: [src_sequence_length, batch_size]\
             or [tgt_sequence_length, src_sequence_length, batch_sizse]
        :param pre_compute: [src_sequence_length, batch_size, hidden_dim]
        :return: [src_sequence_length, batch_size]\
             or [tgt_sequence_length, src_sequence_length, batch_size]

def get_prob(self, src, tgt, mask, pre_compute, return_logits=False):
        '''
        :param s: [src_sequence_length, batch_size, src_dim]
        :param h: [batch_size, tgt_dim] or [tgt_sequence_length, batch_size, tgt_dim]
        :param mask: [src_sequence_length, batch_size]\
             or [tgt_sequence_length, src_sequence_length, batch_sizse]
        :param pre_compute: [src_sequence_length, batch_size, hidden_dim]
        :return: [src_sequence_length, batch_size]\
             or [tgt_sequence_length, src_sequence_length, batch_size]
        '''
        s_shape = src.get_shape().as_list()
        h_shape = tgt.get_shape().as_list()
        src_dim = s_shape[-1]
        tgt_dim = h_shape[-1]
        assert src_dim is not None, 'src dimension must be defined'
        assert tgt_dim is not None, 'tgt dimension must be defined'

        self._define_params(src_dim, tgt_dim)

        if len(h_shape) == 2:
            tgt = tf.expand_dims(tgt, 0)
        if pre_compute is None:
            pre_compute = self.get_pre_compute(src)

        buf0 = pre_compute
        buf1 = tf.tensordot(tgt, self.var['U'], axes=[[2], [0]])
        buf2 = tf.tanh(tf.expand_dims(buf0, 0) + tf.expand_dims(buf1, 1))

        if not self.is_vanilla:
            xh1 = tgt
            xh2 = tgt
            s1 = src
            if self.need_padding:
                xh1 = tf.tensordot(xh1, self.var['V_t'], 1)
                xh2 = tf.tensordot(xh2, self.var['S_t'], 1)
                s1 = tf.tensordot(s1, self.var['V_s'], 1)
            if not self.is_identity_transform:
                xh1 = tf.tensordot(xh1, self.var['T'], 1)
                xh2 = tf.tensordot(xh2, self.var['T'], 1)
            buf3 = tf.expand_dims(s1, 0) * tf.expand_dims(xh1, 1)
            buf3 = tf.tanh(tf.tensordot(buf3, self.var['V'], axes=[[3], [0]]))
            buf = tf.reshape(tf.tanh(buf2 + buf3), shape=tf.shape(buf3))
        else:
            buf = buf2
        v = self.var['v']
        e = tf.tensordot(buf, v, [[3], [0]])
        e = tf.squeeze(e, axis=[3])
        tmp = tf.reshape(e + (mask - 1) * 10000.0, shape=tf.shape(e))
        prob = tf.nn.softmax(tmp, 1)
        if len(h_shape) == 2:
            prob = tf.squeeze(prob, axis=[0])
            tmp = tf.squeeze(tmp, axis=[0])
        if return_logits:
            return prob, tmp
        return prob