jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def update_step(obs, x, P, H, R, y_sigma, mask):
    """
    Update state x and P after the observation,
    outputs filtered state and covariance
    """
    # create masks
    if mask is not None:
        H = mask_observation_operator(H, mask)
        R = mask_observation_noise_diag(R, mask)

    # Update state mean and covariance p(x | y), Joseph Form
    # Kalman gain, a more stable implementation than naive P @ H^T @ y_sigma^{-1}
    L = torch.cholesky(
        y_sigma + 1e-6 * torch.eye(y_sigma.shape[-1], device=y_sigma.device)
    )
    K = torch.triangular_solve(H @ P.transpose(1, 2), L, upper=False)[0]
    K = torch.triangular_solve(K, L.transpose(1, 2))[0].transpose(1, 2)

    v = obs - batch_mat_predict(H, x)
    x = x + torch.einsum("bso,bo->bs", K, v)

    # P = (torch.eye(*P.shape[1:]) - K @ self.H) @ P @ (torch.eye(*P.shape[1:]) - K @ self.H).T + K @ self.trans_noise @ K.T
    identity = torch.eye(*P.shape[1:], device=P.device)
    t1 = identity - K @ H
    P = t1.matmul(P).matmul(t1.transpose(1, 2)) + K.matmul(R).matmul(K.transpose(1, 2))
    return x, P
//  [python]
def update_step(obs, x, P, H, R, y_sigma, mask):
    """
    Update state x and P after the observation,
    outputs filtered state and covariance
    """
    # create masks
    if mask is not None:
        H = mask_observation_operator(H, mask)
        R = mask_observation_noise_diag(R, mask)

    # Update state mean and covariance p(x | y), Joseph Form
    # Kalman gain, a more stable implementation than naive P @ H^T @ y_sigma^{-1}
    L = torch.linalg.cholesky(
        y_sigma + 1e-6 * torch.eye(y_sigma.shape[-1], device=y_sigma.device)
    )
    K = torch.triangular_solve(H @ P.transpose(1, 2), L, upper=False)[0]
    K = torch.triangular_solve(K, L.transpose(1, 2))[0].transpose(1, 2)

    v = obs - batch_mat_predict(H, x)
    x = x + torch.einsum("bso,bo->bs", K, v)

    # P = (torch.eye(*P.shape[1:]) - K @ self.H) @ P @ (torch.eye(*P.shape[1:]) - K @ self.H).T + K @ self.trans_noise @ K.T
    identity = torch.eye(*P.shape[1:], device=P.device)
    t1 = identity - K @ H
    P = t1.matmul(P).matmul(t1.transpose(1, 2)) + K.matmul(R).matmul(K.transpose(1, 2))
    return x, P
//  [python]
def sample_inputs(self, device, dtype, requires_grad=False):
        """
        This function generates Cholesky factors of positive-definite (non-singular) Hermitian (symmetric) matrices
        for cholesky_inverse.
        """
        from torch.testing._internal.common_utils import random_hermitian_pd_matrix
        inputs = (
            torch.zeros(0, 0, dtype=dtype, device=device),  # 0x0 matrix
            torch.zeros(0, 2, 2, dtype=dtype, device=device),  # zero batch of matrices
            random_hermitian_pd_matrix(S, dtype=dtype, device=device),  # single matrix
            random_hermitian_pd_matrix(S, 2, dtype=dtype, device=device),  # batch of matrices
        )
        test_cases = (torch.cholesky(a) for a in inputs)
        out = []
        for a in test_cases:
            a.requires_grad = requires_grad
            out.append(SampleInput(a))
            out.append(SampleInput(a, kwargs=dict(upper=True)))
        return out
//  [python]
def sample_inputs(self, device, dtype, requires_grad=False):
        """
        This function generates Cholesky factors of positive-definite (non-singular) Hermitian (symmetric) matrices
        for cholesky_inverse.
        """
        from torch.testing._internal.common_utils import random_hermitian_pd_matrix
        inputs = (
            torch.zeros(0, 0, dtype=dtype, device=device),  # 0x0 matrix
            torch.zeros(0, 2, 2, dtype=dtype, device=device),  # zero batch of matrices
            random_hermitian_pd_matrix(S, dtype=dtype, device=device),  # single matrix
            random_hermitian_pd_matrix(S, 2, dtype=dtype, device=device),  # batch of matrices
        )
        test_cases = (torch.linalg.cholesky(a) for a in inputs)
        out = []
        for a in test_cases:
            a.requires_grad = requires_grad
            out.append(SampleInput(a))
            out.append(SampleInput(a, kwargs=dict(upper=True)))
        return out
//  [python]
def logm_invB_A(B, A):
#     inputs: A/B.shape = [h, w, d, 3, 3]
#     output: shape = [h, w, d, 3, 3]
    G = torch.cholesky(B)
    inv_G = torch.inverse(G)
    W = torch.einsum("...ij,...jk,...lk->...il", inv_G, A, inv_G)
    lamda, Q = torch.linalg.eig(W)#, eigenvectors=True
    lamda, Q = lamda.real, Q.real
    log_lamda = torch.diag_embed(torch.log(lamda))
    V = torch.einsum('...ji,...jk->...ik', inv_G, Q)
    inv_V = torch.inverse(V)
    return torch.einsum('...ij,...jk,...kl->...il', V, log_lamda, inv_V)
//  [python]
def logm_invB_A(B, A):
#     inputs: A/B.shape = [h, w, d, 3, 3]
#     output: shape = [h, w, d, 3, 3]
    G = torch.linalg.cholesky(B)
    inv_G = torch.inverse(G)
    W = torch.einsum("...ij,...jk,...lk->...il", inv_G, A, inv_G)
    lamda, Q = torch.linalg.eig(W)#, eigenvectors=True
    lamda, Q = lamda.real, Q.real
    log_lamda = torch.diag_embed(torch.log(lamda))
    V = torch.einsum('...ji,...jk->...ik', inv_G, Q)
    inv_V = torch.inverse(V)
    return torch.einsum('...ij,...jk,...kl->...il', V, log_lamda, inv_V)
//  [python]
def vsgp_multiclass(num_steps, whiten):
    # adapted from http://gpflow.readthedocs.io/en/latest/notebooks/multiclass.html
    pyro.set_rng_seed(0)
    X = torch.rand(100, 1)
    K = (-0.5 * (X - X.t()).pow(2) / 0.01).exp() + torch.eye(100) * 1e-6
    f = torch.cholesky(K).matmul(torch.randn(100, 3))
    y = f.argmax(dim=-1)

    kernel = gp.kernels.Sum(
        gp.kernels.Matern32(1), gp.kernels.WhiteNoise(1, variance=torch.tensor(0.01))
    )
    likelihood = gp.likelihoods.MultiClass(num_classes=3)
    Xu = X[::5].clone()

    gpmodule = gp.models.VariationalSparseGP(
        X, y, kernel, Xu, likelihood, latent_shape=torch.Size([3]), whiten=whiten
    )

    gpmodule.Xu.requires_grad_(False)
    gpmodule.kernel.kern1.variance_unconstrained.requires_grad_(False)

    optimizer = torch.optim.Adam(gpmodule.parameters(), lr=0.0001)
    gp.util.train(gpmodule, optimizer, num_steps=num_steps)
