def one_step(self, l, images):
        '''
        l = [n_examples, 5]
        image = [n_examples, height, width]
        '''

        tol = 1e-4
        g_x = self.B * (l[:, 0] + 1) / 2.
        g_y = self.A * (l[:, 1] + 1) / 2.
        delta = (max(self.A, self.B) - 1) / (self.N - 1) * T.exp(l[:, 2])
        sigma = T.exp(l[:, 3])

        mu_x = g_x.dimshuffle([0, 'x']) +\
            (self.mu_ind - self.N / 2. + 0.5) * delta.dimshuffle([0, 'x'])
        mu_y = g_y.dimshuffle([0, 'x']) +\
            (self.mu_ind - self.N / 2. + 0.5) * delta.dimshuffle([0, 'x'])

        F_x = T.exp(-((self.B_ind - mu_x.dimshuffle([0, 1, 'x']))**2) / (
            2 * (sigma.dimshuffle([0, 'x', 'x'])))**2)
        F_x = F_x / (F_x.sum(axis=-1).dimshuffle(0, 1, 'x') + tol)

        # Compute Y filter banks##
        F_y = T.exp(-((self.A_ind - mu_y.dimshuffle([0, 1, 'x']))**2) / (
            2 * (sigma.dimshuffle([0, 'x', 'x'])))**2)
        F_y = F_y / (F_y.sum(axis=-1).dimshuffle(0, 1, 'x') + tol)

        read = T.batched_dot(T.batched_dot(F_y, images), F_x.dimshuffle([0, 2, 1]))
        return read, g_x, g_y, delta, sigma

 def energy(self):
     rho_x = rho(self.x)
     rho_h = rho(self.h)
     squared_norm = ( T.batched_dot(self.x,self.x) + T.batched_dot(self.h,self.h) ) / 2
     uni_terms    = - T.dot(rho_x, self.bx) - T.dot(rho_h, self.bh)
     bi_terms     = - T.batched_dot( T.dot(rho_x, self.W1), rho_h )
     return squared_norm + uni_terms + bi_terms

    def get_output_for(self, inputs, **kwargs):

        # seq_input: (batch_size, seq_size, n_hidden_con)
        # seq_mask: (batch_size, seq_size)
        # condition: (batch_size, n_hidden_con)
        seq_input, seq_mask, condition = inputs

        if self.gate_covariance:
            update = T.nnet.sigmoid(
                T.sum(seq_input * self.w_gate, axis=-1, keepdims=True) +
                self.b_gate)
            seq_input *= update

        length_seq = seq_input.shape[1]
        if self.covariance_decay:
            decay = T.arange(1, length_seq+1)
            decay = (self.covariance_decay +
                     (length_seq-decay) * (1 - self.covariance_decay))
            decay = T.sqrt(decay)
            decay = decay.dimshuffle('x', 0, 'x')
            seq_input *= decay

        seq_input *= T.shape_padright(seq_mask)
        # (batch_size, n_hidden_question, n_hidden_question)
        covariance = T.batched_dot(seq_input.dimshuffle(0, 2, 1), seq_input)
        # (batch_size, n_hidden_question), equivalent to the following line:
        # att = T.sum(covariance * condition.dimshuffle((0, 'x', 1)), axis=2)
        att = 1000 * T.batched_dot(covariance, condition.dimshuffle((0, 1)))

        if not self.covariance_decay:
            att /= T.sum(seq_mask, axis=1, keepdims=True)
        # norm2_att = T.sum(att * condition, axis=1, keepdims=True)
        # att = 1000 * att / norm2_att

        return att

