def step(self, i_t, x_t, z_t, att_p, y_p, c_p, *other_args):
    # See Unit.scan() for seqs.
    # args: seqs (x_t = unit.xc, z_t, i_t), outputs (# unit.n_act, y_p, c_p, ...), non_seqs (none)
    other_outputs = []
    #att_p = theano.printing.Print('att in lstms', attrs=['__str__'])(att_p)
    if self.recurrent_transform:
      state_vars = other_args[:len(self.recurrent_transform.state_vars)]
      self.recurrent_transform.set_sorted_state_vars(state_vars)
      z_r, r_updates = self.recurrent_transform.step(y_p)
      z_t += z_r
      for v in self.recurrent_transform.get_sorted_state_vars():
        other_outputs += [r_updates[v]]
    maxatt = att_p.repeat(z_t.shape[1]).reshape((z_t.shape[0],z_t.shape[1]))#.dimshuffle(1,0)
    #maxatt = theano.printing.Print('maxatt',attrs=['__str__','shape'])(maxatt)
    z_t = T.switch(maxatt>0,z_t,z_t + T.dot(y_p, self.W_re))
    #z_t += T.dot(y_p, self.W_re)
    #z_t = theano.printing.Print('z_t lstms',attrs=['shape'])(z_t)

    partition = z_t.shape[1] // 4
    ingate = T.nnet.sigmoid(z_t[:,:partition])
    forgetgate = ((T.nnet.sigmoid(z_t[:,partition:2*partition])).T * (1.-att_p)).T
    outgate = T.nnet.sigmoid(z_t[:,2*partition:3*partition])
    input = T.tanh(z_t[:,3*partition:4*partition])
    #c_t = ((forgetgate * c_p + ingate * input).T * (1.-T.max(att_p,axis=-1))).T
    c_t = forgetgate * c_p + ingate * input
    y_t = outgate * T.tanh(c_t)
    i_output = T.outer(i_t, self.o_output)
    i_h = T.outer(i_t, self.o_h)
    # return: next outputs (# unit.n_act, y_t, c_t, ...)
    return (y_t * i_output, c_t * i_h + c_p * (1 - i_h)) + tuple(other_outputs)

 def full(self, X, Xs=None):
     X, Xs = self._slice(X, Xs)
     scf_x = self.scaling_func(X, self.args)
     if Xs is None:
         return tt.outer(scf_x, scf_x) * self.cov_func(X)
     else:
         scf_xs = self.scaling_func(Xs, self.args)
         return tt.outer(scf_x, scf_xs) * self.cov_func(X, Xs)

 def contrastive_divergence_1(self, v1):
     '''Determine the weight updates according to CD-1'''
     h1 = self.sample_h_given_v(v1)
     v2 = self.sample_v_given_h(h1)
     h2p = self.propup(v2)
     return (T.outer(v1, h1) - T.outer(v2, h2p),
             v1 - v2,
             h1 - h2p)

def image_step_val(Imat, htm1mat, ctm1mat, 
                   Wcnn, Wxi, Whi, bi, Wxf, Whf, bf, 
                   Wxc, Whc, bc, Wxo, Who, bo, Why, by, forbatch):
    xtmat = theano.dot(Imat, Wcnn)
    itmat = sigma(theano.dot(xtmat,Wxi) + theano.dot(htm1mat,Whi) + T.outer(forbatch,bi) )
    ftmat = sigma(theano.dot(xtmat,Wxf) + theano.dot(htm1mat,Whf) + T.outer(forbatch,bf) )
    ctmat = ftmat * ctm1mat + itmat*act(theano.dot(xtmat,Wxc)+theano.dot(htm1mat,Whc)+T.outer(forbatch,bc) )
    otmat = sigma(theano.dot(xtmat,Wxo) + theano.dot(htm1mat,Who) + T.outer(forbatch,bo) )
    htmat = otmat * act(ctmat)
#    yt = T.concatenate([addzero,tempyt],axis=0)
    return htmat, ctmat    

def psb(inverse_hessian, weight_delta, gradient_delta, **options):
    gradient_delta_t = gradient_delta.T
    param = weight_delta - inverse_hessian.dot(gradient_delta)

    devider = (1. / T.dot(gradient_delta, gradient_delta))
    param1 = T.outer(param, gradient_delta) + T.outer(gradient_delta, param)
    param2 = (
        T.dot(gradient_delta, param) *
        T.outer(gradient_delta, gradient_delta_t)
    )

    return inverse_hessian + param1 * devider - param2 * devider ** 2

