Python tensorflow.cholesky() Examples

def build_model(hps, kernel, z_pos, x, n_particles, full_cov=False):
    """
    Build the SVGP model.
    Note that for inference, we only need the diagonal part of Cov[Y], as
    ELBO equals sum over individual observations.
    For visualization etc we may want a full covariance. Thus the argument
    `full_cov`.
    """
    bn = zs.BayesianNet()
    Kzz_chol = tf.cholesky(kernel(z_pos, z_pos))
    fz = bn.multivariate_normal_cholesky(
        'fz', tf.zeros([hps.n_z], dtype=hps.dtype), Kzz_chol,
        n_samples=n_particles)
    # f(X)|f(Z) follows GP(0, K) gp_conditional
    fx_given_fz = bn.stochastic(
        'fx', gp_conditional(z_pos, fz, x, full_cov, kernel, Kzz_chol))
    # Y|f(X) ~ N(f(X), noise_level * I)
    noise_level = tf.get_variable(
        'noise_level', shape=[], dtype=hps.dtype,
        initializer=tf.constant_initializer(0.05))
    noise_level = tf.nn.softplus(noise_level)
    bn.normal('y', mean=fx_given_fz, std=noise_level, group_ndims=1)
    return bn 

def conditional_ND(self, Xnew, full_cov=False):
        ## modified from GPR
        Kx = self.kern.K(self._X_mean, Xnew)
        K = self.kern.K(self._X_mean) + tf.eye(tf.shape(self._X_mean)[0], dtype=settings.float_type) * self._lik_variance
        L = tf.cholesky(K)
        A = tf.matrix_triangular_solve(L, Kx, lower=True)
        V = tf.matrix_triangular_solve(L, self._Y - self.mean_function(self._X_mean))
        fmean = tf.matmul(A, V, transpose_a=True) + self.mean_function(Xnew)
        if full_cov:
            fvar = self.kern.K(Xnew) - tf.matmul(A, A, transpose_a=True)
            shape = tf.stack([1, 1, tf.shape(self._Y)[1]])
            fvar = tf.tile(tf.expand_dims(fvar, 2), shape)
        else:
            fvar = self.kern.Kdiag(Xnew) - tf.reduce_sum(tf.square(A), 0)
            fvar = tf.tile(tf.reshape(fvar, (-1, 1)), [1, tf.shape(self._Y)[1]])
        return fmean, fvar 

def __init__(self, kern, X, num_outputs, mean_function, input_prop_dim=None, **kwargs):
        """
        A dense layer with fixed inputs. NB X does not change here, and must be the inputs. Minibatches not possible
        """
        Layer.__init__(self, input_prop_dim, **kwargs)
        self.num_data = X.shape[0]
        q_mu = np.zeros((self.num_data, num_outputs))
        self.q_mu = Parameter(q_mu)
        self.q_mu.prior = Gaussian_prior(0., 1.)
        self.kern = kern
        self.mean_function = mean_function

        self.num_outputs = num_outputs

        Ku = self.kern.compute_K_symm(X) + np.eye(self.num_data) * settings.jitter
        self.Lu = tf.constant(np.linalg.cholesky(Ku))
        self.X = tf.constant(X) 

def build_likelihood_terms(self):
        Kdiag = self.kern.Kdiag(self.X)
        Kuu = make_Kuu(self.kern, self.a, self.b, self.ms)
        sigma2 = self.likelihood.variance

        # Compute intermediate matrices
        P = self.KufKfu / sigma2 + Kuu.get()
        L = tf.cholesky(P)
        log_det_P = tf.reduce_sum(tf.log(tf.square(tf.diag_part(L))))
        c = tf.matrix_triangular_solve(L, self.KufY) / sigma2

        # compute log marginal bound
        ND = tf.cast(tf.size(self.Y), float_type)
        D = tf.cast(tf.shape(self.Y)[1], float_type)
        return (-0.5 * ND * tf.log(2 * np.pi * sigma2),
                -0.5 * D * log_det_P,
                0.5 * D * Kuu.logdet(),
                -0.5 * self.tr_YTY / sigma2,
                0.5 * tf.reduce_sum(tf.square(c)),
                -0.5 * tf.reduce_sum(Kdiag)/sigma2,
                0.5 * Kuu.trace_KiX(self.KufKfu) / sigma2) 