//  [python]
def vsgp_multiclass(num_steps, whiten):
    # adapted from http://gpflow.readthedocs.io/en/latest/notebooks/multiclass.html
    pyro.set_rng_seed(0)
    X = torch.rand(100, 1)
    K = (-0.5 * (X - X.t()).pow(2) / 0.01).exp() + torch.eye(100) * 1e-6
    f = torch.linalg.cholesky(K).matmul(torch.randn(100, 3))
    y = f.argmax(dim=-1)

    kernel = gp.kernels.Sum(
        gp.kernels.Matern32(1), gp.kernels.WhiteNoise(1, variance=torch.tensor(0.01))
    )
    likelihood = gp.likelihoods.MultiClass(num_classes=3)
    Xu = X[::5].clone()

    gpmodule = gp.models.VariationalSparseGP(
        X, y, kernel, Xu, likelihood, latent_shape=torch.Size([3]), whiten=whiten
    )

    gpmodule.Xu.requires_grad_(False)
    gpmodule.kernel.kern1.variance_unconstrained.requires_grad_(False)

    optimizer = torch.optim.Adam(gpmodule.parameters(), lr=0.0001)
    gp.util.train(gpmodule, optimizer, num_steps=num_steps)
//  [python]
def batch_cholesky_inverse(matrix):
    """
	Args: 	matrix, 4-D tensor, [N, C, M, M].
			matrix must be a symmetric positive define matrix.
	"""
    chol_low = torch.linalg.cholesky(matrix)
    # chol_low = torch.cholesky(matrix, upper=False)
    chol_low_inv = batch_low_tri_inv(chol_low)
    return torch.matmul(chol_low_inv.transpose(-2, -1), chol_low_inv)
//  [python]
def variational_loss(self, x, y, noise, mu_hat, A_hat):
        """
        Computes L_H for the data in the current batch.
        :param x: auxiliary data for current batch (batch, 1 + 1 + M)
        :param y: mean vector for current latent channel, output of the encoder network (batch, 1)
        :param noise: variance vector for current latent channel, output of the encoder network (batch, 1)
        :param mu_hat:
        :param A_hat:
        :return: sum_term, KL_term (variational loss = sum_term + KL_term)  (1,)
        """
        b = x.shape[0]
        m = self.inducing_index_points.shape[0]

        K_mm = self.kernel_matrix(self.inducing_index_points, self.inducing_index_points) # (m,m)
        K_mm_inv = torch.linalg.inv(_add_diagonal_jitter(K_mm, self.jitter)) # (m,m)

        K_nn = self.kernel_matrix(x, x, x_inducing=False, y_inducing=False, diag_only=True) # (b)

        K_nm = self.kernel_matrix(x, self.inducing_index_points, x_inducing=False)  # (b, m)
        K_mn = torch.transpose(K_nm, 0, 1)

        S = A_hat

        # KL term
        mean_vector = torch.matmul(K_nm, torch.matmul(K_mm_inv, mu_hat))

        K_mm_chol = torch.cholesky(_add_diagonal_jitter(K_mm, self.jitter))
        S_chol = torch.cholesky(_add_diagonal_jitter(A_hat, self.jitter))
        K_mm_log_det = 2 * torch.sum(torch.log(torch.diagonal(K_mm_chol)))
        S_log_det = 2 * torch.sum(torch.log(torch.diagonal(S_chol)))

        KL_term = 0.5 * (K_mm_log_det - S_log_det - m +
                             torch.trace(torch.matmul(K_mm_inv, A_hat)) +
                             torch.sum(mu_hat * torch.matmul(K_mm_inv, mu_hat)))

        # diag(K_tilde), (b, )
        precision = 1 / noise

        K_tilde_terms = precision * (K_nn - torch.diagonal(torch.matmul(K_nm, torch.matmul(K_mm_inv, K_mn))))

        # k_i \cdot k_i^T, (b, m, m)
        lambda_mat = torch.matmul(K_nm.unsqueeze(2), torch.transpose(K_nm.unsqueeze(2), 1, 2))

        # K_mm_inv \cdot k_i \cdot k_i^T \cdot K_mm_inv, (b, m, m)
        lambda_mat = torch.matmul(K_mm_inv, torch.matmul(lambda_mat, K_mm_inv))

        # Trace terms, (b,)
        trace_terms = precision * torch.einsum('bii->b', torch.matmul(S, lambda_mat))

        # L_3 sum part, (1,)
        L_3_sum_term = -0.5 * (torch.sum(K_tilde_terms) + torch.sum(trace_terms) +
                                torch.sum(torch.log(noise)) + b * np.log(2 * np.pi) +
                                torch.sum(precision * (y - mean_vector) ** 2))

        return L_3_sum_term, KL_term
