def compute_cost_log_in_parallel(original_rnn_outputs, labels, func, x_ends, y_ends):
	mask = T.log(1 - T.or_(T.eq(labels, T.zeros_like(labels)), T.eq(labels, shift_matrix(labels, 2))))

	initial_state = T.log(T.zeros_like(labels))
	initial_state = T.set_subtensor(initial_state[:,0], 0)

	def select_probabilities(rnn_outputs, label):
		return rnn_outputs[:,label]	

	rnn_outputs, _ = theano.map(select_probabilities, [original_rnn_outputs, labels])
	rnn_outputs = T.log(rnn_outputs.dimshuffle((1,0,2)))

	def forward_step(probabilities, last_probabilities):
		all_forward_probabilities = T.stack(
			last_probabilities + probabilities,
			log_shift_matrix(last_probabilities, 1) + probabilities,
			log_shift_matrix(last_probabilities, 2) + probabilities + mask,
		)

		result = func(all_forward_probabilities, 0)
		return result

	forward_probabilities, _ = theano.scan(fn = forward_step, sequences = rnn_outputs, outputs_info = initial_state)
	forward_probabilities = forward_probabilities.dimshuffle((1,0,2))

	def compute_cost(forward_probabilities, x_end, y_end):
		return -func(forward_probabilities[x_end-1,y_end-2:y_end])

	return theano.map(compute_cost, [forward_probabilities, x_ends, y_ends])[0]

        def unit(parent_x, child_h, child_c, child_exists):
            (h_i, h_o, h_u), _ = theano.map(
                fn=lambda Ui, Uo, Uu, h, exists:
                    (exists * T.dot(Ui, h), exists * T.dot(Uo, h), exists * T.dot(Uu, h)),
                sequences=[self.U_i, self.U_o, self.U_u, child_h, child_exists])

            i = T.nnet.sigmoid(T.dot(self.W_i, parent_x) + h_i.sum(axis=0) + self.b_i)
            o = T.nnet.sigmoid(T.dot(self.W_o, parent_x) + h_o.sum(axis=0) + self.b_o)
            u = T.tanh(T.dot(self.W_u, parent_x) + h_u.sum(axis=0) + self.b_u)

            def _sub_f(U):
                sub_h_f, _ = theano.map(
                    fn=lambda sub_U, h, exists: exists * T.dot(sub_U, h),
                    sequences=[U, child_h, child_exists])
                return sub_h_f.sum(axis=0)

            h_f, _ = theano.map(
                fn=lambda U: _sub_f(U),
                sequences=[self.U_f])
            f = (T.nnet.sigmoid(
                    T.dot(self.W_f, parent_x).dimshuffle('x', 0) + h_f +
                    self.b_f.dimshuffle('x', 0)) *
                 child_exists.dimshuffle(0, 'x'))

            c = i * u + T.sum(f * child_c, axis=0)
            h = o * T.tanh(c)
            return h, c

    def step(visible, filtered_hidden_mean_m1, filtered_hidden_cov_m1):
        A, B = transition, emission                         # (h, h), (h, v)

        # Shortcuts for the filtered mean and covariance from the previous
        # time step.
        f_m1 = filtered_hidden_mean_m1                      # (n, h)
        F_m1 = filtered_hidden_cov_m1                       # (n, h, h)

        # Calculate mean of joint.
        hidden_mean = T.dot(f_m1, A) + hnm                  # (n, h)

        visible_mean = T.dot(hidden_mean, B) + vnm          # (n, v)

        # Calculate covariance of joint.
        hidden_cov = stacked_dot(
            A.T, stacked_dot(F_m1, A))                      # (n, h, h)

        hidden_cov += hnc

        visible_cov = stacked_dot(                          # (n, v, v)
            B.T, stacked_dot(hidden_cov, B))
        visible_cov += vnc

        visible_hidden_cov = stacked_dot(hidden_cov, B)     # (n, h, v)

        visible_error = visible - visible_mean              # (n, v)

        inv_visible_cov, _ = theano.map(
            lambda x: matrix_inverse(x), visible_cov)       # (n, v, v)

        # I don't know a better name for this monster.
        visible_hidden_cov_T = visible_hidden_cov.dimshuffle(0, 2, 1)   # (n, v, h)
        D = stacked_dot(inv_visible_cov, visible_hidden_cov_T)

        f = (D * visible_error.dimshuffle(0, 1, 'x')        # (n, h)
            ).sum(axis=1)
        f += hidden_mean

        F = hidden_cov
        F -= stacked_dot(visible_hidden_cov, D)

        log_l = (inv_visible_cov *                          # (n,)
            visible_error.dimshuffle(0, 1, 'x') *
            visible_error.dimshuffle(0,'x', 1)).sum(axis=(1, 2))
        log_l *= -.5

        dets, _ = theano.map(lambda x: det(x), visible_cov)

        log_l -= 0.5 * T.log(dets)
        log_l -= np.log(2 * np.pi)

        return f, F, log_l

	def pv_function(self, tensor_input):
		indexf_matrix = theano.shared(
									np.zeros(
										[self.max_length, self.max_length], 
										dtype=np.int32
										),
									name = 'indexf_matrix',
									borrow=True
									)
		
		pf_matrix = theano.shared(
								np.zeros(
										[self.max_length, self.max_length], 
										dtype=theano.config.floatX
										),
								name = 'pf_matrix',
								borrow=True
								)
		pf_matrix = T.set_subtensor(pf_matrix[0, 0:tensor_input.shape[0]], 1.0)
		
		vf_matrix = theano.shared(
								np.zeros(
										(self.max_length, self.max_length, self.size), 
										dtype=theano.config.floatX
										),
								name = 'vf_matrix',
								borrow=True
								)
		results, updates = theano.map(
				fn = lambda i, L, t_tensor_input: L[t_tensor_input[i]],
				sequences=[T.arange(tensor_input.shape[0])],
				non_sequences=[self.L, tensor_input],
				name = 'vf_matrix prepare'
				)
		vf_matrix = T.set_subtensor(vf_matrix[0, 0:tensor_input.shape[0]], results)
		
		for i in range(1,self.max_length):
			results, updates = theano.map(
				fn = self._pv_function,
				sequences=[T.arange(self.max_length-i)],
				non_sequences = [i, pf_matrix, vf_matrix],
				#name = 'pv function'
				)
			
			indexf_matrix = T.set_subtensor(indexf_matrix[i, 0:self.max_length-i], results[0])
			pf_matrix = T.set_subtensor(pf_matrix[i, 0:self.max_length-i], results[1])
			vf_matrix = T.set_subtensor(vf_matrix[i, 0:self.max_length-i], results[2])
			
		return indexf_matrix, pf_matrix, vf_matrix

    def decode_to_probs(self, activations, relative_position, low_bound, high_bound):
        squashed = T.reshape(activations, (-1,self.RAW_ENCODING_WIDTH))
        n_parallel = squashed.shape[0]
        probs = T.nnet.softmax(squashed)


        def _scan_fn(cprobs, cpos):

            if self.with_artic:
                abs_probs = cprobs[:2]
                rel_probs = cprobs[2:]
            else:
                rel_probs = cprobs
                abs_probs = T.ones((2,))

            aligned = T.roll(rel_probs, (cpos-low_bound)%12)

            num_tile = int(math.ceil((high_bound-low_bound)/self.WINDOW_SIZE))

            tiled = T.tile(aligned, (num_tile,))[:(high_bound-low_bound)]

            full = T.concatenate([abs_probs, tiled], 0)
            return full

        # probs = theano.printing.Print("probs",['shape'])(probs)
        # relative_position = theano.printing.Print("relative_position",['shape'])(relative_position)
        from_scan, _ = theano.map(fn=_scan_fn, sequences=[probs, T.flatten(relative_position)])
        # from_scan = theano.printing.Print("from_scan",['shape'])(from_scan)
        newshape = T.concatenate([activations.shape[:-1],[2+high_bound-low_bound]],0)
        fixed = T.reshape(from_scan, newshape, ndim=activations.ndim)
        return fixed

    def compute_tree(self, emb_x, tree):
        self.recursive_unit = self.create_recursive_unit()
        self.leaf_unit = self.create_leaf_unit()
        num_nodes = tree.shape[0]  # num internal nodes
        num_leaves = self.num_words - num_nodes

        # compute leaf hidden states
        leaf_h, _ = theano.map(
            fn=self.leaf_unit,
            sequences=[emb_x[:num_leaves]])

        # use recurrence to compute internal node hidden states
        def _recurrence(cur_emb, node_info, t, node_h, last_h):
            child_exists = node_info > -1
            child_h = node_h[node_info - child_exists * t] * child_exists.dimshuffle(0, 'x')
            parent_h = self.recursive_unit(cur_emb, child_h, child_exists)
            node_h = T.concatenate([node_h,
                                    parent_h.reshape([1, self.hidden_dim])])
            return node_h[1:], parent_h

        dummy = theano.shared(self.init_vector([self.hidden_dim]))
        (_, parent_h), _ = theano.scan(
            fn=_recurrence,
            outputs_info=[leaf_h, dummy],
            sequences=[emb_x[num_leaves:], tree, T.arange(num_nodes)],
            n_steps=num_nodes)

        return T.concatenate([leaf_h, parent_h], axis=0)

 def getSample(self, Y, nSamp = 1):    
     def get_layers(ii):
         output = lasagne.layers.get_output(lasagne.layers.get_all_layers(self.sbn_nn), inputs = Y)
         return output[::-1]
     
     output,_ = theano.map(get_layers, T.arange(nSamp))
     return output

    def get_reward_sequences(self, env_state_sessions, agent_action_sessions):
        """Computes the rewards given to agent at each time step for each batch.

        :param env_state_sessions: Environment state [batch_i,seq_i,state_units] history for all sessions.
        :type env_state_sessions: theano tensor [batch_i,seq_i,state_units]

        :param agent_action_sessions: Actions chosen by agent at each tick for all sessions.
        :type agent_action_sessions: int[batch_i,seq_i]

        :return rewards: What reward was given to an agent for corresponding action from state in that batch.
        :rtype: float[batch_i,seq_i]
        """
        env_state_sessions = check_list(env_state_sessions)
        n_states = len(env_state_sessions)
        agent_action_sessions = check_list(agent_action_sessions)
        n_actions = len(agent_action_sessions)

        def compute_reward(batch_i, *args):
            session_states, session_actions = unpack_list(args, [n_states, n_actions])
            return self.get_reward(session_states, session_actions, batch_i)

        sequences = [T.arange(agent_action_sessions[0].shape[0], ), ] + env_state_sessions + agent_action_sessions

        rewards, updates = theano.map(compute_reward, sequences=sequences)

        assert len(updates) == 0
        return rewards.reshape(agent_action_sessions[0].shape)  # reshape bach to original