def __init__(self, x, y, kern, mean_function=gpflow.mean_functions.Zero(),
                 scale=1., name='GPRCached'):
        """Initialize GP and cholesky decomposition."""
        # Make sure gpflow is imported
        if not isinstance(gpflow, ModuleType):
            raise gpflow

        # self.scope_name = scope.original_name_scope
        gpflow.gpr.GPR.__init__(self, x, y, kern, mean_function, name)

        # Create new dataholders for the cached data
        # TODO zero-dim dataholders cause strange allocator errors in
        # tensorflow with MKL
        dtype = config.np_dtype
        self.cholesky = gpflow.param.DataHolder(np.empty((0, 0), dtype=dtype),
                                                on_shape_change='pass')
        self.alpha = gpflow.param.DataHolder(np.empty((0, 0), dtype=dtype),
                                             on_shape_change='pass')
        self._scale = scale
        self.update_cache() 

def _define_full_covariance_probs(self, shard_id, shard):
    """Defines the full covariance probabilties per example in a class.

    Updates a matrix with dimension num_examples X num_classes.

    Args:
      shard_id: id of the current shard.
      shard: current data shard, 1 X num_examples X dimensions.
    """
    diff = shard - self._means
    cholesky = tf.cholesky(self._covs + self._min_var)
    log_det_covs = 2.0 * tf.reduce_sum(tf.log(tf.matrix_diag_part(cholesky)), 1)
    x_mu_cov = tf.square(
        tf.matrix_triangular_solve(
            cholesky, tf.transpose(
                diff, perm=[0, 2, 1]), lower=True))
    diag_m = tf.transpose(tf.reduce_sum(x_mu_cov, 1))
    self._probs[shard_id] = -0.5 * (
        diag_m + tf.to_float(self._dimensions) * tf.log(2 * np.pi) +
        log_det_covs) 

def build_backward_variance(self, Yvar):
        """
        Additional method for scaling variance backward (used in :class:`.Normalizer`). Can process both the diagonal
        variances returned by predict_f, as well as full covariance matrices.

        :param Yvar: size N x N x P or size N x P
        :return: Yvar scaled, same rank and size as input
        """
        rank = tf.rank(Yvar)
        # Because TensorFlow evaluates both fn1 and fn2, the transpose can't be in the same line. If a full cov
        # matrix is provided fn1 turns it into a rank 4, then tries to transpose it as a rank 3.
        # Splitting it in two steps however works fine.
        Yvar = tf.cond(tf.equal(rank, 2), lambda: tf.matrix_diag(tf.transpose(Yvar)), lambda: Yvar)
        Yvar = tf.cond(tf.equal(rank, 2), lambda: tf.transpose(Yvar, perm=[1, 2, 0]), lambda: Yvar)

        N = tf.shape(Yvar)[0]
        D = tf.shape(Yvar)[2]
        L = tf.cholesky(tf.square(tf.transpose(self.A)))
        Yvar = tf.reshape(Yvar, [N * N, D])
        scaled_var = tf.reshape(tf.transpose(tf.cholesky_solve(L, tf.transpose(Yvar))), [N, N, D])
        return tf.cond(tf.equal(rank, 2), lambda: tf.reduce_sum(scaled_var, axis=1), lambda: scaled_var) 