//  [python]
def variational_loss(self, x, y, noise, mu_hat, A_hat):
        """
        Computes L_H for the data in the current batch.
        :param x: auxiliary data for current batch (batch, 1 + 1 + M)
        :param y: mean vector for current latent channel, output of the encoder network (batch, 1)
        :param noise: variance vector for current latent channel, output of the encoder network (batch, 1)
        :param mu_hat:
        :param A_hat:
        :return: sum_term, KL_term (variational loss = sum_term + KL_term)  (1,)
        """
        b = x.shape[0]
        m = self.inducing_index_points.shape[0]

        K_mm = self.kernel_matrix(self.inducing_index_points, self.inducing_index_points) # (m,m)
        K_mm_inv = torch.linalg.inv(_add_diagonal_jitter(K_mm, self.jitter)) # (m,m)

        K_nn = self.kernel_matrix(x, x, x_inducing=False, y_inducing=False, diag_only=True) # (b)

        K_nm = self.kernel_matrix(x, self.inducing_index_points, x_inducing=False)  # (b, m)
        K_mn = torch.transpose(K_nm, 0, 1)

        S = A_hat

        # KL term
        mean_vector = torch.matmul(K_nm, torch.matmul(K_mm_inv, mu_hat))

        K_mm_chol = torch.linalg.cholesky(_add_diagonal_jitter(K_mm, self.jitter))
        S_chol = torch.linalg.cholesky(_add_diagonal_jitter(A_hat, self.jitter))
        K_mm_log_det = 2 * torch.sum(torch.log(torch.diagonal(K_mm_chol)))
        S_log_det = 2 * torch.sum(torch.log(torch.diagonal(S_chol)))

        KL_term = 0.5 * (K_mm_log_det - S_log_det - m +
                             torch.trace(torch.matmul(K_mm_inv, A_hat)) +
                             torch.sum(mu_hat * torch.matmul(K_mm_inv, mu_hat)))

        # diag(K_tilde), (b, )
        precision = 1 / noise

        K_tilde_terms = precision * (K_nn - torch.diagonal(torch.matmul(K_nm, torch.matmul(K_mm_inv, K_mn))))

        # k_i \cdot k_i^T, (b, m, m)
        lambda_mat = torch.matmul(K_nm.unsqueeze(2), torch.transpose(K_nm.unsqueeze(2), 1, 2))

        # K_mm_inv \cdot k_i \cdot k_i^T \cdot K_mm_inv, (b, m, m)
        lambda_mat = torch.matmul(K_mm_inv, torch.matmul(lambda_mat, K_mm_inv))

        # Trace terms, (b,)
        trace_terms = precision * torch.einsum('bii->b', torch.matmul(S, lambda_mat))

        # L_3 sum part, (1,)
        L_3_sum_term = -0.5 * (torch.sum(K_tilde_terms) + torch.sum(trace_terms) +
                                torch.sum(torch.log(noise)) + b * np.log(2 * np.pi) +
                                torch.sum(precision * (y - mean_vector) ** 2))

        return L_3_sum_term, KL_term
//  [python]
def _call(self, x):
        return torch.cholesky(x)
//  [python]
def _call(self, x):
        return torch.linalg.cholesky(x)
//  [python]
def random_mvt(df_shape, loc_shape, cov_shape, dim):
    """
    Generate a random MultivariateStudentT distribution for testing.
    """
    rank = dim + dim
    df = torch.rand(df_shape, requires_grad=True).exp()
    loc = torch.randn(loc_shape + (dim,), requires_grad=True)
    cov = torch.randn(cov_shape + (dim, rank), requires_grad=True)
    cov = cov.matmul(cov.transpose(-1, -2))
    scale_tril = torch.cholesky(cov)
    return MultivariateStudentT(df, loc, scale_tril)
//  [python]
def random_mvt(df_shape, loc_shape, cov_shape, dim):
    """
    Generate a random MultivariateStudentT distribution for testing.
    """
    rank = dim + dim
    df = torch.rand(df_shape, requires_grad=True).exp()
    loc = torch.randn(loc_shape + (dim,), requires_grad=True)
    cov = torch.randn(cov_shape + (dim, rank), requires_grad=True)
    cov = cov.matmul(cov.transpose(-1, -2))
    scale_tril = torch.linalg.cholesky(cov)
    return MultivariateStudentT(df, loc, scale_tril)
//  [python]
def convert_cbc_terms_to_socp_terms(bfe, e, V, bfv, v, extravars,
                                        testing=True):
        assert (torch.linalg.eigvals(V).real.abs() > 0).all()
        assert v.abs() > 0
        m = bfe.shape[-1]
        bfc = bfe.new_zeros((m+extravars))
        if testing:
            u0 = torch.zeros((m,))
            u0_hom = torch.cat((torch.tensor([1.]), u0))
        with torch.no_grad():
            # [1, u] Asq [1; u]
            Asq = torch.cat(
                (
                    torch.cat((torch.tensor([[v]]),         (bfv / 2).reshape(1, -1)), dim=-1),
                    torch.cat(((bfv / 2).reshape(-1, 1),    V), dim=-1)
                ),
                dim=-2)
            if testing:
                np.testing.assert_allclose(u0_hom @ Asq @ u0_hom,
                                        u0 @ V @ u0 + bfv @ u0 + v)

            # [1, u] Asq [1; u] = |L[1; u]|_2 = |[0, A] [y_1; y_2; u] + b|_2
            n_constraints = m+1
            L = torch.cholesky(Asq) # (m+1) x (m+1)

        if testing:
            np.testing.assert_allclose(L @ L.T, Asq, rtol=1e-2, atol=1e-3)

        A = torch.zeros((n_constraints, m + extravars))
        A[:, extravars:] = L.T[:, 1:]
        bfb = L.T[:, 0] # (m+1)
        if testing:
            y_u0 = torch.cat((torch.zeros(extravars), u0))
            np.testing.assert_allclose(A @ y_u0 + bfb, L.T @ u0_hom)
            np.testing.assert_allclose(u0_hom @ Asq @ u0_hom, u0_hom @ L @ L.T @ u0_hom, rtol=1e-2, atol=1e-3)
            np.testing.assert_allclose(u0 @ V @ u0 + bfv @ u0 + v, (A @ y_u0 + bfb) @ (A @ y_u0 + bfb), rtol=1e-2, atol=1e-3)
        if extravars >= 1: # "I assumed atleast δ "
            bfc[extravars-1] = 1 # For delta the relaxation factor
        bfc[extravars:] = bfe # only affine terms need to be negative?
        d = e
        return A, bfb, bfc, d
//  [python]
def convert_cbc_terms_to_socp_terms(bfe, e, V, bfv, v, extravars,
                                        testing=True):
        assert (torch.linalg.eigvals(V).real.abs() > 0).all()
        assert v.abs() > 0
        m = bfe.shape[-1]
        bfc = bfe.new_zeros((m+extravars))
        if testing:
            u0 = torch.zeros((m,))
            u0_hom = torch.cat((torch.tensor([1.]), u0))
        with torch.no_grad():
            # [1, u] Asq [1; u]
            Asq = torch.cat(
                (
                    torch.cat((torch.tensor([[v]]),         (bfv / 2).reshape(1, -1)), dim=-1),
                    torch.cat(((bfv / 2).reshape(-1, 1),    V), dim=-1)
                ),
                dim=-2)
            if testing:
                np.testing.assert_allclose(u0_hom @ Asq @ u0_hom,
                                        u0 @ V @ u0 + bfv @ u0 + v)

            # [1, u] Asq [1; u] = |L[1; u]|_2 = |[0, A] [y_1; y_2; u] + b|_2
            n_constraints = m+1
            L = torch.linalg.cholesky(Asq) # (m+1) x (m+1)

        if testing:
            np.testing.assert_allclose(L @ L.T, Asq, rtol=1e-2, atol=1e-3)

        A = torch.zeros((n_constraints, m + extravars))
        A[:, extravars:] = L.T[:, 1:]
        bfb = L.T[:, 0] # (m+1)
        if testing:
            y_u0 = torch.cat((torch.zeros(extravars), u0))
            np.testing.assert_allclose(A @ y_u0 + bfb, L.T @ u0_hom)
            np.testing.assert_allclose(u0_hom @ Asq @ u0_hom, u0_hom @ L @ L.T @ u0_hom, rtol=1e-2, atol=1e-3)
            np.testing.assert_allclose(u0 @ V @ u0 + bfv @ u0 + v, (A @ y_u0 + bfb) @ (A @ y_u0 + bfb), rtol=1e-2, atol=1e-3)
        if extravars >= 1: # "I assumed atleast δ "
            bfc[extravars-1] = 1 # For delta the relaxation factor
        bfc[extravars:] = bfe # only affine terms need to be negative?
        d = e
        return A, bfb, bfc, d