def cosine_similarity(x, y, eps=1e-6):
    r"""
    Cosine similarity between a vector and each row of a base matrix.

    Parameters
    ----------
    x: a 1D Theano variable
        Vector to compare to each row of the matrix y.
    y: a 2D Theano variable
        Matrix to be compared to
    eps: float
        Precision of the operation (necessary for differentiability).

    Return
    ------
    z: a 1D Theano variable
        A vector whose components are the cosine similarities
        between x and each row of y.
    """
    def _cosine_similarity(x, y, eps=1e-6):
        y = y.dimshuffle(1, 0)
        z = T.dot(x, y)
        z /= T.sqrt(T.sum(x * x) * T.sum(y * y, axis=0) + eps)

        return z

    def step(x_b, y_b):
        return _cosine_similarity(x_b, y_b, eps)
    z, _ = theano.map(step, sequences=[x, y])

    return z

    def get_reward_sequences(self,env_state_sessions,agent_action_sessions):
        """
        computes the rewards given to agent at each time step for each batch
        parameters:
            env_state_seq - environment state [batch_i,seq_i,state_units] history for all sessions
            agent_action_seq - int[batch_i,seq_i]
        returns:
            rewards float[batch_i,seq_i] - what reward was given to an agent for corresponding action from state in that batch

        """
        
        
        
        def compute_reward(batch_i,session_states,session_actions):
            return self.get_reward(session_states,session_actions,batch_i)



        sequences = [
            T.arange(env_state_sessions.shape[0],),
            env_state_sessions,
            agent_action_sessions,
        ]

        rewards,updates = theano.map(compute_reward,
                              sequences=sequences)
        assert len(updates)==0
        return rewards.reshape(agent_action_sessions.shape) #reshape bach to original