def train():
    train_set, valid_set, test_set = loadData()
    x,y = train_set
    m,n_input = x.shape
    width = 28
    height = 28
    n_hidden = 49
    
    learning_rate = .1
    
    #set up shared variables
    W = theano.shared(numpy.random.uniform(-4 * numpy.sqrt(6. / (n_hidden + n_input)),4 * numpy.sqrt(6. / (n_hidden + n_input)),(n_hidden,width*height)),name="W")
    b_v = theano.shared(numpy.zeros((width*height,)),name="b_v")
    b_h = theano.shared(numpy.zeros((n_hidden,)),name="b_h")
    
    theano_rng = T.shared_randomstreams.RandomStreams(numpy.random.randint(2 ** 30))
    
    v_input = T.fvector("v_input")
    
    #1. sample hidden units
    h_prob = T.nnet.sigmoid(T.dot(v_input,W.T)+b_h)
    h_sample = theano_rng.binomial(size=(n_hidden,), n=1, p=h_prob)
    #2. calculate positive gradient
    g_p = T.outer(v_input,h_sample)
    #3. make reconstruction
    v_prob_reconstruction = T.nnet.sigmoid(T.dot(h_sample,W)+b_v)
    v_reconstruction = theano_rng.binomial(size=(n_input,), n=1, p=v_prob_reconstruction)
    h_prob_reconstruction = T.nnet.sigmoid(T.dot(v_reconstruction,W.T)+b_h)
    h_reconstruction = theano_rng.binomial(size=(n_hidden,), n=1, p=h_prob_reconstruction)
    #4. calculate negative gradient
    g_n = T.outer(v_reconstruction,h_reconstruction)
    #FUNCTIONS FOR TESTING
    #f_h_prob = theano.function(inputs=[v_input,],outputs=[h_prob,])
    #f_h_sample = theano.function(inputs=[v_input,],outputs=[h_sample,])
    #f_g_p = theano.function(inputs=[v_input,],outputs=[g_p,])
    #f_v_prob_reconstruction = theano.function(inputs=[v_input,],outputs=[v_prob_reconstruction,])
    #f_v_reconstruction = theano.function(inputs=[v_input,],outputs=[v_reconstruction,])
    #f_h_prob_reconstruction = theano.function(inputs=[v_input,],outputs=[h_prob_reconstruction,])
    #f_h_reconstruction = theano.function(inputs=[v_input,],outputs=[h_reconstruction,])
    #f_g_n = theano.function(inputs=[v_input,],outputs=[g_n,])
    
    learn = theano.function(inputs=[v_input,],updates=[(W,W+learning_rate*(g_p-g_n).T)])
    
    for i in range(300001):
        if i > 0:
            if i%10000 == 0:
                print "Epcoh: ",i
                display_weights(W,width,height,i)
        learn(x[i%m,:])
    
    with open('weights.pkl', 'wb') as output:
        pickle.dump(W.get_value(), output, pickle.HIGHEST_PROTOCOL)

        def times_reflection(input, n_hidden, reflection):
            input_re = input[:, :n_hidden]
            input_im = input[:, n_hidden:]
            reflect_re = reflection[n_hidden:]
            reflect_im = reflection[:n_hidden]
            
            vstarv = (reflect_re**2 + reflect_im**2).sum()
            input_re_reflect = input_re - 2 / vstarv * (T.outer(T.dot(input_re, reflect_re), reflect_re) +
                                                        T.outer(T.dot(input_im, reflect_im), reflect_im))
            input_im_reflect = input_im - 2 / vstarv * (-T.outer(T.dot(input_re, reflect_im), reflect_im) +
                                                        T.outer(T.dot(input_im, reflect_re), reflect_re))

            return T.concatenate([input_re_reflect, input_im_reflect], axis=1)      

 def _step(x_t, i_t, c_tm1, y_tm1):
   #z_t = T.dot(x_t, W) + T.dot(y_tm1, V_h) + b
   z_t = x_t + T.dot(y_tm1, V_h)
   partition = z_t.shape[1] / 4
   ingate = T.nnet.sigmoid(z_t[:,:partition])
   forgetgate = T.nnet.sigmoid(z_t[:,partition:2*partition])
   outgate = T.nnet.sigmoid(z_t[:,2*partition:3*partition])
   input = T.tanh(z_t[:,3*partition:4*partition])
   c_t = forgetgate * c_tm1 + ingate * input
   y_t = outgate * T.tanh(c_t)
   i_output = T.outer(i_t, o_output)
   i_h = T.outer(i_t, o_h)
   return c_t * i_h + c_tm1 * (1 - i_h), y_t * i_output

 def __init__(self, C, D):
     self.W = theano.shared(np.ones((C,D), dtype='float32'))
     t_M = T.matrix('M', dtype='float32')
     t_vM = T.vector('M', dtype='float32')
     t_Y = T.vector('Y', dtype='float32')
     t_I = T.vector('I', dtype='float32')
     t_s = T.vector('s', dtype='float32')
     t_eps = T.scalar('epsilon', dtype='float32')
     self.input_integration = theano.function(
         [t_Y],
         T.dot(T.log(self.W),t_Y),
         allow_input_downcast=True
         )
     self.M_summation = theano.function(
         [t_M],
         T.sum(t_M, axis=0),
         allow_input_downcast=True
         )
     self.recurrent_softmax = theano.function(
         [t_I,t_vM],
         t_vM*T.exp(t_I)/T.sum(t_vM*T.exp(t_I)),
         allow_input_downcast=True
         )
     self.weight_update = theano.function(
         [t_Y,t_s,t_eps],
         self.W,
         updates={
             self.W:
             self.W + t_eps*(T.outer(t_s,t_Y) - t_s[:,np.newaxis]*self.W)
             },
         allow_input_downcast=True
         )
     self.epsilon = None
     self._Y = None
     self._s = None

 def __init__(self, C, D, use_unlabeled):
     self.W = theano.shared(np.ones((C,D), dtype='float32'))
     t_eps = T.scalar('epsilon', dtype='float32')
     t_Y = T.vector('Y', dtype='float32')
     t_s = T.vector('s', dtype='float32')
     self.activation_unlabeled = theano.function(
         [t_Y],
         T.sum(t_Y*self.W/T.sum(self.W, axis=0), axis=1),
         allow_input_downcast=True
         )
     self.activation_normalization = theano.function(
         [t_s],
         t_s/T.sum(t_s),
         allow_input_downcast=True
         )
     self.weight_update = theano.function(
         [t_Y,t_s,t_eps],
         self.W,
         updates={
             self.W:
             self.W + t_eps*(T.outer(t_s,t_Y) - t_s[:,np.newaxis]*self.W)
             },
         allow_input_downcast=True
         )
     self.epsilon = None
     self._Y = None
     self._s = None
     self._delta = np.eye(C, dtype='float32')
     self._C = C
     self._use_unlabeled = use_unlabeled
     self._skipupdate = False

    def learningstep_m1(self, Y, L, M, W, epsilon):
        """Perform a single learning step.

        This is a faster learning step for the case of
        mini-batch-size = 1.

        Keyword arguments:
        the keyword arguments must be the same as given in
        self.input_parameters(mode) for mode='train'.
        """
        # Input integration:
        I = T.dot(T.log(W),Y)
        # recurrent term:
        vM = theano.ifelse.ifelse(
            T.eq(L,-1), # if no label is provided
            T.sum(M, axis=0),
            M[L,:]
            )
        # numeric trick to prevent overflow in the exp-function:
        max_exponent = 88. - T.log(I.shape[0]).astype('float32')
        scale = theano.ifelse.ifelse(T.gt(I[T.argmax(I)], max_exponent),
            I[T.argmax(I)] - max_exponent, 0.)
        # activation: recurrent softmax with overflow protection
        s = vM*T.exp(I-scale)/T.sum(vM*T.exp(I-scale))
        s.name = 's_%d.%d[t]'%(self._nmultilayer,self._nlayer)
        # weight update
        W_new = W + epsilon*(T.outer(s,Y) - s[:,np.newaxis]*W)
        W_new.name = 'W_%d.%d[t]'%(self._nmultilayer,self._nlayer)
        return s, W_new

 def one_iter(W_i, V_i, b_i, a, v_lt_i, p_lt_i, log_likelihood):
     h_i = self.sigmoid(a)
     p_i = self.sigmoid(T.dot(h_i, V_i) + b_i)
     v_i = 1. * (theano_rng.uniform([num_samples]) <= p_i)
     log_likelihood += v_i * T.log(p_i) + (1 - v_i) * T.log(1 - p_i)
     a += T.outer(v_i, W_i)
     return a, v_i, p_i, log_likelihood

    def grad(self, inputs, output_gradients):
        """
        Reverse-mode gradient updates for matrix solve operation c = A \ b.

        Symbolic expression for updates taken from [1]_.

        References
        ----------
        ..[1] M. B. Giles, "An extended collection of matrix derivative results
          for forward and reverse mode automatic differentiation",
          http://eprints.maths.ox.ac.uk/1079/

        """
        A, b = inputs
        c = self(A, b)
        c_bar = output_gradients[0]
        trans_map = {
            'lower_triangular': 'upper_triangular',
            'upper_triangular': 'lower_triangular'
        }
        trans_solve_op = Solve(
            # update A_structure and lower to account for a transpose operation
            A_structure=trans_map.get(self.A_structure, self.A_structure),
            lower=not self.lower
        )
        b_bar = trans_solve_op(A.T, c_bar)
        # force outer product if vector second input
        A_bar = -tensor.outer(b_bar, c) if c.ndim == 1 else -b_bar.dot(c.T)
        if self.A_structure == 'lower_triangular':
            A_bar = tensor.tril(A_bar)
        elif self.A_structure == 'upper_triangular':
            A_bar = tensor.triu(A_bar)
        return [A_bar, b_bar]