//  [python]
def __init__(self, x: Tensor, R: Tensor = None):
        super().__init__()
        self.register_buffer("mu", x.mean(0))
        if R is None:
            R = sample_covariance(x)
        self.register_buffer("L", torch.cholesky(R))
//  [python]
def __init__(self, x: Tensor, R: Tensor = None):
        super().__init__()
        self.register_buffer("mu", x.mean(0))
        if R is None:
            R = sample_covariance(x)
        self.register_buffer("L", torch.linalg.cholesky(R))
//  [python]
def test_mean_var(batch_shape):
    dim = 2
    loc = torch.randn(batch_shape + (dim,))
    A = torch.randn(batch_shape + (dim, dim + dim))
    scale_tril = torch.cholesky(A.matmul(A.transpose(-2, -1)))
    df = torch.randn(batch_shape).exp() + 4
    num_samples = 100000
    d = MultivariateStudentT(df, loc, scale_tril)
    samples = d.sample(sample_shape=(num_samples,))
    expected_mean = samples.mean(0)
    expected_variance = samples.var(0)
    assert_equal(d.mean, expected_mean, prec=0.1)
    assert_equal(d.variance, expected_variance, prec=0.2)

    assert_equal(
        MultivariateStudentT(0.5, loc, scale_tril).mean,
        torch.full(batch_shape + (dim,), float("nan")),
    )
    assert_equal(
        MultivariateStudentT(0.5, loc, scale_tril).variance,
        torch.full(batch_shape + (dim,), float("nan")),
    )
    assert_equal(
        MultivariateStudentT(1.5, loc, scale_tril).variance,
        torch.full(batch_shape + (dim,), float("inf")),
    )
//  [python]
def test_mean_var(batch_shape):
    dim = 2
    loc = torch.randn(batch_shape + (dim,))
    A = torch.randn(batch_shape + (dim, dim + dim))
    scale_tril = torch.linalg.cholesky(A.matmul(A.transpose(-2, -1)))
    df = torch.randn(batch_shape).exp() + 4
    num_samples = 100000
    d = MultivariateStudentT(df, loc, scale_tril)
    samples = d.sample(sample_shape=(num_samples,))
    expected_mean = samples.mean(0)
    expected_variance = samples.var(0)
    assert_equal(d.mean, expected_mean, prec=0.1)
    assert_equal(d.variance, expected_variance, prec=0.2)

    assert_equal(
        MultivariateStudentT(0.5, loc, scale_tril).mean,
        torch.full(batch_shape + (dim,), float("nan")),
    )
    assert_equal(
        MultivariateStudentT(0.5, loc, scale_tril).variance,
        torch.full(batch_shape + (dim,), float("nan")),
    )
    assert_equal(
        MultivariateStudentT(1.5, loc, scale_tril).variance,
        torch.full(batch_shape + (dim,), float("inf")),
    )
//  [python]
def test_log_prob(batch_shape, dim):
    loc = torch.randn(batch_shape + (dim,))
    A = torch.randn(batch_shape + (dim, dim + dim))
    scale_tril = torch.cholesky(A.matmul(A.transpose(-2, -1)))
    x = torch.randn(batch_shape + (dim,))
    df = torch.randn(batch_shape).exp() + 2
    actual_log_prob = MultivariateStudentT(df, loc, scale_tril).log_prob(x)

    if dim == 1:
        expected_log_prob = (
            StudentT(df.unsqueeze(-1), loc, scale_tril[..., 0]).log_prob(x).sum(-1)
        )
        assert_equal(actual_log_prob, expected_log_prob)

    # test the fact MVT(df, loc, scale)(x) = int MVN(loc, scale / m)(x) Gamma(df/2,df/2)(m) dm
    num_samples = 100000
    gamma_samples = Gamma(df / 2, df / 2).sample(sample_shape=(num_samples,))
    mvn_scale_tril = scale_tril / gamma_samples.sqrt().unsqueeze(-1).unsqueeze(-1)
    mvn = MultivariateNormal(loc, scale_tril=mvn_scale_tril)
    expected_log_prob = mvn.log_prob(x).logsumexp(0) - math.log(num_samples)
    assert_equal(actual_log_prob, expected_log_prob, prec=0.01)
//  [python]
def test_log_prob(batch_shape, dim):
    loc = torch.randn(batch_shape + (dim,))
    A = torch.randn(batch_shape + (dim, dim + dim))
    scale_tril = torch.linalg.cholesky(A.matmul(A.transpose(-2, -1)))
    x = torch.randn(batch_shape + (dim,))
    df = torch.randn(batch_shape).exp() + 2
    actual_log_prob = MultivariateStudentT(df, loc, scale_tril).log_prob(x)

    if dim == 1:
        expected_log_prob = (
            StudentT(df.unsqueeze(-1), loc, scale_tril[..., 0]).log_prob(x).sum(-1)
        )
        assert_equal(actual_log_prob, expected_log_prob)

    # test the fact MVT(df, loc, scale)(x) = int MVN(loc, scale / m)(x) Gamma(df/2,df/2)(m) dm
    num_samples = 100000
    gamma_samples = Gamma(df / 2, df / 2).sample(sample_shape=(num_samples,))
    mvn_scale_tril = scale_tril / gamma_samples.sqrt().unsqueeze(-1).unsqueeze(-1)
    mvn = MultivariateNormal(loc, scale_tril=mvn_scale_tril)
    expected_log_prob = mvn.log_prob(x).logsumexp(0) - math.log(num_samples)
    assert_equal(actual_log_prob, expected_log_prob, prec=0.01)