def _build_inverse_projection(self, X, W, mu, std):
        n, d_latent = self.shape
        variance = tf.square(std)

        M = tf.matmul(tf.transpose(W), W) + tf.diag(tf.ones((d_latent))*variance)
        L = tf.cholesky(M)
        
        # pred_Z = M^-1 W' r' as per (6) in tipping & bishop JRSS
        # https://www.microsoft.com/en-us/research/wp-content/uploads/2016/02/bishop-ppca-jrss.pdf
        r = X-mu
        WTr = tf.transpose(tf.matmul(r, W)) # W' (t - mu)  
        tmp = tf.matrix_triangular_solve(L, WTr, lower=True)
        pred_z = tf.transpose(tf.matrix_triangular_solve(tf.transpose(L), tmp, lower=False))        
        
        return pred_z, L, std 

def _to_normal2d(output_batch) -> ds.MultivariateNormalTriL:
    """
    :param output_batch: (n_samples, 5)
    :return
    """

    # mean of x and y
    x_mean = Lambda(lambda o: o[:, 0])(output_batch)
    y_mean = Lambda(lambda o: o[:, 1])(output_batch)

    # std of x and y
    # std is must be 0 or positive
    x_std = Lambda(lambda o: K.exp(o[:, 2]))(output_batch)
    y_std = Lambda(lambda o: K.exp(o[:, 3]))(output_batch)

    # correlation coefficient
    # correlation coefficient range is [-1, 1]
    cor = Lambda(lambda o: K.tanh(o[:, 4]))(output_batch)

    loc = Concatenate()([
        Lambda(lambda x_mean: K.expand_dims(x_mean, 1))(x_mean),
        Lambda(lambda y_mean: K.expand_dims(y_mean, 1))(y_mean)
    ])

    x_var = Lambda(lambda x_std: K.square(x_std))(x_std)
    y_var = Lambda(lambda y_std: K.square(y_std))(y_std)
    xy_cor = Multiply()([x_std, y_std, cor])

    cov = Lambda(lambda inputs: K.stack(inputs, axis=0))(
        [x_var, xy_cor, xy_cor, y_var])
    cov = Lambda(lambda cov: K.permute_dimensions(cov, (1, 0)))(cov)
    cov = Reshape((2, 2))(cov)

    scale_tril = Lambda(lambda cov: tf.cholesky(cov))(cov)
    mvn = ds.MultivariateNormalTriL(loc, scale_tril)

    return mvn 

def _build_likelihood_terms(self):
        Kdiag = reduce(tf.multiply, [k.Kdiag(self.X[:, i:i+1]) for i, k in enumerate(self.kerns)])
        Kuu = [make_Kuu(k, a, b, self.ms) for k, a, b, in zip(self.kerns, self.a, self.b)]
        Kuu_solid = kron([Kuu_d.get() for Kuu_d in Kuu])
        Kuu_inv_solid = kron([Kuu_d.inv().get() for Kuu_d in Kuu])
        sigma2 = self.likelihood.variance

        # Compute intermediate matrices
        P = self.KufKfu / sigma2 + Kuu_solid
        L = tf.cholesky(P)
        log_det_P = tf.reduce_sum(tf.log(tf.square(tf.diag_part(L))))
        c = tf.matrix_triangular_solve(L, self.KufY) / sigma2

        Kuu_logdets = [K.logdet() for K in Kuu]
        N_others = [float(np.prod(self.Ms)) / M for M in self.Ms]
        Kuu_logdet = reduce(tf.add, [N*logdet for N, logdet in zip(N_others, Kuu_logdets)])

        # compute log marginal bound
        ND = tf.cast(tf.size(self.Y), float_type)
        D = tf.cast(tf.shape(self.Y)[1], float_type)
        return (-0.5 * ND * tf.log(2 * np.pi * sigma2),
                -0.5 * D * log_det_P,
                0.5 * D * Kuu_logdet,
                -0.5 * self.tr_YTY / sigma2,
                0.5 * tf.reduce_sum(tf.square(c)),
                -0.5 * tf.reduce_sum(Kdiag)/sigma2,
                0.5 * tf.reduce_sum(Kuu_inv_solid * self.KufKfu) / sigma2) 