def theano_scan_color(writer, draw_fn):
	with writer as writer_buf:
		writer_buf_reshaped = writer_buf.reshape((Screen.screen_vane_count, Screen.screen_max_magnitude, 3))
		vane_matrix = [[[float(vane), float(vane), float(vane)] for px in range(Screen.screen_max_magnitude)]
					   for vane in range(Screen.screen_vane_count)]
		px_matrix =   [[[float(px),float(px),float(px)] for px in range(Screen.screen_max_magnitude)]
					   for vane in range(Screen.screen_vane_count)]
		col_matrix =  [[[float(0), float(1), float(2)] for px in range(Screen.screen_max_magnitude)]
					   for vane in range(Screen.screen_vane_count)]
		vane_vec = T.as_tensor(vane_matrix)
		px_vec = T.as_tensor(px_matrix)
		col_vec = T.as_tensor(col_matrix)
		step = T.fscalar('step')

		draw_fn_with_step = draw_fn(step)
		f, _ = theano.map(draw_fn_with_step, [vane_vec, px_vec, col_vec])

		fn_actual = theano.function([step], f, allow_input_downcast=True, on_unused_input='ignore')

		step_actual = 0
		while True:
			writer.frame_ready()
			start = time.time()
			writer_buf_reshaped[:] = fn_actual(step_actual)
			step_actual -= 1
			done = time.time()
			fps = 1.0/(done - start)
			if fps < TARGET_FPS:
				logging.warning('Frame rate is %f, which is lower than target %d', fps, TARGET_FPS)

 def attend(self, y_p):
   updates = self.default_updates()
   for g in range(self.attrs['glimpse']):
     for i in range(len(self.base)-1,-1,-1):
       factor = T.constant(self.base[i].attrs['factor'][0], 'int32') if i > 0 else 1
       B, C, I, h_p, _ = self.get(y_p, i, g)
       if i == len(self.base) - 1:
         z_i = self.distance(C, h_p)
       else:
         length = T.cast(T.max(T.sum(I,axis=0))+1,'int32')
         ext = T.cast(T.minimum(ext/factor,T.min(length)),'int32')
         def pick(i_t, ext):
           pad = T.minimum(i_t+ext, B.shape[0]) - ext
           return T.concatenate([T.zeros((pad,), 'int8'), T.ones((ext,), 'int8'), T.zeros((B.shape[0]-pad-ext+1,), 'int8')], axis=0)
         idx, _ = theano.map(pick, sequences = [pos/factor], non_sequences = [ext])
         idx = (idx.dimshuffle(1,0)[:-1].flatten() > 0).nonzero()
         C = C.reshape((C.shape[0]*C.shape[1],C.shape[2]))[idx].reshape((ext,C.shape[1],C.shape[2]))
         z_i = self.distance(C, h_p)
         I = I.reshape((I.shape[0]*I.shape[1],))[idx].reshape((ext,I.shape[1]))
       if i > 0:
         pos = T.argmax(self.softmax(z_i,I),axis=0) * factor
         ext = factor
       else:
         w_i = self.softmax(z_i,I)
     B = B.reshape((B.shape[0]*B.shape[1],B.shape[2]))[idx].reshape((ext,B.shape[1],B.shape[2]))
     proto = T.sum(B * w_i.dimshuffle(0,1,'x').repeat(B.shape[2],axis=2),axis=0)
     for i in range(len(self.base)):
       self.glimpses[i].append(proto)
   return T.dot(proto, self.custom_vars['W_att_in_0']), updates

    def test_function5(self):

        w = theano.shared(1.0, name="w")


        def joke(a, b):
            k = w * a
            # g = 0.01 * T.grad((k - 1)**2, w)
            return k, {w: w - 1.0}

        x = T.dscalar("x")
        hs, _ = theano.scan(joke, sequences=[np.array([1.0, 2.0, 3.0])], outputs_info=[np.float64(1.0)] )

        print hs, _

        def upd(h):
            return T.grad(hs[h], w)

        gs, up = theano.map(upd, sequences=[T.arange(hs.shape[0])])
        print gs, up
        # print hs, _
        # print gs, up

        func = theano.function(inputs=[], outputs=gs, updates= [])
        print func()
        print w.get_value()