//  [python]
def init_L(self):
        with torch.no_grad():
            kx = self.k(self.Z, self.X).evaluate()
            self.kzz = self.k(self.Z).evaluate()
        L_set = False
        for reg in [1e-7, 1e-6, 1e-5, 1e-4, 1e-3, 1e-2]:  # ,1e-1,1.0,5.0,10.]:
            try:
                with torch.no_grad():
                    chol_L = torch.cholesky(self.kzz + kx @ kx.t() / self.sigma + self.eye * reg)
                    L = torch.cholesky(torch.inverse(chol_L))
                    # L = torch.cholesky(torch.inverse(self.kzz  + self.eye*self.sigma))
                    self.L = torch.nn.Parameter(L)
                    L_set = True
            except Exception as e:
                gc.collect()
                print('cholesky init error: ', reg)
                print(e)
                # torch.cuda.empty_cache()
        if not L_set:
            print("Welp this is awkward, it don't want to factorize")

            gc.collect()
            # torch.cuda.empty_cache()
            L = torch.randn_like(self.kzz) * 0.1
            self.L = torch.nn.Parameter(L)
        inv_set = False
        for reg in [1e-7, 1e-6, 1e-5, 1e-4, 1e-3, 1e-2]:  # , 1e-1, 1.0,2.5,5.0,10.]:
            self.reg = reg
            try:
                if not self.parametrize_Z:
                    with torch.no_grad():
                        self.register_buffer('inv_L', torch.cholesky(self.kzz + self.eye * self.reg))
                        self.register_buffer('inv_Z', torch.inverse(self.inv_L))
                    inv_set = True
            except Exception as e:
                gc.collect()
                print(e)
                print('cholesky inv error: ', reg)
                pass
        if not inv_set:
            gc.collect()
            self.parametrize_Z = True
//  [python]
def init_L(self):
        with torch.no_grad():
            kx = self.k(self.Z, self.X).evaluate()
            self.kzz = self.k(self.Z).evaluate()
        L_set = False
        for reg in [1e-7, 1e-6, 1e-5, 1e-4, 1e-3, 1e-2]:  # ,1e-1,1.0,5.0,10.]:
            try:
                with torch.no_grad():
                    chol_L = torch.linalg.cholesky(self.kzz + kx @ kx.t() / self.sigma + self.eye * reg)
                    L = torch.linalg.cholesky(torch.cholesky_inverse(chol_L))
                    # L = torch.linalg.cholesky(torch.inverse(self.kzz  + self.eye*self.sigma))
                    self.L = torch.nn.Parameter(L)
                    L_set = True
            except Exception as e:
                gc.collect()
                print('cholesky init error: ', reg)
                print(e)
                # torch.cuda.empty_cache()
        if not L_set:
            print("Welp this is awkward, it don't want to factorize")

            gc.collect()
            # torch.cuda.empty_cache()
            L = torch.randn_like(self.kzz) * 0.1
            self.L = torch.nn.Parameter(L)
        inv_set = False
        for reg in [1e-7, 1e-6, 1e-5, 1e-4, 1e-3, 1e-2]:  # , 1e-1, 1.0,2.5,5.0,10.]:
            self.reg = reg
            try:
                if not self.parametrize_Z:
                    with torch.no_grad():
                        self.register_buffer('inv_L', torch.linalg.cholesky(self.kzz + self.eye * self.reg))
                        self.register_buffer('inv_Z', torch.cholesky_inverse(self.inv_L))
                    inv_set = True
            except Exception as e:
                gc.collect()
                print(e)
                print('cholesky inv error: ', reg)
                pass
        if not inv_set:
            gc.collect()
            self.parametrize_Z = True
//  [python]
def compute_loss(self, inputs, outputs, inputs_var=None, outputs_var=None):
        # TODO checking if inputs/outputs was changed
        self.inputs_tr = inputs
        self.outputs_tr = outputs
        self.already_set_train_data = True

        Sigma = self.kernel_list[0](inputs[0], inputs[0]) + JITTER * torch.eye(inputs[0].size(0), device=list(self.parameters())[0].device)
        _noise = self._get_noise_according_exp_format()

        Sigma = Sigma + _noise.pow(-1) * torch.eye(inputs[0].size(0), device=list(self.parameters())[0].device)

        L = torch.cholesky(Sigma)
        #option 1 (use this if torch supports)
        # Gamma,_ = torch.triangular_solve(self.Y, L, upper = False)
        #option 2

        gamma = L.inverse() @ outputs[0]       # we can use this as an alternative because L is a lower triangular matrix.

        y_num, y_dimension = outputs[0].shape
        nll =  0.5 * (gamma ** 2).sum() +  L.diag().log().sum() * y_dimension  \
            + 0.5 * y_num * torch.log(2 * torch.tensor(PI, device=list(self.parameters())[0].device)) * y_dimension
        return nll
//  [python]
def compute_loss(self, inputs, outputs, inputs_var=None, outputs_var=None):
        # TODO checking if inputs/outputs was changed
        self.inputs_tr = inputs
        self.outputs_tr = outputs
        self.already_set_train_data = True

        Sigma = self.kernel_list[0](inputs[0], inputs[0]) + JITTER * torch.eye(inputs[0].size(0), device=list(self.parameters())[0].device)
        _noise = self._get_noise_according_exp_format()

        Sigma = Sigma + _noise.pow(-1) * torch.eye(inputs[0].size(0), device=list(self.parameters())[0].device)

        L = torch.linalg.cholesky(Sigma)
        #option 1 (use this if torch supports)
        # Gamma,_ = torch.triangular_solve(self.Y, L, upper = False)
        #option 2

        gamma = L.inverse() @ outputs[0]       # we can use this as an alternative because L is a lower triangular matrix.

        y_num, y_dimension = outputs[0].shape
        nll =  0.5 * (gamma ** 2).sum() +  L.diag().log().sum() * y_dimension  \
            + 0.5 * y_num * torch.log(2 * torch.tensor(PI, device=list(self.parameters())[0].device)) * y_dimension
        return nll
//  [python]
def predict_with_var(self, inputs, vars=None):
        if self.already_set_train_data is False:
            assert False, "gp model model hasn't been trained. predict failed"
    
        with torch.no_grad():
            Sigma = self.kernel_list[0](self.inputs_tr[0], self.inputs_tr[0]) + JITTER * torch.eye(self.inputs_tr[0].size(0), device=list(self.parameters())[0].device)
            _noise = self._get_noise_according_exp_format()
            Sigma = Sigma + _noise.pow(-1) * torch.eye(self.inputs_tr[0].size(0), device=list(self.parameters())[0].device)
            self.k_l_inv = Sigma

            kx = self.kernel_list[0](self.inputs_tr[0], inputs[0])
            self.kstar = kx
            L = torch.cholesky(Sigma)
            LinvKx,_ = torch.triangular_solve(kx, L, upper = False)

            u = kx.t() @ torch.cholesky_solve(self.outputs_tr[0], L)

            var_diag = self.kernel_list[0](inputs[0], inputs[0]).diag().view(-1, 1) - (LinvKx**2).sum(dim = 0).view(-1, 1)
            _noise = self._get_noise_according_exp_format()
            var_diag = var_diag + _noise.pow(-1)
            var_diag = var_diag.expand_as(u)
        return u, var_diag