def _build_predict(self, Xnew, full_cov=False):
        Kuu = [make_Kuu(k, a, b, self.ms) for k, a, b, in zip(self.kerns, self.a, self.b)]
        Kuu_solid = kron([Kuu_d.get() for Kuu_d in Kuu])
        Kuu_inv_solid = kron([Kuu_d.inv().get() for Kuu_d in Kuu])
        sigma2 = self.likelihood.variance

        # Compute intermediate matrices
        P = self.KufKfu / sigma2 + Kuu_solid
        L = tf.cholesky(P)
        c = tf.matrix_triangular_solve(L, self.KufY) / sigma2

        Kus = [make_Kuf(k, Xnew[:, i:i+1], a, b, self.ms)
               for i, (k, a, b) in enumerate(zip(self.kerns, self.a, self.b))]
        Kus = tf.transpose(make_kvs([tf.transpose(Kus_d) for Kus_d in Kus]))
        tmp = tf.matrix_triangular_solve(L, Kus)
        mean = tf.matmul(tf.transpose(tmp), c)
        KiKus = tf.matmul(Kuu_inv_solid, Kus)
        if full_cov:
            raise NotImplementedError
        else:
            var = reduce(tf.multiply, [k.Kdiag(Xnew[:, i:i+1]) for i, k in enumerate(self.kerns)])
            var += tf.reduce_sum(tf.square(tmp), 0)
            var -= tf.reduce_sum(KiKus * Kus, 0)
            shape = tf.stack([1, tf.shape(self.Y)[1]])
            var = tf.tile(tf.expand_dims(var, 1), shape)
        return mean, var 

def _build_likelihood(self):
        num_inducing = tf.size(self.ms)
        num_data = tf.shape(self.Y)[0]
        output_dim = tf.shape(self.Y)[1]

        err = self.Y - self.mean_function(self.X)
        Kdiag = self.kern.Kdiag(self.X)
        Kuf = make_Kuf(self.kern, self.X, self.a, self.b, self.ms)
        Kuu = make_Kuu(self.kern, self.a, self.b, self.ms)
        Kuu = Kuu.get()
        sigma = tf.sqrt(self.likelihood.variance)

        # Compute intermediate matrices
        L = tf.cholesky(Kuu)
        A = tf.matrix_triangular_solve(L, Kuf) / sigma
        AAT = tf.matmul(A, tf.transpose(A))

        B = AAT + tf.eye(num_inducing * 2 - 1, dtype=float_type)
        LB = tf.cholesky(B)
        log_det_B = 2. * tf.reduce_sum(tf.log(tf.diag_part(LB)))
        c = tf.matrix_triangular_solve(LB, tf.matmul(A, err)) / sigma

        # compute log marginal bound
        ND = tf.cast(num_data * output_dim, float_type)
        D = tf.cast(output_dim, float_type)
        bound = -0.5 * ND * tf.log(2 * np.pi * self.likelihood.variance)
        bound += -0.5 * D * log_det_B
        bound += -0.5 * tf.reduce_sum(tf.square(err))/self.likelihood.variance
        bound += 0.5 * tf.reduce_sum(tf.square(c))
        bound += -0.5 * tf.reduce_sum(Kdiag)/self.likelihood.variance
        bound += 0.5 * tf.reduce_sum(tf.diag_part(AAT))

        return bound 