def get_square_norm_gradients_scan(D_by_layer, cost, accum = 0):

    # This returns a theano variable that will be of shape (minibatch_size, ).
    # It will contain, for each training example, the associated square-norm of the total gradient.
    # If you take the element-wise square-root afterwards, you will get
    # the associated 2-norms, which is what you want for importance sampling.

    for (layer_name, D) in D_by_layer.items():

        backprop_output = tensor.grad(cost, D['output'])

        if D.has_key('weight'):
            A = D['input']
            B = backprop_output
            S, _ =  theano.scan(fn=lambda A, B: tensor.sqr(tensor.outer(A,B)).sum(),
                                        sequences=[A,B])
            accum = accum + S

        if D.has_key('bias'):

            B = backprop_output
            S, _ =  theano.scan(fn=lambda B: tensor.sqr(B).sum(),
                                        sequences=[B])
            accum = accum + S
        
    return accum

def compute_psi1(lls, lsf, xmean, xvar, z):

    if xmean.ndim == 1:
        xmean = xmean[ None, : ]

    ls = T.exp(lls)
    sf = T.exp(lsf)
    lspxvar = ls + xvar
    constterm1 = ls / lspxvar
    constterm2 = T.prod(T.sqrt(constterm1), 1)
    r2_psi1 = T.outer(T.sum(xmean * xmean / lspxvar, 1), T.ones_like(z[ : , 0 : 1 ])) \
        - np.float32(2) * T.dot(xmean / lspxvar, T.transpose(z)) + \
        T.dot(np.float32(1.0) / lspxvar, T.transpose(z)**2)
    psi1 = sf * T.outer(constterm2, T.ones_like(z[ : , 0 : 1 ])) * T.exp(-np.float32(0.5) * r2_psi1)

    return psi1

    def free_energy(self, visible):
        """Computes the free energy of the model.

        :type visible: theano.tensor.TensorType
        :param visible: The state of the visible units (either 1/0, or mean -
            not important).
        
        :rtype: theano.tensor.var.TensorVariable
        :returns: The free energy of the model, given the visible activation.
            Computed as 

            .. math::
               :label: free_energy

                \mathcal{F}(x) = - \log \sum_h e^{-E(x,h)}
        """
        print 'Running free energy.'

        D = TT.sum(visible, axis=1)
        exponent_term = TT.dot(visible, self.W) + TT.outer(D, self.b_hidden)
                        # TT.outer(D, self.b_hidden)
                        # D is a coefficient, b_hidden should

        hidden_term = TT.sum(TT.log(1 + TT.exp(exponent_term)), axis=1)

        # This is the other and more crucial difference between an RBM and a
        #  RSM: multiplying hidedn bias by "document length".
        b_visible_term = TT.dot(visible, self.b_visible)

        free_energy = - hidden_term - b_visible_term
        return free_energy