//  [python]
def predict_with_var(self, inputs, vars=None):
        if self.already_set_train_data is False:
            assert False, "gp model model hasn't been trained. predict failed"
    
        with torch.no_grad():
            Sigma = self.kernel_list[0](self.inputs_tr[0], self.inputs_tr[0]) + JITTER * torch.eye(self.inputs_tr[0].size(0), device=list(self.parameters())[0].device)
            _noise = self._get_noise_according_exp_format()
            Sigma = Sigma + _noise.pow(-1) * torch.eye(self.inputs_tr[0].size(0), device=list(self.parameters())[0].device)
            self.k_l_inv = Sigma

            kx = self.kernel_list[0](self.inputs_tr[0], inputs[0])
            self.kstar = kx
            L = torch.linalg.cholesky(Sigma)
            LinvKx,_ = torch.triangular_solve(kx, L, upper = False)

            u = kx.t() @ torch.cholesky_solve(self.outputs_tr[0], L)

            var_diag = self.kernel_list[0](inputs[0], inputs[0]).diag().view(-1, 1) - (LinvKx**2).sum(dim = 0).view(-1, 1)
            _noise = self._get_noise_according_exp_format()
            var_diag = var_diag + _noise.pow(-1)
            var_diag = var_diag.expand_as(u)
        return u, var_diag
//  [python]
def prepare(self):
        U = torch.cholesky(self.C_inv, upper=True)
        U_inv = torch.inverse(U)
        tilde_target_inputs = F.linear(self.target_inputs, U)
        Q = torch.qr(tilde_target_inputs.transpose(0, 1))[0].transpose(0, 1)  # orthogonal basis
        Q = F.normalize(Q, dim=1)  # orthonormal basis
        self.target_inputs_basis.data = F.linear(Q, U_inv)
//  [python]
def prepare(self):
        U = torch.linalg.cholesky(self.C_inv, upper=True)
        U_inv = torch.linalg.inv(U)
        tilde_target_inputs = F.linear(self.target_inputs, U)
        Q = torch.linalg.qr(tilde_target_inputs.transpose(0, 1))[0].transpose(0, 1)  # orthogonal basis
        Q = F.normalize(Q, dim=1)  # orthonormal basis
        self.target_inputs_basis.data = F.linear(Q, U_inv)
//  [python]
def __init__(
        self,
        output_size: int,
        a: float,
        b: float,
        n_tau: int,
        kernel: str = "matern52",
        learn_inducing_locations: bool = False,
        output_transform: Callable = Identity(),
    ) -> None:
        super().__init__()
        self.w = Parameter(torch.randn(n_tau, output_size))
        tau = torch.linspace(a, b, n_tau)
        if learn_inducing_locations:
            self.tau = Parameter(tau)
        else:
            self.register_buffer("tau", tau)
        if kernel == "matern52":
            self.K = kernels.MaternKernel(nu=2.5)
        elif kernel == "matern32":
            self.K = kernels.MaternKernel(nu=1.5)
        elif kernel == "rbf":
            self.K = kernels.RBFKernel()
        else:
            raise ValueError(f"Unknown kernel {kernel}")
        self.w = Parameter(
            torch.cholesky(
                self.K(tau, tau).evaluate().inverse() + 1e-4 * torch.eye(len(tau))
            )
            @ torch.randn(n_tau, output_size)
        )
        self.b = Parameter(torch.zeros(output_size))
        self.output_transform = output_transform
//  [python]
def __init__(
        self,
        output_size: int,
        a: float,
        b: float,
        n_tau: int,
        kernel: str = "matern52",
        learn_inducing_locations: bool = False,
        output_transform: Callable = Identity(),
    ) -> None:
        super().__init__()
        self.w = Parameter(torch.randn(n_tau, output_size))
        tau = torch.linspace(a, b, n_tau)
        if learn_inducing_locations:
            self.tau = Parameter(tau)
        else:
            self.register_buffer("tau", tau)
        if kernel == "matern52":
            self.K = kernels.MaternKernel(nu=2.5)
        elif kernel == "matern32":
            self.K = kernels.MaternKernel(nu=1.5)
        elif kernel == "rbf":
            self.K = kernels.RBFKernel()
        else:
            raise ValueError(f"Unknown kernel {kernel}")
        self.w = Parameter(
            torch.linalg.cholesky(
                self.K(tau, tau).evaluate().inverse() + 1e-4 * torch.eye(len(tau))
            )
            @ torch.randn(n_tau, output_size)
        )
        self.b = Parameter(torch.zeros(output_size))
        self.output_transform = output_transform
//  [python]
plt.tight_layout()
plt.savefig('example_hessian.jpg')

eig0 = eigsh(hess.cpu().numpy().reshape(v.numel(), v.numel()), k=1,
             which="SA")[0].item()
tau = max(-1.5 * eig0, 0)
L = torch.cholesky(
    hess.reshape(v.numel(), v.numel()) +
    tau * torch.eye(v.numel(), dtype=dtype, device=device))
//  [python]
plt.tight_layout()
plt.savefig('example_hessian.jpg')

eig0 = eigsh(hess.cpu().numpy().reshape(v.numel(), v.numel()), k=1,
             which="SA")[0].item()
tau = max(-1.5 * eig0, 0)
L = torch.linalg.cholesky(
    hess.reshape(v.numel(), v.numel()) +
    tau * torch.eye(v.numel(), dtype=dtype, device=device))
//  [python]
def __init__(
        self,
        f: Callable,
        L: Union[Callable, Tensor],
        Q: Tensor,
        t_span: Tuple,
        x0: Tensor,
        dt: float,
    ) -> None:
        super().__init__()
        bs = x0.shape[0]
        Q_chol = torch.cholesky(Q) * math.sqrt(dt)
        t = torch.arange(t_span[0], t_span[1] + dt, dt)
        xk = x0.clone()
        x_list = [xk]
        for tk in t[:-1]:
            dbeta = torch.randn(bs, Q_chol.shape[0]) @ Q_chol.t()
            xk = xk + f(tk, xk) * dt
            if callable(L):
                noise_samples = torch.einsum("ijk,ik->ij", L(tk, xk), dbeta)
            else:
                noise_samples = dbeta @ L.t()
            xk += noise_samples
            x_list.append(xk.clone())
        self.x = torch.stack(x_list, dim=1)
        if x0.shape[0] > 1:
            self.mean = self.x.mean(0)
            self.var = sample_covariance(self.x.transpose(1, 0))
        self.t = t