def _build_predict(self, Xnew, full_cov=False):
        num_inducing = tf.size(self.ms)

        err = self.Y - self.mean_function(self.X)
        Kuf = make_Kuf(self.kern, self.X, self.a, self.b, self.ms)
        Kuu = make_Kuu(self.kern, self.a, self.b, self.ms)
        Kuu = Kuu.get()
        sigma = tf.sqrt(self.likelihood.variance)

        # Compute intermediate matrices
        L = tf.cholesky(Kuu)
        A = tf.matrix_triangular_solve(L, Kuf) / sigma
        AAT = tf.matmul(A, tf.transpose(A))

        B = AAT + tf.eye(num_inducing * 2 - 1, dtype=float_type)
        LB = tf.cholesky(B)
        c = tf.matrix_triangular_solve(LB, tf.matmul(A, err)) / sigma

        Kus = make_Kuf(self.kern, Xnew, self.a, self.b, self.ms)
        tmp1 = tf.matrix_triangular_solve(L, Kus, lower=True)
        tmp2 = tf.matrix_triangular_solve(LB, tmp1, lower=True)
        mean = tf.matmul(tf.transpose(tmp2), c)
        if full_cov:
            var = self.kern.K(Xnew) + \
                tf.matmul(tf.transpose(tmp2), tmp2) - \
                tf.matmul(tf.transpose(tmp1), tmp1)
            var = var[:, :, None] * tf.ones(self.Y.shape[1], dtype=float_type)
        else:
            var = self.kern.Kdiag(Xnew) + \
                tf.reduce_sum(tf.square(tmp2), 0) - \
                tf.reduce_sum(tf.square(tmp1), 0)
            var = var[:, None]#  * tf.ones(self.Y.shape[1], dtype=float_type)
        return mean + self.mean_function(Xnew), var 

def orthonorm_op(x, epsilon=1e-7):
    '''
    Computes a matrix that orthogonalizes the input matrix x

    x:      an n x d input matrix
    eps:    epsilon to prevent nonzero values in the diagonal entries of x

    returns:    a d x d matrix, ortho_weights, which orthogonalizes x by
                right multiplication
    '''
    x_2 = K.dot(K.transpose(x), x)
    x_2 += K.eye(K.int_shape(x)[1])*epsilon
    L = tf.cholesky(x_2)
    ortho_weights = tf.transpose(tf.matrix_inverse(L)) * tf.sqrt(tf.cast(tf.shape(x)[0], dtype=K.floatx()))
    return ortho_weights 

def _elbo_helper(self):
        # compute kernel stuff
        kern = self.kern; f = self.q_mu;
        num_data = tf.shape(self.Z)[0]  # M
        num_func = tf.shape(f)[1]  # K
        Kmn = kern.Kzx(self.Z, self.X)
        Kmm = kern.Kzz(self.Z) + tf.eye(num_data, dtype=float_type) * settings.numerics.jitter_level
        Lm = tf.cholesky(Kmm)
        # Compute the projection matrix A
        A = tf.matrix_triangular_solve(Lm, Kmn, lower=True)
        # compute the covariance due to the conditioning
        fvar = kern.Kdiag_tr() - tf.reduce_sum(tf.square(A), 0)
        shape = tf.stack([num_func, 1])
        fvar = tf.tile(tf.expand_dims(fvar, 0), shape)  # K x N x N or K x N
        # another backsubstitution in the unwhitened case
        if not self.whiten:
            A = tf.matrix_triangular_solve(tf.transpose(Lm), A, lower=False)
        # construct the conditional mean
        fmean = tf.matmul(A, f, transpose_a=True)
        if self.q_sqrt is not None:
            if self.q_sqrt.get_shape().ndims == 2:
                LTA = A * tf.expand_dims(tf.transpose(self.q_sqrt), 2)  # K x M x N
            elif self.q_sqrt.get_shape().ndims == 3:
                L = tf.matrix_band_part(tf.transpose(self.q_sqrt, (2, 0, 1)), -1, 0)  # K x M x M
                A_tiled = tf.tile(tf.expand_dims(A, 0), tf.stack([num_func, 1, 1]))
                LTA = tf.matmul(L, A_tiled, transpose_a=True)  # K x M x N
            else:  # pragma: no cover
                raise ValueError("Bad dimension for q_sqrt: %s" %
                                 str(self.q_sqrt.get_shape().ndims))
            fvar = fvar + tf.reduce_sum(tf.square(LTA), 1)  # K x N
        fvar = tf.transpose(fvar)  # N x K or N x N x K
        return fmean + self.mean_function(self.X), fvar 