def compute_psi2(lls, lsf, z, input_means, input_vars):

    ls = T.exp(lls)
    sf = T.exp(lsf)
    b = ls / casting(2.0)
    term_1 = T.prod(T.sqrt(b / (b + input_vars)), 1)

    scale = T.sqrt(4 * (2 * b[ None, : ] + 0 * input_vars))
    scaled_z = z[ None, : , : ] / scale[ : , None , : ]
    scaled_z_minus_m = scaled_z
    r2b = T.sum(scaled_z_minus_m**2, 2)[ :, None, : ] + T.sum(scaled_z_minus_m**2, 2)[ :, : , None ] - \
        2 * T.batched_dot(scaled_z_minus_m, np.transpose(scaled_z_minus_m, [ 0, 2, 1 ]))
    term_2 = T.exp(-r2b)

    scale = T.sqrt(4 * (2 * b[ None, : ] + 2 * input_vars))
    scaled_z = z[ None, : , : ] / scale[ : , None , : ]
    scaled_m = input_means / scale
    scaled_m = T.tile(scaled_m[ : , None, : ], [ 1, z.shape[ 0 ], 1])
    scaled_z_minus_m = scaled_z - scaled_m
    r2b = T.sum(scaled_z_minus_m**2, 2)[ :, None, : ] + T.sum(scaled_z_minus_m**2, 2)[ :, : , None ] + \
        2 * T.batched_dot(scaled_z_minus_m, np.transpose(scaled_z_minus_m, [ 0, 2, 1 ]))
    term_3 = T.exp(-r2b)
    
    psi2_computed = sf**casting(2.0) * term_1[ :, None, None ] * term_2 * term_3

    return T.transpose(psi2_computed, [ 1, 2, 0 ])

 def propup_given_h_lag(self, vt, h_lag, hbias):
     if h_lag == self.h0:
         x = T.batched_dot(vt, self.W) + T.addbroadcast(
             T.dot(h_lag, self.Wt) + hbias, 0, 1)
     else:
         x = T.batched_dot(vt, self.W) + hbias + T.dot(h_lag, self.Wt)
     return [x, T.nnet.sigmoid(x)]

    def read(self, images, center_y, center_x, delta, sigma):
        """
        Parameters
        ----------
        images : T.matrix    (shape: batch_size x img_size)
            Batch of images. Internally it will be reshaped to be a 
            (batch_size, img_height, img_width)-shaped stack of images.
        center_y : T.vector (shape: batch_size)
        center_x : T.vector (shape: batch_size)
        delta : T.vector    (shape: batch_size)
        sigma : T.vector    (shape: batch_size)

        Returns
        -------
        window : T.matrix   (shape: batch_size x N**2)
        """
        N = self.N
        batch_size = images.shape[0]

        # Reshape input into proper 2d images
        I = images.reshape( (batch_size, self.img_height, self.img_width) )

        # Get separable filterbank
        FY, FX = self.filterbank_matrices(center_y, center_x, delta, sigma)

        # apply to the batch of images
        W = T.batched_dot(T.batched_dot(FY, I), FX.transpose([0,2,1]))

        return W.reshape((batch_size, N*N))

  def fwd(self, x, V, A, L):
    """
    x : signal
    V : eigenvectors
    A : area 
    L : eigenvalues
    """
    V = V[:,:self.K]
    L = L[:self.K]

    L = L.dimshuffle('x','x',0)

    rho = T.sqrt(T.sum(A))
   
    # Q x 1 x K, a window for each input function
    ghat = self.activation_interp(
            T.batched_dot(T.tile(L, [self.nin,1,1]), self.Winterp))
    # Q x K x N
    V_ = T.tile(V.dimshuffle('x',1,0), [self.nin, 1, 1])
    # Q x K x N
    tmp = (ghat * V).dimshuffle(0,2,1)
    
    # Q x N x N
    transl = rho * T.batched_dot(V_.dimshuffle(0,2,1), tmp)
    transl = A.dimshuffle('x',0,'x') * transl
    
    # Q x K x N
    tmp = (V.dimshuffle(0,'x',1) * x.dimshuffle(0,1,'x')).dimshuffle(1,2,0)
    # Q x K x N
    desc = rho * T.batched_dot(tmp, transl)
    desc = T.abs_(desc)
    
    desc = desc.dimshuffle(2,0,'x',1) # BC01 format : N x Q x 1 x K
    return self.activation(theano.tensor.nnet.conv.conv2d(desc, self.W).flatten(2) + self.b)

 def h_given_h_lag_vt(self, vt, h_lag, hbias):
     if h_lag == self.h0:
         x = T.batched_dot(vt, self.W) + T.addbroadcast(
             T.dot(h_lag, self.Wt) + hbias.dimshuffle('x', 0), 0, 1)
     else:
         x = T.batched_dot(vt, self.W) + \
             T.dot(h_lag, self.Wt) + hbias.dimshuffle('x', 0)
     return [x, T.nnet.sigmoid(x)]

 def defmodel(self):
     lhs = T.ivector("lhs")
     rhs, nrhs = T.ivectors("rhs","nrhs")
     lhsemb = self.entembs[lhs, :]
     rhsemb = self.W[rhs, :]
     nrhsemb = self.W[nrhs, :]
     pdot = T.batched_dot(lhsemb, rhsemb)
     ndot = T.batched_dot(lhsemb, nrhsemb)
     return pdot, ndot, [lhs, rhs, nrhs]

 def energy_function():
     squared_norm = (
         T.batched_dot(self.x, self.x) + T.batched_dot(self.h, self.h) + T.batched_dot(self.y, self.y)
     ) / 2.0
     uni_terms = -T.dot(self.rho_x, self.bx) - T.dot(self.rho_h, self.bh) - T.dot(self.rho_y, self.by)
     bi_terms = -T.batched_dot(T.dot(self.rho_x, self.W1), self.rho_h) - T.batched_dot(
         T.dot(self.rho_h, self.W2), self.rho_y
     )
     return squared_norm + uni_terms + bi_terms

    def step(self, x, states):
        h_tild_tm1 = states[0]

        B_U = states[1]
        B_W = states[2]

        if self.consume_less == 'cpu':
            x_i = x[:, :self.output_dim]
            x_f = x[:, self.output_dim: 2 * self.output_dim]
            x_c = x[:, 2 * self.output_dim: 3 * self.output_dim]
            x_o = x[:, 3 * self.output_dim: 4 * self.output_dim]
            x_new = x[:, 4 * self.output_dim:]
        else:
            x_i = K.dot(x * B_W[0], self.W_i) + self.b_i
            x_f = K.dot(x * B_W[1], self.W_f) + self.b_f
            x_c = K.dot(x * B_W[2], self.W_c) + self.b_c
            x_o = K.dot(x * B_W[3], self.W_o) + self.b_o
            x_new = x

        # self.C_tape -> BT, t-1, k
        # self.H_tape -> BT, t-1, k

        # x -> BT, k 
        # h_tild_tm1 -> BT, k       

        if self.H_tape is None:
            self.H_tape = K.zeros_like(h_tild_tm1).dimshuffle((0,'x',1))
            self.C_tape = K.zeros_like(h_tild_tm1).dimshuffle((0,'x',1))

        # s_t -> BT, t-1, 1
        t = K.shape(self.C_tape)[1]

        sum1 = K.dot(self.H_tape, self.W_h)
        sum2 = K.dot(K.repeat_elements(x_new.dimshuffle((0,'x',1)),t, axis=1), self.W_x)
        sum3 = K.dot(K.repeat_elements(h_tild_tm1.dimshuffle((0,'x',1)),t, axis=1), self.W_h_tilde)
        tanhed_sum = K.tanh(sum1 + sum2 + sum3)    
        a_t = K.dot(tanhed_sum, self.v)[:,:,0]
        s_t = K.softmax(a_t)

        h_tilde_t = T.batched_dot(self.H_tape.dimshuffle((0,2,1)), s_t.dimshuffle((0,1,'x')))[:,:,0]
        c_tilde_t = T.batched_dot(self.C_tape.dimshuffle((0,2,1)), s_t.dimshuffle((0,1,'x')))[:,:,0]

        i = self.inner_activation(x_i + K.dot(h_tilde_t * B_U[0], self.U_i))
        f = self.inner_activation(x_f + K.dot(h_tilde_t * B_U[1], self.U_f))
        c_t = f * c_tilde_t + i * self.activation(x_c + K.dot(h_tilde_t * B_U[2], self.U_c))
        o = self.inner_activation(x_o + K.dot(h_tilde_t * B_U[3], self.U_o))

        h_t = o * self.activation(c_t)

        # Add to Tape
        self.C_tape = K.concatenate([self.C_tape, c_t.dimshuffle((0,'x',1))], axis=1)
        self.H_tape = K.concatenate([self.H_tape, h_t.dimshuffle((0,'x',1))], axis=1)

        return h_t, [h_tilde_t]

    def batched_cos_sim(s):
        """ from (x,y,z)-shaped pair, produce (x,y)-shaped pair that replaces the z-vector pairs by their cosine similarities """
        import theano
        import theano.tensor as T

        return theano.scan(
            fn=lambda xm, ym: T.batched_dot(xm, ym) / T.sqrt(T.batched_dot(xm, xm) * T.batched_dot(ym, ym)),
            outputs_info=None,
            sequences=s,
            non_sequences=None,
        )[0]

 def free_energy_given_hid_lag(self, vt, h_lag, hbias, vbias):
     if h_lag == self.h0:
         wx_b = T.batched_dot(vt, self.W) +\
             T.addbroadcast(T.dot(h_lag, self.Wt) + hbias, 0, 1)
         vbias_term = T.batched_dot(vt, vbias)
         hidden_term = T.sum(T.log(1 + T.exp(wx_b)), axis=2)
     else:
         wx_b = T.batched_dot(vt, self.W) + T.dot(h_lag, self.Wt) + \
             hbias.dimshuffle('x', 0)
         vbias_term = T.batched_dot(vt, vbias)
         hidden_term = T.sum(T.log(1 + T.exp(wx_b)), axis=2)
     return -hidden_term - vbias_term