//  [python]
def __init__(
        self,
        f: Callable,
        L: Union[Callable, Tensor],
        Q: Tensor,
        t_span: Tuple,
        x0: Tensor,
        dt: float,
    ) -> None:
        super().__init__()
        bs = x0.shape[0]
        Q_chol = torch.linalg.cholesky(Q) * math.sqrt(dt)
        t = torch.arange(t_span[0], t_span[1] + dt, dt)
        xk = x0.clone()
        x_list = [xk]
        for tk in t[:-1]:
            dbeta = torch.randn(bs, Q_chol.shape[0]) @ Q_chol.t()
            xk = xk + f(tk, xk) * dt
            if callable(L):
                noise_samples = torch.einsum("ijk,ik->ij", L(tk, xk), dbeta)
            else:
                noise_samples = dbeta @ L.t()
            xk += noise_samples
            x_list.append(xk.clone())
        self.x = torch.stack(x_list, dim=1)
        if x0.shape[0] > 1:
            self.mean = self.x.mean(0)
            self.var = sample_covariance(self.x.transpose(1, 0))
        self.t = t
//  [python]
def num_concepts(self, value):
        self._num_concepts = value

        if value == 1 or not self.multi_concepts_use_cholesky:
            return

        # for multiple concepts
        # need cholesky decomposed L_t_inv
        # Appendix B

        try:
            L = torch.cholesky(self.C_inv)
        except:
            print('unable to perform cholesky. please make sure input covariance matrix is properly calculated')
            exit()

        L_T = L.T
        L_T_inv = torch.inverse(L_T)

        self.register_buffer('L_T', L_T, persistent = False)
        self.register_buffer('L_T_inv', L_T_inv, persistent = False)
//  [python]
def num_concepts(self, value):
        self._num_concepts = value

        if value == 1 or not self.multi_concepts_use_cholesky:
            return

        # for multiple concepts
        # need cholesky decomposed L_t_inv
        # Appendix B

        try:
            L = torch.linalg.cholesky(self.C_inv)
        except:
            print('unable to perform cholesky. please make sure input covariance matrix is properly calculated')
            exit()

        L_T = L.T
        L_T_inv = torch.inverse(L_T)

        self.register_buffer('L_T', L_T, persistent = False)
        self.register_buffer('L_T_inv', L_T_inv, persistent = False)
//  [python]
def invert_and_cholesky(x: Tensor) -> Tensor:
    """Inverts and Cholesky decomposes the input."""
    inv, info = torch.linalg.inv_ex(x)
    if info != 0 or not inv.isfinite().all():
        inv = make_invertible_and_invert(x)

    inv_chol, info = torch.linalg.cholesky_ex(inv)
    if info != 0 or not inv_chol.isfinite().all():
        inv = check_and_make_pd(inv)
        inv_chol = torch.cholesky(inv)
    return inv_chol
//  [python]
def invert_and_cholesky(x: Tensor) -> Tensor:
    """Inverts and Cholesky decomposes the input."""
    inv, info = torch.linalg.inv_ex(x)
    if info != 0 or not inv.isfinite().all():
        inv = make_invertible_and_invert(x)

    inv_chol, info = torch.linalg.cholesky_ex(inv)
    if info != 0 or not inv_chol.isfinite().all():
        inv = check_and_make_pd(inv)
        inv_chol = torch.linalg.cholesky(inv)
    return inv_chol
//  [python]
def _condition_RRt(self, RRt, D):
        """
        Args:
            RRt (torch.tensor): of size (N x N) the original full covariance
            D (torch.tensor): of dtype float and size (1xN) containing 1.0 for known indices else 0.0.
        """
        self.register_buffer("D_atom", self._C_atomic(D))
        self.register_buffer("zero_indices", torch.nonzero((1 - self.D_atom[0]))[:, 0])
        self.register_buffer("nonzero_indices", torch.nonzero(self.D_atom[0])[:, 0])

        self.register_buffer("S11", RRt[self.zero_indices][:, self.zero_indices])
        self.register_buffer("S12", RRt[self.zero_indices][:, self.nonzero_indices])
        self.register_buffer("S21", RRt[self.nonzero_indices][:, self.zero_indices])
        self.register_buffer("S22", RRt[self.nonzero_indices][:, self.nonzero_indices])

        S_clamp = self.S11 - ((self.S12 @ torch.linalg.pinv(self.S22) @ self.S21))
        R_clamp = torch.cholesky(S_clamp)
        self.register_buffer("RRt_clamp_restricted", R_clamp @ R_clamp.t())
        RRt_clamp = self._scatter(
            torch.zeros_like(RRt), self.zero_indices, self.RRt_clamp_restricted
        )
        R_clamp = self._scatter(torch.zeros_like(RRt), self.zero_indices, R_clamp)
        return R_clamp, RRt_clamp
//  [python]
def _condition_RRt(self, RRt, D):
        """
        Args:
            RRt (torch.tensor): of size (N x N) the original full covariance
            D (torch.tensor): of dtype float and size (1xN) containing 1.0 for known indices else 0.0.
        """
        self.register_buffer("D_atom", self._C_atomic(D))
        self.register_buffer("zero_indices", torch.nonzero((1 - self.D_atom[0]))[:, 0])
        self.register_buffer("nonzero_indices", torch.nonzero(self.D_atom[0])[:, 0])

        self.register_buffer("S11", RRt[self.zero_indices][:, self.zero_indices])
        self.register_buffer("S12", RRt[self.zero_indices][:, self.nonzero_indices])
        self.register_buffer("S21", RRt[self.nonzero_indices][:, self.zero_indices])
        self.register_buffer("S22", RRt[self.nonzero_indices][:, self.nonzero_indices])

        S_clamp = self.S11 - ((self.S12 @ torch.linalg.pinv(self.S22) @ self.S21))
        R_clamp = torch.linalg.cholesky(S_clamp)
        self.register_buffer("RRt_clamp_restricted", R_clamp @ R_clamp.t())
        RRt_clamp = self._scatter(
            torch.zeros_like(RRt), self.zero_indices, self.RRt_clamp_restricted
        )
        R_clamp = self._scatter(torch.zeros_like(RRt), self.zero_indices, R_clamp)
        return R_clamp, RRt_clamp