def _sample(self, mean, cov):
        L = tf.cholesky(cov)
        eps = tf.random_normal(shape=self.shape, dtype=self.dtype)
        return tf.matmul(L, eps) + mean 

def _sample_and_entropy(self, mean, cov, **kwargs):
        L = tf.cholesky(cov)
        eps = tf.random_normal(shape=self.shape, dtype=self.dtype)
        sample = tf.matmul(L, eps) + mean
        entropy = util.dists.multivariate_gaussian_entropy(L=L)
        return sample, entropy 

def __init__(self, mean, cov, d=None):

        mean = tf.convert_to_tensor(mean)
        cov = tf.convert_to_tensor(cov)
        
        try:
            d1, = util.extract_shape(mean)
            mean = tf.reshape(mean, (d1,1))
        except:
            d1,k = util.extract_shape(mean)
            assert(k == 1)

        d2,_ = util.extract_shape(cov)
        assert(d1==d2)
        if d is None:
            d = d1
        else:
            assert(d==d1)
            
        super(MVGaussianMeanCov, self).__init__(d=d)
        
        self._mean = mean
        self._cov = cov
        
        self._L_cov = tf.cholesky(cov)
        self._entropy = util.dists.multivariate_gaussian_entropy(L=self._L_cov)

        L_prec_transpose = util.triangular_inv(self._L_cov)
        self._L_prec = tf.transpose(L_prec_transpose)
        self._prec = tf.matmul(self._L_prec, L_prec_transpose)
        self._prec_mean = tf.matmul(self._prec, self._mean) 

def __init__(self, prec_mean, prec, d=None):

        prec_mean = tf.convert_to_tensor(prec_mean)
        prec = tf.convert_to_tensor(prec)
        
        try:
            d1, = util.extract_shape(prec_mean)
            prec_mean = tf.reshape(prec_mean, (d1,1))
        except:
            d1,k = util.extract_shape(prec_mean)
            assert(k == 1)

            
        d2,_ = util.extract_shape(prec)
        assert(d1==d2)
        if d is None:
            d = d1
        else:
            assert(d==d1)

        super(MVGaussianNatural, self).__init__(d=d)

        self._prec_mean = prec_mean
        self._prec = prec
        
        self._L_prec = tf.cholesky(prec)
        self._entropy = util.dists.multivariate_gaussian_entropy(L_prec=self._L_prec)

        # want to solve prec * mean = prec_mean for mean.
        # this is equiv to (LL') * mean = prec_mean.
        # since tf doesn't have a cholSolve shortcut, just
        # do it directly:
        #   solve L y = prec_mean
        # to get y = (L' * mean), then
        #   solve L' mean = y
        y = tf.matrix_triangular_solve(self._L_prec, self._prec_mean, lower=True, adjoint=False)
        self._mean = tf.matrix_triangular_solve(self._L_prec, y, lower=True, adjoint=True)

        L_cov_transpose = util.triangular_inv(self._L_prec)
        self._L_cov = tf.transpose(L_cov_transpose)
        self._cov = tf.matmul(self._L_cov, L_cov_transpose) 

def _logp(self, result, Z, mu, std):
        n, d_output = self.shape
        
        cov = tf.matmul(Z, tf.transpose(Z)) + tf.diag(tf.ones(n,)*std) 
        L = tf.cholesky(cov)
        r = result - mu
        out_cols = tf.unpack(r, axis=1)
        lps = [multivariate_gaussian_log_density(col, mu=0., L=L) for col in out_cols]
        return tf.reduce_sum(tf.pack(lps)) 