def MemLayer(incomings, params, linear=0):
  '''
  incomings = (u, u_shape, A, A_shape, C, C_shape)
  '''
  ((u, u_shape), (A, A_shape), (C, C_shape)) = incomings
  p = T.switch(linear, T.batched_dot(A, u), nnet.softmax(T.batched_dot(A, u)))
  p_shape = A_shape[:2]
  # C.shape = (batch_size, num_sen, embed_size), u.shape = (batch_size, embed_size)
  # p.shape = (batch_size, num_sen, 1)
  #return (p, u_shape)
  O = (C * p[:, :, None]).sum(axis = 1)

  return ((O, u_shape), (p, p_shape))

    def write(self, windows, center_y, center_x, delta, sigma):
        N = self.N
        batch_size = windows.shape[0]

        # Reshape input into proper 2d windows
        W = windows.reshape( (batch_size, N, N) )

        # Get separable filterbank
        FY, FX = self.filterbank_matrices(center_y, center_x, delta, sigma)

        # apply...
        I = T.batched_dot(T.batched_dot(FY.transpose([0,2,1]), W), FX)

        return I.reshape( (batch_size, self.img_height*self.img_width) )

    def factorization(self, batchSize, argsEmbA, argsEmbB, wC, wC1, wC2):
        # l = batchSize
        # k = self.k  # embed size
        # r = self.r  # relation number
        # argEmbedsA = self.A[argsA.flatten()]  # [l,k]
        # argEmbedsB = self.A[argsB.flatten()]  # [l,k]

        # first = T.tensordot(relationProbs, self.C, axes=[[1], [2]])  # [l,r] * [k,k,r] = [l, k, k]
        Afirst = T.batched_tensordot(wC, argsEmbA, axes=[[1], [1]])  # + self.Cb  # [l, k, k] * [l, k] = [l, k]
        Asecond = T.batched_dot(Afirst, argsEmbB)  # [l, k] * [l, k] = [l]
        # entropy = T.sum(T.log(relationProbs) * relationProbs, axis=1)  # [l,r] * [l,r] = [l]
        spFirst = T.batched_dot(wC1, argsEmbA)
        spSecond = T.batched_dot(wC2, argsEmbB)
        return Asecond + spFirst + spSecond

    def _initialize_posterior_distribution(self, RecognitionParams):

        # Now actually compute the precisions (from their square roots)
        self.Lambda = T.batched_dot(self.LambdaChol, self.LambdaChol.dimshuffle(0,2,1))

        # dynamics matrix & initialize the innovations precision, xDim x xDim
        self.A         = theano.shared(value=RecognitionParams['A'].astype(theano.config.floatX)        ,name='A'        )
        self.QinvChol  = theano.shared(value=RecognitionParams['QinvChol'].astype(theano.config.floatX) ,name='QinvChol' )
        self.Q0invChol = theano.shared(value=RecognitionParams['Q0invChol'].astype(theano.config.floatX),name='Q0invChol')

        self.Qinv  = T.dot(self.QinvChol,self.QinvChol.T)
        self.Q0inv = T.dot(self.Q0invChol,self.Q0invChol.T)

        ################## put together the total precision matrix ######################

        AQinvA = T.dot(T.dot(self.A.T, self.Qinv), self.A)

        # for now we (suboptimally) replicate a bunch of times
        AQinvrep = Tsla.kron(T.ones([self.Tt-1,1,1]),-T.dot(self.A.T, self.Qinv)) # off-diagonal blocks (upper triangle)

        AQinvArep = Tsla.kron(T.ones([self.Tt-2,1,1]), AQinvA+self.Qinv)
        AQinvArepPlusQ = T.concatenate([T.shape_padleft(self.Q0inv + AQinvA), AQinvArep, T.shape_padleft(self.Qinv)])

        # This is our inverse covariance matrix: diagonal (AA) and off-diagonal (BB) blocks.
        self.AA = self.Lambda + AQinvArepPlusQ
        self.BB = AQinvrep

        # symbolic recipe for computing the the diagonal (V) and
        # off-diagonal (VV) blocks of the posterior covariance
        self.V, self.VV, self.S = compute_sym_blk_tridiag(self.AA, self.BB)

        # now compute the posterior mean
        LambdaMu = T.batched_dot(self.Lambda, self.Mu) # scale by precision (no need for transpose; lambda is symmetric)

        #self.old_postX = compute_sym_blk_tridiag_inv_b(self.S,self.V,LambdaMu) # apply inverse

        # compute cholesky decomposition
        self.the_chol = blk_tridag_chol(self.AA, self.BB)
        # intermediary (mult by R^T) -
        ib = blk_chol_inv(self.the_chol[0], self.the_chol[1], LambdaMu)
        # final result (mult by R)-
        self.postX = blk_chol_inv(self.the_chol[0], self.the_chol[1], ib, lower=False, transpose=True)

        # The determinant of the covariance is the square of the determinant of the cholesky factor.
        # Determinant of the Cholesky factor is the product of the diagonal elements of the block-diagonal.
        def comp_log_det(L):
            return T.log(T.diag(L)).sum()
        self.ln_determinant = -2*theano.scan(fn=comp_log_det, sequences=self.the_chol[0])[0].sum()

 def apply(self, idx, inp):
     '''
     :param idx: vector of indices, one per sample
     :param inp: matrix (nb_samples, dims)
     :return:
     '''
     return T.batched_dot(inp, self.W[idx-self.idxoffset])

 def get_output_for(self, inputs, **kwargs):
     #input[0]:(BS,max_senlen,emb_size),input[1]:(BS,1,emb_size),input[2]:(BS,max_sentlen)
     # activation0=(T.dot(inputs[0],self.W_h)).reshape([self.batch_size,self.max_sentlen])+self.b_h.repeat(self.batch_size,0).repeat(self.max_sentlen,1)
     # activation1=T.dot(inputs[1],self.W_q).reshape([self.batch_size]).dimshuffle(0,'x')
     # activation2=T.batched_dot(T.dot(inputs[0],self.W_o),inputs[1].reshape([self.batch_size,self.embedding_size,1])).reshape([self.batch_size,self.max_sentlen])
     activation2=T.batched_dot(inputs[0],inputs[1].reshape([self.batch_size,self.embedding_size,1])).reshape([self.batch_size,self.max_sentlen])
     norm2=T.sqrt(T.sum(T.mul(inputs[0],inputs[0]),axis=2))+0.0000001
     activation2=activation2/norm2
     # activation=(self.nonlinearity(activation0)+self.nonlinearity(activation1)+activation2).reshape([self.batch_size,self.max_sentlen])#.dimshuffle(0,'x',2)#.repeat(self.max_sentlen,axis=1)
     activation2=(activation2).reshape([self.batch_size,self.max_sentlen])#.dimshuffle(0,'x',2)#.repeat(self.max_sentlen,axis=1)
     # final=T.dot(activation,self.W_o) #(BS,max_sentlen)
     activation3=T.batched_dot(inputs[0],inputs[1].reshape([self.batch_size,self.embedding_size,1])).reshape([self.batch_size,self.max_sentlen])
     # if inputs[2] is not None:
     #     final=inputs[2]*final-(1-inputs[2])*1000000
     alpha=lasagne.nonlinearities.softmax(activation2) #(BS,max_sentlen)
     return alpha

    def __init__(self, x, y, l, window, opt, lr, init_emb, dim_emb, dim_hidden, n_vocab, L2_reg, unit,
                 sim='cos', n_layers=1, activation=tanh):
        self.tr_inputs = [x, y, l]
        self.pr_inputs = [x, y, l]

        self.x = x  # 1D: batch_size * l * 2, 2D: window; elem=word_id
        self.y = y  # 1D: batch_size; elem=label
        self.l = l  # scalar: elem=sentence length

        batch_size = y.shape[0]
        n_cands = x.shape[0] / batch_size / l

        self.pad = build_shared_zeros((1, dim_emb))
        if init_emb is None:
            self.emb = theano.shared(sample_weights(n_vocab - 1, dim_emb))
        else:
            self.emb = theano.shared(init_emb)
        self.E = T.concatenate([self.pad, self.emb], 0)
        self.W_out = theano.shared(sample_weights(dim_hidden, dim_hidden))
        self.params = [self.emb, self.W_out]

        """ Input Layer """
        e = self.E[x]  # e: 1D: batch_size * l * 2, 2D: window, 3D: dim_emb
        x_in = e.reshape((batch_size * n_cands, l, -1))

        """ Intermediate Layer """
        # h: 1D: n_batch * n_cands, 2D: dim_emb
        h, params = cnn.layers(x_in, window, dim_emb, dim_hidden, n_layers, activation)
        self.params.extend(params)

        """ Output Layer """
        h = h.reshape((batch_size, n_cands, -1))
        h_1 = h[T.arange(batch_size), 0]
        h_2 = h[T.arange(batch_size), 1:]
        if sim == 'cos':
            y_score = cosign_similarity(h_1, h_2)
        else:
            y_score = T.batched_dot(T.dot(h_1, self.W_out), h_2.dimshuffle(0, 2, 1))
        y_score_hat = T.max(y_score, 1)

        """ Objective Function """
        self.nll = max_margin_loss(y_score_hat, y_score[T.arange(batch_size), y])
        self.L2_sqr = regularization(self.params)
        self.cost = self.nll + L2_reg * self.L2_sqr / 2.

        """ Optimization """
        if opt == 'adagrad':
            self.update = ada_grad(cost=self.cost, params=self.params, lr=lr)
        elif opt == 'ada_delta':
            self.update = ada_delta(cost=self.cost, params=self.params)
        elif opt == 'adam':
            self.update = adam(cost=self.cost, params=self.params, lr=lr)
        else:
            self.update = sgd(cost=self.cost, params=self.params, lr=lr)

        """ Predicts """
        y_hat = T.argmax(y_score, 1)

        """ Check Accuracies """
        self.correct = T.eq(y_hat, y)