//  [python]
def calculate_factorizations(x:torch.Tensor, y:torch.Tensor, models: "list[ExactGPModelMonoTask]"):
        """
        Compute iK and beta using the points in memory, which are needed to make predictions with the gaussian processes.
        These two variables only depends on data in memory, and not on input distribution,
        so they separated from other computation such that they can be computed outside the optimisation function,
        which is computed multiple times at each iteration

        Function inspired from
        https://github.com/nrontsis/PILCO/blob/6a962c8e4172f9e7f29ed6e373c4be2dd4b69cb7/pilco/models/mgpr.py#L81,
        reimplemented from tensorflow to pytorch
        Args:
            x (torch.Tensor): matrix containing the states and actions. Dim=(Nm, Ns + Na + Nc)
            y (torch.Tensor): matrix containing the states change. Dim=(Nm, Ns)
            models (list of gpytorch.models.ExactGP): list containing the gp models used to predict each state change.
                                                        Len=Ns
            Ns: number of states
            Na: number of actions
            Nc: number of context variables. At the moment, only 1 if time used
            Nm: number of points in memory

        Returns:
            iK (torch.Tensor): needed by the gaussian processes models to compute the predictions
            beta (torch.Tensor): needed by the gaussian processes models to compute the predictions

        """
        K = torch.stack([model.covar_module(x).evaluate() for model in models])
        batched_eye = torch.eye(K.shape[1]).repeat(K.shape[0], 1, 1)
        L = torch.cholesky(K + torch.stack([model.likelihood.noise for model in models])[:, None] * batched_eye)
        iK = torch.cholesky_solve(batched_eye, L)
        Y_ = (y.t())[:, :, None]
        beta = torch.cholesky_solve(Y_, L)[:, :, 0]
        return iK, beta
//  [python]
def calculate_factorizations(x:torch.Tensor, y:torch.Tensor, models: "list[ExactGPModelMonoTask]"):
        """
        Compute iK and beta using the points in memory, which are needed to make predictions with the gaussian processes.
        These two variables only depends on data in memory, and not on input distribution,
        so they separated from other computation such that they can be computed outside the optimisation function,
        which is computed multiple times at each iteration

        Function inspired from
        https://github.com/nrontsis/PILCO/blob/6a962c8e4172f9e7f29ed6e373c4be2dd4b69cb7/pilco/models/mgpr.py#L81,
        reimplemented from tensorflow to pytorch
        Args:
            x (torch.Tensor): matrix containing the states and actions. Dim=(Nm, Ns + Na + Nc)
            y (torch.Tensor): matrix containing the states change. Dim=(Nm, Ns)
            models (list of gpytorch.models.ExactGP): list containing the gp models used to predict each state change.
                                                        Len=Ns
            Ns: number of states
            Na: number of actions
            Nc: number of context variables. At the moment, only 1 if time used
            Nm: number of points in memory

        Returns:
            iK (torch.Tensor): needed by the gaussian processes models to compute the predictions
            beta (torch.Tensor): needed by the gaussian processes models to compute the predictions

        """
        K = torch.stack([model.covar_module(x).evaluate() for model in models])
        batched_eye = torch.eye(K.shape[1]).repeat(K.shape[0], 1, 1)
        L = torch.linalg.cholesky(K + torch.stack([model.likelihood.noise for model in models])[:, None] * batched_eye)
        iK = torch.cholesky_solve(batched_eye, L)
        Y_ = (y.t())[:, :, None]
        beta = torch.cholesky_solve(Y_, L)[:, :, 0]
        return iK, beta
//  [python]
def batch_var(coords_train_norm, E_train, coords_test_norm, E_test, model, model_level):
  b, m, _ = coords_train_norm.shape
  K_mm = model.cov_modules[model_level](coords_train_norm, E_train)
  noise = model.get_var(model_level) * torch.ones(b,m,device=coords_train_norm.device)
  K_mn = model.cross_cov_modules[model_level](coords_train_norm, E_train, coords_test_norm, E_test)
  K_nn_diag = model.diagonal_cov_modules[model_level](coords_test_norm, E_test)
  L = torch.cholesky(K_mm, upper=False)

  v = torch.triangular_solve(K_mn, L, upper=False)[0]
  pred_var = K_nn_diag - torch.sum(v * v, dim=1)

  return pred_var
//  [python]
def batch_var(coords_train_norm, E_train, coords_test_norm, E_test, model, model_level):
  b, m, _ = coords_train_norm.shape
  K_mm = model.cov_modules[model_level](coords_train_norm, E_train)
  noise = model.get_var(model_level) * torch.ones(b,m,device=coords_train_norm.device)
  K_mn = model.cross_cov_modules[model_level](coords_train_norm, E_train, coords_test_norm, E_test)
  K_nn_diag = model.diagonal_cov_modules[model_level](coords_test_norm, E_test)
  L = torch.linalg.cholesky(K_mm, upper=False)

  v = torch.linalg.solve_triangular(L, K_mn, upper=False)
  pred_var = K_nn_diag - torch.sum(v * v, dim=1)

  return pred_var
//  [python]
def model(self):
        self.set_mode("model")

        # W = (inv(Luu) @ Kuf).T
        # Qff = Kfu @ inv(Kuu) @ Kuf = W @ W.T
        # Fomulas for each approximation method are
        # DTC:  y_cov = Qff + noise,                   trace_term = 0
        # FITC: y_cov = Qff + diag(Kff - Qff) + noise, trace_term = 0
        # VFE:  y_cov = Qff + noise,                   trace_term = tr(Kff-Qff) / noise
        # y_cov = W @ W.T + D
        # trace_term is added into log_prob

        N = self.X.size(0)
        M = self.Xu.size(0)
        Kuu = self.kernel(self.Xu).contiguous()
        Kuu.view(-1)[:: M + 1] += self.jitter  # add jitter to the diagonal
        Luu = torch.cholesky(Kuu)
        Kuf = self.kernel(self.Xu, self.X)
        W = torch.triangular_solve(Kuf, Luu, upper=False)[0].t()

        D = self.noise.expand(N)
        if self.approx == "FITC" or self.approx == "VFE":
            Kffdiag = self.kernel(self.X, diag=True)
            Qffdiag = W.pow(2).sum(dim=-1)
            if self.approx == "FITC":
                D = D + Kffdiag - Qffdiag
            else:  # approx = "VFE"
                trace_term = (Kffdiag - Qffdiag).sum() / self.noise
                trace_term = trace_term.clamp(min=0)

        zero_loc = self.X.new_zeros(N)
        f_loc = zero_loc + self.mean_function(self.X)
        if self.y is None:
            f_var = D + W.pow(2).sum(dim=-1)
            return f_loc, f_var
        else:
            if self.approx == "VFE":
                pyro.factor(self._pyro_get_fullname("trace_term"), -trace_term / 2.0)

            return pyro.sample(
                self._pyro_get_fullname("y"),
                dist.LowRankMultivariateNormal(f_loc, W, D)
                .expand_by(self.y.shape[:-1])
                .to_event(self.y.dim() - 1),
                obs=self.y,
            )