def _build_predict(self, Xnew, full_cov=False):
        Kuu = [make_Kuu(k, ai, bi, self.ms) for k, ai, bi in zip(self.kerns, self.a, self.b)]
        Kuu = BlockDiagMat_many([mat for k in Kuu for mat in [k.A, k.B]])
        sigma2 = self.likelihood.variance

        # Compute intermediate matrices
        P = self.KufKfu / sigma2 + Kuu.get()
        L = tf.cholesky(P)
        c = tf.matrix_triangular_solve(L, self.KufY) / sigma2

        Kus = tf.concat([make_Kuf(k, Xnew[:, i:i+1], a, b, self.ms)
                         for i, (k, a, b) in enumerate(zip(self.kerns, self.a, self.b))], axis=0)
        tmp = tf.matrix_triangular_solve(L, Kus)
        mean = tf.matmul(tf.transpose(tmp), c)
        KiKus = Kuu.solve(Kus)
        if full_cov:
            var = reduce(tf.add, [k.K(Xnew[:, i:i+1]) for i, k in enumerate(self.kerns)])
            var += tf.matmul(tf.transpose(tmp), tmp)
            var -= tf.matmul(tf.transpose(KiKus), Kus)
            shape = tf.stack([1, 1, tf.shape(self.Y)[1]])
            var = tf.tile(tf.expand_dims(var, 2), shape)
        else:
            var = reduce(tf.add, [k.Kdiag(Xnew[:, i:i+1]) for i, k in enumerate(self.kerns)])
            var += tf.reduce_sum(tf.square(tmp), 0)
            var -= tf.reduce_sum(KiKus * Kus, 0)
            shape = tf.stack([1, tf.shape(self.Y)[1]])
            var = tf.tile(tf.expand_dims(var, 1), shape)
        return mean, var 

def _sample(self, prior_mean, prior_cov,
                 transition_mat, transition_mean, transition_cov,
                 observation_mat=None, observation_mean=None, observation_cov=None):

        transition_mean = tf.reshape(transition_mean, shape=(self.D, 1))
        transition_eps = tf.random_normal(shape=(self.T, self.D))
        transition_epses = [tf.reshape(n, shape=(self.D, 1)) for n in tf.unpack(transition_eps)]

        prior_mean = tf.reshape(prior_mean, shape=(self.D, 1))
        prior_cov_L = tf.cholesky(prior_cov)
        state = prior_mean + tf.matmul(prior_cov_L, transition_epses[0])

        transition_cov_L = tf.cholesky(transition_cov)
        if not self._flag_no_obs:
            observation_cov_L = tf.cholesky(observation_cov)
            obs_eps = tf.random_normal(shape=(self.T, self.K))
            obs_epses = [tf.reshape(n, shape=(self.K, 1)) for n in tf.unpack(obs_eps)]
            observation_mean = tf.reshape(observation_mean, shape=(self.K, 1))
            
        output = []
        hidden = []
        for t in range(self.T):
            if not self._flag_no_obs:
                pred_obs = tf.matmul(observation_mat, state) + observation_mean
                output.append(pred_obs + tf.matmul(observation_cov_L, obs_epses[t]))
            else:
                output.append(state)
            hidden.append(state)

            if t < self.T-1:
                state_noise = transition_mean + tf.matmul(transition_cov_L, transition_epses[t+1])
                state = tf.matmul(transition_mat, state) + state_noise

        self._sampled_hidden = hidden
        return tf.pack(tf.squeeze(output)) 

def multivariate_gaussian_entropy(Sigma=None, L=None, L_prec=None):

    if L is None and Sigma is not None:
        L = tf.cholesky(Sigma)
    
    if L is not None:
        half_logdet = tf.reduce_sum(tf.log(tf.diag_part(L)))
        n, _ = extract_shape(L)
    else:
        half_logdet = -tf.reduce_sum(tf.log(tf.diag_part(L_prec)))
        n, _ = extract_shape(L_prec)

    log_2pi = 1.83787706641
    entropy = .5*n*(1 + log_2pi) + half_logdet
    return entropy 