def gaussian_filter_2d_variable_sigma(input, sigmas,
                                      window_radius=None,
                                      border_mode='zero'
                                      ):
    def filter_sigma(idx, kernel):
        dimpattern_w = ('x', 'x', 'x', 0)
        dimpattern_h = ('x', 'x', 0, 'x')
        filter_w = kernel.dimshuffle(dimpattern_w)
        blur_w = T.nnet.conv2d(
            padded_input[idx:idx+1], filter_w,
            border_mode=_get_chained_w_h_conv_border(conv_border, 'w'),
            filter_shape=[1, 1, 1, None])
        filter_h = kernel.dimshuffle(dimpattern_h)
        return T.nnet.conv2d(
            blur_w, filter_h,
            border_mode=_get_chained_w_h_conv_border(conv_border, 'h'),
            filter_shape=[1, 1, None, 1])

    ndim = 4
    assert input.ndim == ndim, \
        "there must be {} dimensions, got {}".format(ndim, input.ndim)
    window_radius = gaussian_kernel_default_radius(sigmas, window_radius)
    padded_input, conv_border = add_border(input, window_radius, border_mode)
    kernel = gaussian_kernel_1d(sigmas, window_radius)
    blur, _ = theano.map(
        filter_sigma,
        sequences=[T.arange(sigmas.shape[0]), kernel])
    return blur.reshape(input.shape)

 def connect(self, S):
     self.S = S
     def step(s_current, h_prev):
         h_t = self.activation(
             T.dot(s_current, self.W_ih) + 
             T.dot(h_prev, self.W_hh)
         )
         y_t = self.activation(
             T.dot(h_t, self.W_ho)
         )
         return h_t, y_t
     [self.H, self.output], _ = theano.scan(
         step,
         sequences = self.S,
         outputs_info = [self.h_init, None]
     )
     
     self.prediction, _ = theano.map(
         lambda x: T.argmax(x),
         sequences = self.output
     )
     self.final_state = self.H[self.H.shape[0] - 1]
     self.outputter = theano.function([self.S], self.output)
     self.predicter = theano.function([self.S], self.prediction)
     self.CONNECTED = True

    def compute_objective_and_gradients(self, nSamp):
        hsamp = self.mrec.getSample(self.Y, nSamp)

        # evaluate the generative model density P_\theta(y_i , h_i)
        p_yh,_ = theano.map(self.mprior.evaluateLogDensity, sequences=hsamp)
        # evaluate the recognition model density Q_\phi(h_i | y_i)
        q_hgy,_ = theano.map(self.mrec.evalLogDensity, sequences=hsamp)

        ff = (p_yh-q_hgy)
        sortidx = ff.argsort(axis=0)
        
        fmax = ff[(sortidx[-1],T.arange(ff.shape[-1]))].dimshuffle('x',0)
        
        f_hy = T.exp(ff - fmax)
        sum_across_samples = f_hy.sum(axis=0, keepdims = True)
        Lhat = T.log(sum_across_samples/nSamp) + fmax
         
        col_idx = T.arange(ff.shape[-1])
                                                                          # This 1e-12 constant is for debugging nans
                                                                          # in other parts of code. We know we'll get 
                                                                          # nans where we'll then overwrite. Use it with 
        # compute cross-validated estimates of Lhat                       # nanguard mode.
        hold_out_except_last = T.log((sum_across_samples - f_hy)/(nSamp-1)) + fmax #+1e-12) + fmax
        f2max_vec = ff[(sortidx[-2],T.arange(ff.shape[-1]))]
        f2max = f2max_vec.dimshuffle('x',0)
        # Do tricky things to keep the numerics in order (avoid a term being \approxeq 0)
        ff_nolast = T.set_subtensor(ff[(sortidx[-1],col_idx)], f2max_vec)
        f_hy_last = T.exp(ff_nolast - f2max)
        # compute held-out sum when we hold out the maximum element
        hold_out_last = T.log((f_hy_last.sum(axis=0, keepdims=True) - f_hy_last)/(nSamp-1)) + f2max    
        # compute final held-out estimates
        hold_out = T.set_subtensor(hold_out_except_last[(sortidx[-1],col_idx)], hold_out_last[(sortidx[-1],col_idx)])
                        
        Lhat_cv = Lhat - hold_out 
        the_ws = f_hy / sum_across_samples

        weighted_q = T.sum((Lhat_cv*q_hgy + the_ws*ff).mean(axis=1))
        #weighted_q = T.sum((Lhat_cv*q_hgy + the_ws*(p_yh-q_hgy)).sum(axis=1))

        # gradients for approximate posterior
        dqhgy = T.grad(cost=weighted_q, wrt = self.mrec.getParams(), consider_constant=([the_ws,Lhat_cv]+hsamp), disconnected_inputs='ignore')
        
        # gradients for prior
        dpyh = T.grad(cost=T.sum((the_ws*ff).mean(axis=1)), wrt = self.mprior.getParams(), consider_constant=hsamp + [the_ws], disconnected_inputs='ignore')        
        #dpyh = T.grad(cost=T.sum((the_ws*(p_yh-q_hgy)).sum(axis=1)), wrt = self.mprior.getParams(), consider_constant=hsamp + [the_ws], disconnected_inputs='ignore')        
        
        return [Lhat.mean(), dpyh, dqhgy]

 def __call__(self,X):
     #out = self.W[:,X]
     def step(x):
         return self.W[:x]
     stk = theano.map( lambda x: self.W[x],X)
     out = T.stacklists(stk[0])
     #return out.dimshuffle('x','x',0,1)
     return out

 def _score_ao_tg(self, tag_ids, word_ids):
     output, _ = theano.map(fn=order0_ll_score_given_word_and_tag, 
                            sequences=[tag_ids, word_ids], 
                            non_sequences=[self._tg_lp_tag_np_table,
                                           self._tg_tag_emb,
                                           self._tg_word_emb],
                            name="_score_ao_tg")
     return T.sum(output)

    def call(self, inputs, mask=None):
        l1 = inputs[0]
        l2 = inputs[1]

        def f(i, l1, l2):
            return T.clip(T.batched_tensordot(l1[i], l2[i], 1), FLOAT_MIN, FLOAT_MAX).astype(FLOATX)

        return theano.map(f, T.arange(l1.shape[0]), non_sequences=[l1, l2])[0]

def hessian_diag1(f, v):
    g = gradient1(f, v)
    idx = tt.arange(g.shape[0], dtype='int32')

    def hess_ii(i):
        return gradient1(g[i], v)[i]

    return theano.map(hess_ii, idx)[0]