def kron(a, b):
    """ Kronecker product

    Same as scipy.linalg.kron(a, b).

    :note: numpy.kron(a, b) != scipy.linalg.kron(a, b)!
        They don't have the same shape and order when
        a.ndim != b.ndim != 2.

    :param a: array_like
    :param b: array_like
    :return: array_like with a.ndim + b.ndim - 2 dimensions.

    """
    a = tensor.as_tensor_variable(a)
    b = tensor.as_tensor_variable(b)
    if (a.ndim + b.ndim <= 2):
        raise TypeError('kron: inputs dimensions must sum to 3 or more. '
                        'You passed %d and %d.' % (a.ndim, b.ndim))
    o = tensor.outer(a, b)
    o = o.reshape(tensor.concatenate((a.shape, b.shape)),
                  a.ndim + b.ndim)
    shf = o.dimshuffle(0, 2, 1, * range(3, o.ndim))
    if shf.ndim == 3:
        shf = o.dimshuffle(1, 0, 2)
        o = shf.flatten()
    else:
        o = shf.reshape((o.shape[0] * o.shape[2],
                         o.shape[1] * o.shape[3]) +
                        tuple([o.shape[i] for i in range(4, o.ndim)]))
    return o

 def forward(self):
     z = self.z0  # sxd
     u = self.u_   # d
     w = self.w_   # d
     b = self.b   # .
     h = self.h   # f
     # h(sxd \dot d + .)  = s
     if not self.batched:
         hwz = h(z.dot(w) + b)  # s
         # sxd + (s \outer d) = sxd
         z1 = z + tt.outer(hwz,  u)  # sxd
         return z1
     else:
         z = z.swapaxes(0, 1)
         # z bxsxd
         # u bxd
         # w bxd
         b = b.dimshuffle(0, 'x')
         # b bx-
         hwz = h(tt.batched_dot(z, w) + b)  # bxs
         # bxsxd + (bxsx- * bx-xd) = bxsxd
         hwz = hwz.dimshuffle(0, 1, 'x')  # bxsx-
         u = u.dimshuffle(0, 'x', 1)  # bx-xd
         z1 = z + hwz * u  # bxsxd
         return z1.swapaxes(0, 1)  # sxbxd

    def compute_probabilistic_matrix(self,X, y, num_cases, k=5):

        z       = T.dot(X, self.A) #Transform x into z space 
        dists   = T.sqr(dist2hy(z,z))
        dists   = T.extra_ops.fill_diagonal(dists, T.max(dists)+1)
        nv      = T.min(dists,axis=1) # value of nearest neighbour 
        dists   = (dists.T - nv).T
        d       = T.extra_ops.fill_diagonal(dists, 0)
   
        #Take only k nearest 
        num     = T.zeros((num_cases, self.num_classes))
        denom   = T.zeros((num_cases,))
        for c_i in xrange(self.num_classes):

            #Mask for class i
            mask_i = T.eq(T.outer(T.ones_like(y),y),c_i)

            #K nearest neighbour within a class i 
            dim_ci = T.sum(mask_i[0])
            d_c_i = T.reshape(d[mask_i.nonzero()],(num_cases,dim_ci))
            k_indice = T.argsort(d_c_i, axis=1)[:,0:k]
            
            kd = T.zeros((num_cases,k))
            for it in xrange(k):
                kd = T.set_subtensor(kd[:,it], d_c_i[T.arange(num_cases),k_indice[:,it]]) 

            #Numerator
            value   = T.exp(-T.mean(kd,axis=1))
            num     = T.set_subtensor(num[:,c_i], value) 
            denom   += value 
            

        p = num / denom.dimshuffle(0,'x')    #prob that point i will be correctly classified    
        return p

 def one_iter(Wi, Vi, bi, rand_i, a, vis_i, post):
     hid  = self.sigmoid(a)
     pi   = self.sigmoid(T.dot(hid, Vi) + bi)
     vis_i = T.cast(rand_i <= pi, floatX)
     post  = post + T.log(pi*vis_i + (1-pi)*(1-vis_i))
     a     = a + T.outer(vis_i, Wi)
     return a, vis_i, post