def cholesky_inverse(A):
  """Compute the inverse of `A` using Choselky decomposition. NOTE: `A` must be
  symmetric positive definite. This method of inversion is not completely stable since
  tf.cholesky is not always stable. Might raise `tf.errors.InvalidArgumentError`
  """
  N     = tf.shape(A)[0]
  L     = tf.cholesky(A)
  L_inv = tf.matrix_triangular_solve(L, tf.eye(N))
  A_inv = tf.matmul(L_inv, L_inv, transpose_a=True)
  return A_inv 

def resample_w(self, cholesky=False):
    sample = tf.random_normal(shape=self.w_mu.shape, dtype=self.dtype)

    # Compute A s.t. A A^T = w_Sigma. Note that SVD and Cholesky give different A
    if cholesky:
      # Use cholesky
      A = tf.cholesky(self.w_Sigma)
    else:
      # Use SVD
      S, U, _ = tf.svd(self.w_Sigma)
      A = tf.matmul(U, tf.diag(tf.sqrt(S)))

    w = self.w_mu + tf.matmul(A, sample)
    return tf.assign(self.w, w, name="resample_w") 

def orthonorm_op(x, epsilon=1e-7):
    '''
    Computes a matrix that orthogonalizes the input matrix x

    x:      an n x d input matrix
    eps:    epsilon to prevent nonzero values in the diagonal entries of x

    returns:    a d x d matrix, ortho_weights, which orthogonalizes x by
                right multiplication
    '''
    x_2 = K.dot(K.transpose(x), x)
    x_2 += K.eye(K.int_shape(x)[1])*epsilon
    L = tf.cholesky(x_2)
    ortho_weights = tf.transpose(tf.matrix_inverse(L)) * tf.sqrt(tf.cast(tf.shape(x)[0], dtype=K.floatx()))
    return ortho_weights 

def _build_cholesky(self):
    tf.logging.info("Computing Kernel")
    self.k_data_data_reg = self.k_data_data + tf.eye(
        self.input_x.shape[0], dtype=tf.float64) * self.stability_eps
    if FLAGS.print_kernel:
      self.k_data_data_reg = tf.Print(
          self.k_data_data_reg, [self.k_data_data_reg],
          message="K_DD = ", summarize=100)
    self.l = tf.cholesky(self.k_data_data_reg)
    self.v = tf.matrix_triangular_solve(self.l, self.y_pl) 

def _objective_grad(xx, obj, grad_loss, grad_eigval, grad_chol):
  """Symbolic form of the gradient of the objective with stop_gradients."""
  del grad_eigval
  del grad_chol
  with tf.name_scope('objective_grad'):
    chol = tf.cholesky(xx)
    choli = tf.linalg.inv(chol)
    rq = tf.matmul(choli, tf.matmul(obj, choli, transpose_b=True))

    dl = tf.diag(tf.matrix_diag_part(choli))
    triu = tf.matrix_band_part(tf.matmul(rq, dl), 0, -1)
    gxx = -1.0 * tf.matmul(choli, triu, transpose_a=True)
    gobj = tf.matmul(choli, dl, transpose_a=True)

    return grad_loss * gxx, grad_loss * gobj 

def _objective(xx, obj):
  """Objective function as custom op so that we can overload gradients."""
  with tf.name_scope('objective'):
    chol = tf.cholesky(xx)
    choli = tf.linalg.inv(chol)

    rq = tf.matmul(choli, tf.matmul(obj, choli, transpose_b=True))
    eigval = tf.matrix_diag_part(rq)
    loss = tf.trace(rq)
    grad = functools.partial(_objective_grad, xx, obj)
  return (loss, eigval, chol), grad 

def build_cholesky_if_needed(self):
        # make sure we only compute this once
        if self.needs_build_cholesky:
            self.Ku = self.feature.Kuu(self.kern, jitter=settings.jitter)
            self.Lu = tf.cholesky(self.Ku)
            self.Ku_tiled = tf.tile(self.Ku[None, :, :], [self.num_outputs, 1, 1])
            self.Lu_tiled = tf.tile(self.Lu[None, :, :], [self.num_outputs, 1, 1])
            self.needs_build_cholesky = False