jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_lkj_prior_batch_log_prob(self, cuda=False):
        device = torch.device("cuda") if cuda else torch.device("cpu")
        prior = LKJPrior(2, torch.tensor([0.5, 1.5], device=device))
        dist = LKJCholesky(2, torch.tensor([0.5, 1.5], device=device))

        S = torch.eye(2, device=device)
        S_chol = torch.linalg.cholesky(S)
        self.assertTrue(approx_equal(prior.log_prob(S), dist.log_prob(S_chol)))
        S = torch.stack([S, torch.tensor([[1.0, 0.5], [0.5, 1]], device=S.device)])
        S_chol = torch.linalg.cholesky(S)
        self.assertTrue(approx_equal(prior.log_prob(S), dist.log_prob(S_chol)))
        with self.assertRaises(ValueError):
            prior.log_prob(torch.eye(3, device=device))
//  [python]
def test_wishart_log_prob(self):
        set_rng_seed(0)  # see Note [Randomized statistical tests]
        ndim = 3
        df = torch.rand([], requires_grad=True) + ndim - 1
        # SciPy allowed ndim -1 < df < ndim for Wishar distribution after version 1.7.0
        if version.parse(scipy.__version__) < version.parse("1.7.0"):
            df += 1.0
        tmp = torch.randn(ndim, 10)
        cov = (torch.matmul(tmp, tmp.t()) / tmp.size(-1)).requires_grad_()
        prec = cov.inverse().requires_grad_()
        scale_tril = torch.cholesky(cov).requires_grad_()

        # check that logprob values match scipy logpdf,
        # and that covariance and scale_tril parameters are equivalent
        dist1 = Wishart(df, cov)
        dist2 = Wishart(df, precision_matrix=prec)
        dist3 = Wishart(df, scale_tril=scale_tril)
        ref_dist = scipy.stats.wishart(df.item(), cov.detach().numpy())

        x = dist1.sample((1000,))
        expected = ref_dist.logpdf(x.transpose(0, 2).numpy())

        self.assertEqual(
            0.0,
            np.mean((dist1.log_prob(x).detach().numpy() - expected) ** 2),
            atol=1e-3,
            rtol=0,
        )
        self.assertEqual(
            0.0,
            np.mean((dist2.log_prob(x).detach().numpy() - expected) ** 2),
            atol=1e-3,
            rtol=0,
        )
        self.assertEqual(
            0.0,
            np.mean((dist3.log_prob(x).detach().numpy() - expected) ** 2),
            atol=1e-3,
            rtol=0,
        )

        # Double-check that batched versions behave the same as unbatched
        df = torch.rand(5, requires_grad=True) + ndim - 1
        # SciPy allowed ndim -1 < df < ndim for Wishar distribution after version 1.7.0
        if version.parse(scipy.__version__) < version.parse("1.7.0"):
            df += 1.0
        tmp = torch.randn(5, ndim, 10)
        cov = (tmp.unsqueeze(-2) * tmp.unsqueeze(-3)).mean(-1).requires_grad_()

        dist_batched = Wishart(df, cov)
        dist_unbatched = [Wishart(df[i], cov[i]) for i in range(df.size(0))]

        x = dist_batched.sample((1000,))
        batched_prob = dist_batched.log_prob(x)
        unbatched_prob = torch.stack(
            [dist_unbatched[i].log_prob(x[:, i]) for i in range(5)]
        ).t()

        self.assertEqual(batched_prob.shape, unbatched_prob.shape)
        self.assertEqual(batched_prob, unbatched_prob, atol=1e-3, rtol=0)
//  [python]
def test_wishart_log_prob(self):
        set_rng_seed(0)  # see Note [Randomized statistical tests]
        ndim = 3
        df = torch.rand([], requires_grad=True) + ndim - 1
        # SciPy allowed ndim -1 < df < ndim for Wishar distribution after version 1.7.0
        if version.parse(scipy.__version__) < version.parse("1.7.0"):
            df += 1.0
        tmp = torch.randn(ndim, 10)
        cov = (torch.matmul(tmp, tmp.t()) / tmp.size(-1)).requires_grad_()
        prec = cov.inverse().requires_grad_()
        scale_tril = torch.linalg.cholesky(cov).requires_grad_()

        # check that logprob values match scipy logpdf,
        # and that covariance and scale_tril parameters are equivalent
        dist1 = Wishart(df, cov)
        dist2 = Wishart(df, precision_matrix=prec)
        dist3 = Wishart(df, scale_tril=scale_tril)
        ref_dist = scipy.stats.wishart(df.item(), cov.detach().numpy())

        x = dist1.sample((1000,))
        expected = ref_dist.logpdf(x.transpose(0, 2).numpy())

        self.assertEqual(
            0.0,
            np.mean((dist1.log_prob(x).detach().numpy() - expected) ** 2),
            atol=1e-3,
            rtol=0,
        )
        self.assertEqual(
            0.0,
            np.mean((dist2.log_prob(x).detach().numpy() - expected) ** 2),
            atol=1e-3,
            rtol=0,
        )
        self.assertEqual(
            0.0,
            np.mean((dist3.log_prob(x).detach().numpy() - expected) ** 2),
            atol=1e-3,
            rtol=0,
        )

        # Double-check that batched versions behave the same as unbatched
        df = torch.rand(5, requires_grad=True) + ndim - 1
        # SciPy allowed ndim -1 < df < ndim for Wishar distribution after version 1.7.0
        if version.parse(scipy.__version__) < version.parse("1.7.0"):
            df += 1.0
        tmp = torch.randn(5, ndim, 10)
        cov = (tmp.unsqueeze(-2) * tmp.unsqueeze(-3)).mean(-1).requires_grad_()

        dist_batched = Wishart(df, cov)
        dist_unbatched = [Wishart(df[i], cov[i]) for i in range(df.size(0))]

        x = dist_batched.sample((1000,))
        batched_prob = dist_batched.log_prob(x)
        unbatched_prob = torch.stack(
            [dist_unbatched[i].log_prob(x[:, i]) for i in range(5)]
        ).t()

        self.assertEqual(batched_prob.shape, unbatched_prob.shape)
        self.assertEqual(batched_prob, unbatched_prob, atol=1e-3, rtol=0)
//  [python]
def test_lkj_cholesky_factor_prior_log_prob(self, cuda=False):
        device = torch.device("cuda") if cuda else torch.device("cpu")
        prior = LKJCholeskyFactorPrior(2, torch.tensor(0.5, device=device))
        dist = LKJCholesky(2, torch.tensor(0.5, device=device))
        S = torch.eye(2, device=device)
        S_chol = torch.cholesky(S)
        self.assertAlmostEqual(prior.log_prob(S_chol), dist.log_prob(S_chol), places=4)
        S = torch.stack([S, torch.tensor([[1.0, 0.5], [0.5, 1]], device=S_chol.device)])
        S_chol = torch.stack([torch.cholesky(Si) for Si in S])
        self.assertTrue(approx_equal(prior.log_prob(S_chol), dist.log_prob(S_chol)))
//  [python]
def test_lkj_cholesky_factor_prior_log_prob(self, cuda=False):
        device = torch.device("cuda") if cuda else torch.device("cpu")
        prior = LKJCholeskyFactorPrior(2, torch.tensor(0.5, device=device))
        dist = LKJCholesky(2, torch.tensor(0.5, device=device))
        S = torch.eye(2, device=device)
        S_chol = torch.linalg.cholesky(S)
        self.assertAlmostEqual(prior.log_prob(S_chol), dist.log_prob(S_chol), places=4)
        S = torch.stack([S, torch.tensor([[1.0, 0.5], [0.5, 1]], device=S_chol.device)])
        S_chol = torch.stack([torch.linalg.cholesky(Si) for Si in S])
        self.assertTrue(approx_equal(prior.log_prob(S_chol), dist.log_prob(S_chol)))
//  [python]
def test_lkj_cholesky_factor_prior_batch_log_prob(self, cuda=False):
        device = torch.device("cuda") if cuda else torch.device("cpu")
        prior = LKJCholeskyFactorPrior(2, torch.tensor([0.5, 1.5], device=device))
        dist = LKJCholesky(2, torch.tensor([0.5, 1.5], device=device))

        S = torch.eye(2, device=device)
        S_chol = torch.cholesky(S)
        self.assertTrue(approx_equal(prior.log_prob(S_chol), dist.log_prob(S_chol)))
        S = torch.stack([S, torch.tensor([[1.0, 0.5], [0.5, 1]], device=S.device)])
        S_chol = torch.stack([torch.cholesky(Si) for Si in S])
        self.assertTrue(approx_equal(prior.log_prob(S_chol), dist.log_prob(S_chol)))
//  [python]
def test_lkj_cholesky_factor_prior_batch_log_prob(self, cuda=False):
        device = torch.device("cuda") if cuda else torch.device("cpu")
        prior = LKJCholeskyFactorPrior(2, torch.tensor([0.5, 1.5], device=device))
        dist = LKJCholesky(2, torch.tensor([0.5, 1.5], device=device))

        S = torch.eye(2, device=device)
        S_chol = torch.linalg.cholesky(S)
        self.assertTrue(approx_equal(prior.log_prob(S_chol), dist.log_prob(S_chol)))
        S = torch.stack([S, torch.tensor([[1.0, 0.5], [0.5, 1]], device=S.device)])
        S_chol = torch.stack([torch.linalg.cholesky(Si) for Si in S])
        self.assertTrue(approx_equal(prior.log_prob(S_chol), dist.log_prob(S_chol)))
//  [python]
def test_lkj_covariance_prior_log_prob(self, cuda=False):
        device = torch.device("cuda") if cuda else torch.device("cpu")
        sd_prior = SmoothedBoxPrior(exp(-1), exp(1))
        if cuda:
            sd_prior = sd_prior.cuda()
        prior = LKJCovariancePrior(2, torch.tensor(0.5, device=device), sd_prior)
        S = torch.eye(2, device=device)

        corr_dist = LKJCholesky(2, torch.tensor(0.5, device=device))
        dist_log_prob = corr_dist.log_prob(S) + sd_prior.log_prob(S.diagonal(dim1=-1, dim2=-2)).sum()
        self.assertAlmostEqual(prior.log_prob(S), dist_log_prob, places=4)

        S = torch.stack([S, torch.tensor([[1.0, 0.5], [0.5, 1]], device=S.device)])
        S_chol = torch.cholesky(S)
        dist_log_prob = corr_dist.log_prob(S_chol) + sd_prior.log_prob(torch.diagonal(S, dim1=-2, dim2=-1))
        self.assertTrue(approx_equal(prior.log_prob(S), dist_log_prob))
//  [python]
def test_lkj_covariance_prior_log_prob(self, cuda=False):
        device = torch.device("cuda") if cuda else torch.device("cpu")
        sd_prior = SmoothedBoxPrior(exp(-1), exp(1))
        if cuda:
            sd_prior = sd_prior.cuda()
        prior = LKJCovariancePrior(2, torch.tensor(0.5, device=device), sd_prior)
        S = torch.eye(2, device=device)

        corr_dist = LKJCholesky(2, torch.tensor(0.5, device=device))
        dist_log_prob = corr_dist.log_prob(S) + sd_prior.log_prob(S.diagonal(dim1=-1, dim2=-2)).sum()
        self.assertAlmostEqual(prior.log_prob(S), dist_log_prob, places=4)

        S = torch.stack([S, torch.tensor([[1.0, 0.5], [0.5, 1]], device=S.device)])
        S_chol = torch.linalg.cholesky(S)
        dist_log_prob = corr_dist.log_prob(S_chol) + sd_prior.log_prob(torch.diagonal(S, dim1=-2, dim2=-1))
        self.assertTrue(approx_equal(prior.log_prob(S), dist_log_prob))
//  [python]
def optnet_factor_kkt(U_S, R, d, n_eq, n_ineq, int_reg=10 ** -6):
    dtype = U_S.dtype
    n_batch = U_S.shape[0]
    zeros = torch.zeros(n_batch, n_ineq, n_eq, dtype=dtype)

    # --- update U_S for d:
    d_diag = torch.diag_embed(1 / d.squeeze(2))

    reg = torch.eye(n_ineq, dtype=dtype).unsqueeze(0)
    mat = torch.cholesky(R + d_diag + int_reg * reg, upper=True)
    mat = torch.cat((zeros, mat), dim=2)

    U1 = U_S[:, [0], :]
    out = torch.cat((U1, mat), dim=1)

    return out
//  [python]
def optnet_factor_kkt(U_S, R, d, n_eq, n_ineq, int_reg=10 ** -6):
    dtype = U_S.dtype
    n_batch = U_S.shape[0]
    zeros = torch.zeros(n_batch, n_ineq, n_eq, dtype=dtype)

    # --- update U_S for d:
    d_diag = torch.diag_embed(1 / d.squeeze(2))

    reg = torch.eye(n_ineq, dtype=dtype).unsqueeze(0)
    mat = torch.linalg.cholesky(R + d_diag + int_reg * reg, upper=True)
    mat = torch.cat((zeros, mat), dim=2)

    U1 = U_S[:, [0], :]
    out = torch.cat((U1, mat), dim=1)

    return out
//  [python]
def test_lkj_covariance_prior_log_prob_hetsd(self, cuda=False):
        device = torch.device("cuda") if cuda else torch.device("cpu")
        a = torch.tensor([exp(-1), exp(-2)], device=device)
        b = torch.tensor([exp(1), exp(2)], device=device)
        sd_prior = SmoothedBoxPrior(a, b)
        prior = LKJCovariancePrior(2, torch.tensor(0.5, device=device), sd_prior)
        corr_dist = LKJCholesky(2, torch.tensor(0.5, device=device))

        S = torch.eye(2, device=device)
        dist_log_prob = corr_dist.log_prob(S) + sd_prior.log_prob(S.diagonal(dim1=-1, dim2=-2)).sum()
        self.assertAlmostEqual(prior.log_prob(S), dist_log_prob, places=4)

        S = torch.stack([S, torch.tensor([[1.0, 0.5], [0.5, 1]], device=S.device)])
        S_chol = torch.cholesky(S)
        dist_log_prob = corr_dist.log_prob(S_chol) + sd_prior.log_prob(torch.diagonal(S, dim1=-2, dim2=-1))
        self.assertTrue(approx_equal(prior.log_prob(S), dist_log_prob))
//  [python]
def test_lkj_covariance_prior_log_prob_hetsd(self, cuda=False):
        device = torch.device("cuda") if cuda else torch.device("cpu")
        a = torch.tensor([exp(-1), exp(-2)], device=device)
        b = torch.tensor([exp(1), exp(2)], device=device)
        sd_prior = SmoothedBoxPrior(a, b)
        prior = LKJCovariancePrior(2, torch.tensor(0.5, device=device), sd_prior)
        corr_dist = LKJCholesky(2, torch.tensor(0.5, device=device))

        S = torch.eye(2, device=device)
        dist_log_prob = corr_dist.log_prob(S) + sd_prior.log_prob(S.diagonal(dim1=-1, dim2=-2)).sum()
        self.assertAlmostEqual(prior.log_prob(S), dist_log_prob, places=4)

        S = torch.stack([S, torch.tensor([[1.0, 0.5], [0.5, 1]], device=S.device)])
        S_chol = torch.linalg.cholesky(S)
        dist_log_prob = corr_dist.log_prob(S_chol) + sd_prior.log_prob(torch.diagonal(S, dim1=-2, dim2=-1))
        self.assertTrue(approx_equal(prior.log_prob(S), dist_log_prob))
//  [python]
def test_lkj_covariance_prior_batch_log_prob(self, cuda=False):
        device = torch.device("cuda") if cuda else torch.device("cpu")
        v = torch.ones(2, 1, device=device)
        sd_prior = SmoothedBoxPrior(exp(-1) * v, exp(1) * v)
        prior = LKJCovariancePrior(2, torch.tensor([0.5, 1.5], device=device), sd_prior)
        corr_dist = LKJCholesky(2, torch.tensor([0.5, 1.5], device=device))

        S = torch.eye(2, device=device)
        dist_log_prob = corr_dist.log_prob(S) + sd_prior.log_prob(S.diagonal(dim1=-1, dim2=-2))
        self.assertLessEqual((prior.log_prob(S) - dist_log_prob).abs().sum(), 1e-4)

        S = torch.stack([S, torch.tensor([[1.0, 0.5], [0.5, 1]], device=S.device)])
        S_chol = torch.cholesky(S)
        dist_log_prob = corr_dist.log_prob(S_chol) + sd_prior.log_prob(torch.diagonal(S, dim1=-2, dim2=-1))
        self.assertLessEqual((prior.log_prob(S) - dist_log_prob).abs().sum(), 1e-4)
//  [python]
def test_lkj_covariance_prior_batch_log_prob(self, cuda=False):
        device = torch.device("cuda") if cuda else torch.device("cpu")
        v = torch.ones(2, 1, device=device)
        sd_prior = SmoothedBoxPrior(exp(-1) * v, exp(1) * v)
        prior = LKJCovariancePrior(2, torch.tensor([0.5, 1.5], device=device), sd_prior)
        corr_dist = LKJCholesky(2, torch.tensor([0.5, 1.5], device=device))

        S = torch.eye(2, device=device)
        dist_log_prob = corr_dist.log_prob(S) + sd_prior.log_prob(S.diagonal(dim1=-1, dim2=-2))
        self.assertLessEqual((prior.log_prob(S) - dist_log_prob).abs().sum(), 1e-4)

        S = torch.stack([S, torch.tensor([[1.0, 0.5], [0.5, 1]], device=S.device)])
        S_chol = torch.linalg.cholesky(S)
        dist_log_prob = corr_dist.log_prob(S_chol) + sd_prior.log_prob(torch.diagonal(S, dim1=-2, dim2=-1))
        self.assertLessEqual((prior.log_prob(S) - dist_log_prob).abs().sum(), 1e-4)
//  [python]
def __init__(self, t, cov1, cov2, cov_target):
        # Compute necessities
        if t == 0.5:
            self.chol_forward = torch.cholesky(cov_target)
        elif t < 0.5:
            self.chol_a = torch.cholesky(cov1)
            self.chol_b = torch.cholesky(cov_target)
            self.chol_forward = (1.0 - 2.0 * t) * self.chol_a + 2.0 * t * self.chol_b
        else:
            self.chol_a = torch.cholesky(cov2)
            self.chol_b = torch.cholesky(cov_target)
            self.chol_forward = (2.0 * t - 1.0) * self.chol_a + 2.0 * (1.0 - t) * self.chol_b
        self.log_jac = torch.slogdet(self.chol_forward).logabsdet
        self.chol_inverse = torch.inverse(self.chol_forward)
//  [python]
def __init__(self, t, cov1, cov2, cov_target):
        # Compute necessities
        if t == 0.5:
            self.chol_forward = torch.linalg.cholesky(cov_target)
        elif t < 0.5:
            self.chol_a = torch.linalg.cholesky(cov1)
            self.chol_b = torch.linalg.cholesky(cov_target)
            self.chol_forward = (1.0 - 2.0 * t) * self.chol_a + 2.0 * t * self.chol_b
        else:
            self.chol_a = torch.linalg.cholesky(cov2)
            self.chol_b = torch.linalg.cholesky(cov_target)
            self.chol_forward = (2.0 * t - 1.0) * self.chol_a + 2.0 * (1.0 - t) * self.chol_b
        self.log_jac = torch.linalg.slogdet(self.chol_forward).logabsdet
        self.chol_inverse = torch.linalg.inv(self.chol_forward)
//  [python]
def to_group(self, x):
        return torch.cholesky(x, upper=True)
//  [python]
def to_group(self, x):
        return torch.linalg.cholesky(x, upper=True)
//  [python]
def optnet_pre_factor_kkt(Q, G, A):
    # S =  [ A Q^{-1} A^T        A Q^{-1} G^T           ]
    #      [ G Q^{-1} A^T        G Q^{-1} G^T + D^{-1} ]
    # S = rbind(cbind(A_invQ_AT,A_invQ_GT),cbind(G_invQ_AT,G%*%invQ_GT +diag(n_ineq)))
    # chol(S)[1,] == U_S[1,]

    # --- prep:
    n_batch = Q.shape[0]
    dtype = Q.dtype
    n_eq = get_ncon(A, dim=1)
    any_eq = n_eq > 0
    n_ineq = get_ncon(G, dim=1)
    n_con = n_eq + n_ineq
    GT = torch.transpose(G, 1, 2)
    AT = torch.transpose(A, 1, 2)

    U_Q = torch.cholesky(Q, upper=True)
    U_S = torch.zeros((n_batch, n_con, n_con), dtype=dtype)

    # --- ineq:
    invQ_GT = torch.cholesky_solve(GT, U_Q, upper=True)
    R = torch.matmul(G, invQ_GT)

    if any_eq:
        # --- eq
        invQ_AT = torch.cholesky_solve(AT, U_Q, upper=True)
        A_invQ_AT = torch.matmul(A, invQ_AT)
        U11 = torch.cholesky(A_invQ_AT, upper=True)

        # --- cross product terms
        G_invQ_AT = torch.matmul(G, invQ_AT)
        U12 = torch.solve(U11, torch.transpose(G_invQ_AT, 1, 2))[0]

        U1 = torch.cat((U11, U12), dim=2)
        zeros = torch.zeros((n_batch, n_ineq, n_con), dtype=dtype)
        U_S = torch.cat((U1, zeros), dim=1)

        R = R - torch.matmul(torch.transpose(U12, 1, 2), U12)

    return U_Q, U_S, R
//  [python]
def optnet_pre_factor_kkt(Q, G, A):
    # S =  [ A Q^{-1} A^T        A Q^{-1} G^T           ]
    #      [ G Q^{-1} A^T        G Q^{-1} G^T + D^{-1} ]
    # S = rbind(cbind(A_invQ_AT,A_invQ_GT),cbind(G_invQ_AT,G%*%invQ_GT +diag(n_ineq)))
    # chol(S)[1,] == U_S[1,]

    # --- prep:
    n_batch = Q.shape[0]
    dtype = Q.dtype
    n_eq = get_ncon(A, dim=1)
    any_eq = n_eq > 0
    n_ineq = get_ncon(G, dim=1)
    n_con = n_eq + n_ineq
    GT = torch.transpose(G, 1, 2)
    AT = torch.transpose(A, 1, 2)

    U_Q = torch.linalg.cholesky(Q, upper=True)
    U_S = torch.zeros((n_batch, n_con, n_con), dtype=dtype)

    # --- ineq:
    invQ_GT = torch.cholesky_solve(GT, U_Q, upper=True)
    R = torch.matmul(G, invQ_GT)

    if any_eq:
        # --- eq
        invQ_AT = torch.cholesky_solve(AT, U_Q, upper=True)
        A_invQ_AT = torch.matmul(A, invQ_AT)
        U11 = torch.linalg.cholesky(A_invQ_AT, upper=True)

        # --- cross product terms
        G_invQ_AT = torch.matmul(G, invQ_AT)
        U12 = torch.linalg.solve(U11, torch.transpose(G_invQ_AT, 1, 2))

        U1 = torch.cat((U11, U12), dim=2)
        zeros = torch.zeros((n_batch, n_ineq, n_con), dtype=dtype)
        U_S = torch.cat((U1, zeros), dim=1)

        R = R - torch.matmul(torch.transpose(U12, 1, 2), U12)

    return U_Q, U_S, R
//  [python]
def forward(ctx, H, b):
        """
        :param ctx:
        :param H: [b, f, 6, 6]
        :param b: [b, f, 6]
        :return:
        """
        chol = torch.cholesky(H)
        xx = torch.cholesky_solve(b, chol)
        ctx.save_for_backward(chol, xx)
        return xx
//  [python]
def forward(ctx, H, b):
        """
        :param ctx:
        :param H: [b, f, 6, 6]
        :param b: [b, f, 6]
        :return:
        """
        chol = torch.linalg.cholesky(H)
        xx = torch.cholesky_solve(b, chol)
        ctx.save_for_backward(chol, xx)
        return xx
//  [python]
def solve_pose(H, b, max_update=1.0):
    u = torch.cholesky(H)
    pose_vec = torch.cholesky_solve(b, u)
    pose_vec = torch.where(torch.isnan(pose_vec), torch.zeros_like(pose_vec), pose_vec)
    pose_vec = pose_vec[..., 0]
    return pose_vec
//  [python]
def solve_pose(H, b, max_update=1.0):
    u = torch.linalg.cholesky(H)
    pose_vec = torch.cholesky_solve(b, u)
    pose_vec = torch.where(torch.isnan(pose_vec), torch.zeros_like(pose_vec), pose_vec)
    pose_vec = pose_vec[..., 0]
    return pose_vec
//  [python]
def _stable_cholesky(input, upper=False, out=None, eps=1e-06):
    try:
        if upper:
            return torch.cholesky(input, upper=True, out=out).mH
        return torch.cholesky(input, upper=False, out=out)
    except RuntimeError:
        input = condition_scm(input, eps)
        if upper:
            return torch.cholesky(input, upper=True, out=out).mH
        return torch.cholesky(input, upper=False, out=out)
//  [python]
def _stable_cholesky(input, upper=False, out=None, eps=1e-06):
    try:
        if upper:
            return torch.linalg.cholesky(input, out=out).mH
        return torch.linalg.cholesky(input, out=out)
    except RuntimeError:
        input = condition_scm(input, eps)
        if upper:
            return torch.linalg.cholesky(input, out=out).mH
        return torch.linalg.cholesky(input, out=out)
//  [python]
def draw_from_gp(inputs, sf, ell, L=1, N=1, n_out=1, eps=1e-5):
    if inputs.ndim == 1:
        inputs = inputs.unsqueeze(1) 
    T = inputs.shape[0]
    cov  = K(inputs,inputs,ell,sf,eps=eps) # T,T
    L_ = torch.cholesky(cov) # psd_safe_cholesky(cov)
    # L,N,T,n_out or N,T,n_out or T,n_out
    return L_ @ torch.randn([L,N,T,n_out],device=inputs.device).squeeze(0).squeeze(0)
//  [python]
def draw_from_gp(inputs, sf, ell, L=1, N=1, n_out=1, eps=1e-5):
    if inputs.ndim == 1:
        inputs = inputs.unsqueeze(1) 
    T = inputs.shape[0]
    cov  = K(inputs,inputs,ell,sf,eps=eps) # T,T
    L_ = torch.linalg.cholesky(cov) # psd_safe_cholesky(cov)
    # L,N,T,n_out or N,T,n_out or T,n_out
    return L_ @ torch.randn([L,N,T,n_out],device=inputs.device).squeeze(0).squeeze(0)
//  [python]
def test_pyro_catch_error(self) -> None:
        def potential_fn(z: Dict[str, torch.Tensor]) -> torch.Tensor:
            # pyre-fixme[16]: Module `distributions` has no attribute
            #  `MultivariateNormal`.
            mvn = pyro.distributions.MultivariateNormal(
                loc=torch.zeros(2),
                covariance_matrix=z["K"],
            )
            return mvn.log_prob(torch.zeros(2))

        # Test base case where everything is fine
        z = {"K": torch.eye(2)}
        grads, val = potential_grad(potential_fn, z)
        self.assertTrue(torch.allclose(grads["K"], -0.5 * torch.eye(2)))
        norm_mvn = torch.distributions.Normal(0, 1)
        self.assertTrue(torch.allclose(val, 2 * norm_mvn.log_prob(torch.zeros(1))))

        # Default behavior should catch the ValueError when trying to instantiate
        # the MVN and return NaN instead
        z = {"K": torch.ones(2, 2)}
        _, val = potential_grad(potential_fn, z)
        self.assertTrue(torch.isnan(val))

        # Default behavior should catch the LinAlgError when peforming a
        # Cholesky decomposition and return NaN instead
        def potential_fn_chol(z: Dict[str, torch.Tensor]) -> torch.Tensor:
            return torch.cholesky(z["K"])

        _, val = potential_grad(potential_fn_chol, z)
        self.assertTrue(torch.isnan(val))

        # Default behavior should not catch other errors
        def potential_fn_rterr_foo(z: Dict[str, torch.Tensor]) -> torch.Tensor:
            raise RuntimeError("foo")

        with self.assertRaisesRegex(RuntimeError, "foo"):
            potential_grad(potential_fn_rterr_foo, z)
//  [python]
def test_pyro_catch_error(self) -> None:
        def potential_fn(z: Dict[str, torch.Tensor]) -> torch.Tensor:
            # pyre-fixme[16]: Module `distributions` has no attribute
            #  `MultivariateNormal`.
            mvn = pyro.distributions.MultivariateNormal(
                loc=torch.zeros(2),
                covariance_matrix=z["K"],
            )
            return mvn.log_prob(torch.zeros(2))

        # Test base case where everything is fine
        z = {"K": torch.eye(2)}
        grads, val = potential_grad(potential_fn, z)
        self.assertTrue(torch.allclose(grads["K"], -0.5 * torch.eye(2)))
        norm_mvn = torch.distributions.Normal(0, 1)
        self.assertTrue(torch.allclose(val, 2 * norm_mvn.log_prob(torch.zeros(1))))

        # Default behavior should catch the ValueError when trying to instantiate
        # the MVN and return NaN instead
        z = {"K": torch.ones(2, 2)}
        _, val = potential_grad(potential_fn, z)
        self.assertTrue(torch.isnan(val))

        # Default behavior should catch the LinAlgError when peforming a
        # Cholesky decomposition and return NaN instead
        def potential_fn_chol(z: Dict[str, torch.Tensor]) -> torch.Tensor:
            return torch.linalg.cholesky(z["K"])

        _, val = potential_grad(potential_fn_chol, z)
        self.assertTrue(torch.isnan(val))

        # Default behavior should not catch other errors
        def potential_fn_rterr_foo(z: Dict[str, torch.Tensor]) -> torch.Tensor:
            raise RuntimeError("foo")

        with self.assertRaisesRegex(RuntimeError, "foo"):
            potential_grad(potential_fn_rterr_foo, z)
//  [python]
def precompute(self, batch_size: int = 1):
        """
        Precompute inducing points
        Args:
            batch_size (int, optional):
        """

        Kz = self.increment_weight(self.Z)  # num_step, num_inducing, dim
        self.Kz = Kz.transpose(2, 1).unsqueeze(-1)  # num_step, dim, num_inducing, 1

        # kernel function is computed as in Eq. 8
        Kzz = torch.matmul(
            self.Kz, self.Kz.transpose(-1, -2)
        )  # num_step, dim, num_inducing, num_ducing
        Kzz = Kzz.sum(dim=0)  # dim, num_inducing , num_inducing
        self.Lz = torch.cholesky(
            Kzz + torch.eye(self.num_inducing).to(self.Z) * 1e-6
        )
        eps = torch.randn(self.dim, self.num_inducing, batch_size).to(self.Z)
        self.delta_Bz = torch.matmul(self.Lz, eps)
//  [python]
def precompute(self, batch_size: int = 1):
        """
        Precompute inducing points
        Args:
            batch_size (int, optional):
        """

        Kz = self.increment_weight(self.Z)  # num_step, num_inducing, dim
        self.Kz = Kz.transpose(2, 1).unsqueeze(-1)  # num_step, dim, num_inducing, 1

        # kernel function is computed as in Eq. 8
        Kzz = torch.matmul(
            self.Kz, self.Kz.transpose(-1, -2)
        )  # num_step, dim, num_inducing, num_ducing
        Kzz = Kzz.sum(dim=0)  # dim, num_inducing , num_inducing
        self.Lz = torch.linalg.cholesky(
            Kzz + torch.eye(self.num_inducing).to(self.Z) * 1e-6
        )
        eps = torch.randn(self.dim, self.num_inducing, batch_size).to(self.Z)
        self.delta_Bz = torch.matmul(self.Lz, eps)
//  [python]
def test_multivariate_normal_lazy(self, cuda=False):
        device = torch.device("cuda") if cuda else torch.device("cpu")
        for dtype in (torch.float, torch.double):
            mean = torch.tensor([0, 1, 2], device=device, dtype=dtype)
            covmat = torch.diag(torch.tensor([1, 0.75, 1.5], device=device, dtype=dtype))
            covmat_chol = torch.cholesky(covmat)
            mvn = MultivariateNormal(mean=mean, covariance_matrix=DenseLinearOperator(covmat))
            self.assertTrue(torch.is_tensor(mvn.covariance_matrix))
            self.assertIsInstance(mvn.lazy_covariance_matrix, LinearOperator)
            self.assertAllClose(mvn.variance, torch.diag(covmat))
            self.assertAllClose(mvn.covariance_matrix, covmat)
            self.assertAllClose(mvn._unbroadcasted_scale_tril, covmat_chol)
            mvn_plus1 = mvn + 1
            self.assertAllClose(mvn_plus1.mean, mvn.mean + 1)
            self.assertAllClose(mvn_plus1.covariance_matrix, mvn.covariance_matrix)
            self.assertAllClose(mvn_plus1._unbroadcasted_scale_tril, covmat_chol)
            mvn_times2 = mvn * 2
            self.assertAllClose(mvn_times2.mean, mvn.mean * 2)
            self.assertAllClose(mvn_times2.covariance_matrix, mvn.covariance_matrix * 4)
            self.assertAllClose(mvn_times2._unbroadcasted_scale_tril, covmat_chol * 2)
            mvn_divby2 = mvn / 2
            self.assertAllClose(mvn_divby2.mean, mvn.mean / 2)
            self.assertAllClose(mvn_divby2.covariance_matrix, mvn.covariance_matrix / 4)
            self.assertAllClose(mvn_divby2._unbroadcasted_scale_tril, covmat_chol / 2)
            # TODO: Add tests for entropy, log_prob, etc. - this an issue b/c it
            # uses using root_decomposition which is not very reliable
            # self.assertAlmostEqual(mvn.entropy().item(), 4.3157, places=4)
            # self.assertAlmostEqual(mvn.log_prob(torch.zeros(3)).item(), -4.8157, places=4)
            # self.assertTrue(
            #     torch.allclose(
            #         mvn.log_prob(torch.zeros(2, 3)), -4.8157 * torch.ones(2))
            #     )
            # )
            conf_lower, conf_upper = mvn.confidence_region()
            self.assertAllClose(conf_lower, mvn.mean - 2 * mvn.stddev)
            self.assertAllClose(conf_upper, mvn.mean + 2 * mvn.stddev)
            self.assertTrue(mvn.sample().shape == torch.Size([3]))
            self.assertTrue(mvn.sample(torch.Size([2])).shape == torch.Size([2, 3]))
            self.assertTrue(mvn.sample(torch.Size([2, 4])).shape == torch.Size([2, 4, 3]))
//  [python]
def test_multivariate_normal_lazy(self, cuda=False):
        device = torch.device("cuda") if cuda else torch.device("cpu")
        for dtype in (torch.float, torch.double):
            mean = torch.tensor([0, 1, 2], device=device, dtype=dtype)
            covmat = torch.diag(torch.tensor([1, 0.75, 1.5], device=device, dtype=dtype))
            covmat_chol = torch.linalg.cholesky(covmat)
            mvn = MultivariateNormal(mean=mean, covariance_matrix=DenseLinearOperator(covmat))
            self.assertTrue(torch.is_tensor(mvn.covariance_matrix))
            self.assertIsInstance(mvn.lazy_covariance_matrix, LinearOperator)
            self.assertAllClose(mvn.variance, torch.diag(covmat))
            self.assertAllClose(mvn.covariance_matrix, covmat)
            self.assertAllClose(mvn._unbroadcasted_scale_tril, covmat_chol)
            mvn_plus1 = mvn + 1
            self.assertAllClose(mvn_plus1.mean, mvn.mean + 1)
            self.assertAllClose(mvn_plus1.covariance_matrix, mvn.covariance_matrix)
            self.assertAllClose(mvn_plus1._unbroadcasted_scale_tril, covmat_chol)
            mvn_times2 = mvn * 2
            self.assertAllClose(mvn_times2.mean, mvn.mean * 2)
            self.assertAllClose(mvn_times2.covariance_matrix, mvn.covariance_matrix * 4)
            self.assertAllClose(mvn_times2._unbroadcasted_scale_tril, covmat_chol * 2)
            mvn_divby2 = mvn / 2
            self.assertAllClose(mvn_divby2.mean, mvn.mean / 2)
            self.assertAllClose(mvn_divby2.covariance_matrix, mvn.covariance_matrix / 4)
            self.assertAllClose(mvn_divby2._unbroadcasted_scale_tril, covmat_chol / 2)
            # TODO: Add tests for entropy, log_prob, etc. - this an issue b/c it
            # uses using root_decomposition which is not very reliable
            # self.assertAlmostEqual(mvn.entropy().item(), 4.3157, places=4)
            # self.assertAlmostEqual(mvn.log_prob(torch.zeros(3)).item(), -4.8157, places=4)
            # self.assertTrue(
            #     torch.allclose(
            #         mvn.log_prob(torch.zeros(2, 3)), -4.8157 * torch.ones(2))
            #     )
            # )
            conf_lower, conf_upper = mvn.confidence_region()
            self.assertAllClose(conf_lower, mvn.mean - 2 * mvn.stddev)
            self.assertAllClose(conf_upper, mvn.mean + 2 * mvn.stddev)
            self.assertTrue(mvn.sample().shape == torch.Size([3]))
            self.assertTrue(mvn.sample(torch.Size([2])).shape == torch.Size([2, 3]))
            self.assertTrue(mvn.sample(torch.Size([2, 4])).shape == torch.Size([2, 4, 3]))
//  [python]
def test_multivariate_normal_batch_lazy(self, cuda=False):
        device = torch.device("cuda") if cuda else torch.device("cpu")
        for dtype in (torch.float, torch.double):
            mean = torch.tensor([0, 1, 2], device=device, dtype=dtype).repeat(2, 1)
            covmat = torch.diag(torch.tensor([1, 0.75, 1.5], device=device, dtype=dtype)).repeat(2, 1, 1)
            covmat_chol = torch.cholesky(covmat)
            mvn = MultivariateNormal(mean=mean, covariance_matrix=DenseLinearOperator(covmat))
            self.assertTrue(torch.is_tensor(mvn.covariance_matrix))
            self.assertIsInstance(mvn.lazy_covariance_matrix, LinearOperator)
            self.assertAllClose(mvn.variance, torch.diagonal(covmat, dim1=-2, dim2=-1))
            self.assertAllClose(mvn._unbroadcasted_scale_tril, covmat_chol)
            mvn_plus1 = mvn + 1
            self.assertAllClose(mvn_plus1.mean, mvn.mean + 1)
            self.assertAllClose(mvn_plus1.covariance_matrix, mvn.covariance_matrix)
            self.assertAllClose(mvn_plus1._unbroadcasted_scale_tril, covmat_chol)
            mvn_times2 = mvn * 2
            self.assertAllClose(mvn_times2.mean, mvn.mean * 2)
            self.assertAllClose(mvn_times2.covariance_matrix, mvn.covariance_matrix * 4)
            self.assertAllClose(mvn_times2._unbroadcasted_scale_tril, covmat_chol * 2)
            mvn_divby2 = mvn / 2
            self.assertAllClose(mvn_divby2.mean, mvn.mean / 2)
            self.assertAllClose(mvn_divby2.covariance_matrix, mvn.covariance_matrix / 4)
            self.assertAllClose(mvn_divby2._unbroadcasted_scale_tril, covmat_chol / 2)
            # TODO: Add tests for entropy, log_prob, etc. - this an issue b/c it
            # uses using root_decomposition which is not very reliable
            # self.assertTrue(torch.allclose(mvn.entropy(), 4.3157 * torch.ones(2)))
            # self.assertTrue(
            #     torch.allclose(mvn.log_prob(torch.zeros(2, 3)), -4.8157 * torch.ones(2))
            # )
            # self.assertTrue(
            #     torch.allclose(mvn.log_prob(torch.zeros(2, 2, 3)), -4.8157 * torch.ones(2, 2))
            # )
            conf_lower, conf_upper = mvn.confidence_region()
            self.assertAllClose(conf_lower, mvn.mean - 2 * mvn.stddev)
            self.assertAllClose(conf_upper, mvn.mean + 2 * mvn.stddev)
            self.assertTrue(mvn.sample().shape == torch.Size([2, 3]))
            self.assertTrue(mvn.sample(torch.Size([2])).shape == torch.Size([2, 2, 3]))
            self.assertTrue(mvn.sample(torch.Size([2, 4])).shape == torch.Size([2, 4, 2, 3]))
//  [python]
def test_multivariate_normal_batch_lazy(self, cuda=False):
        device = torch.device("cuda") if cuda else torch.device("cpu")
        for dtype in (torch.float, torch.double):
            mean = torch.tensor([0, 1, 2], device=device, dtype=dtype).repeat(2, 1)
            covmat = torch.diag(torch.tensor([1, 0.75, 1.5], device=device, dtype=dtype)).repeat(2, 1, 1)
            covmat_chol = torch.linalg.cholesky(covmat)
            mvn = MultivariateNormal(mean=mean, covariance_matrix=DenseLinearOperator(covmat))
            self.assertTrue(torch.is_tensor(mvn.covariance_matrix))
            self.assertIsInstance(mvn.lazy_covariance_matrix, LinearOperator)
            self.assertAllClose(mvn.variance, torch.diagonal(covmat, dim1=-2, dim2=-1))
            self.assertAllClose(mvn._unbroadcasted_scale_tril, covmat_chol)
            mvn_plus1 = mvn + 1
            self.assertAllClose(mvn_plus1.mean, mvn.mean + 1)
            self.assertAllClose(mvn_plus1.covariance_matrix, mvn.covariance_matrix)
            self.assertAllClose(mvn_plus1._unbroadcasted_scale_tril, covmat_chol)
            mvn_times2 = mvn * 2
            self.assertAllClose(mvn_times2.mean, mvn.mean * 2)
            self.assertAllClose(mvn_times2.covariance_matrix, mvn.covariance_matrix * 4)
            self.assertAllClose(mvn_times2._unbroadcasted_scale_tril, covmat_chol * 2)
            mvn_divby2 = mvn / 2
            self.assertAllClose(mvn_divby2.mean, mvn.mean / 2)
            self.assertAllClose(mvn_divby2.covariance_matrix, mvn.covariance_matrix / 4)
            self.assertAllClose(mvn_divby2._unbroadcasted_scale_tril, covmat_chol / 2)
            # TODO: Add tests for entropy, log_prob, etc. - this an issue b/c it
            # uses using root_decomposition which is not very reliable
            # self.assertTrue(torch.allclose(mvn.entropy(), 4.3157 * torch.ones(2)))
            # self.assertTrue(
            #     torch.allclose(mvn.log_prob(torch.zeros(2, 3)), -4.8157 * torch.ones(2))
            # )
            # self.assertTrue(
            #     torch.allclose(mvn.log_prob(torch.zeros(2, 2, 3)), -4.8157 * torch.ones(2, 2))
            # )
            conf_lower, conf_upper = mvn.confidence_region()
            self.assertAllClose(conf_lower, mvn.mean - 2 * mvn.stddev)
            self.assertAllClose(conf_upper, mvn.mean + 2 * mvn.stddev)
            self.assertTrue(mvn.sample().shape == torch.Size([2, 3]))
            self.assertTrue(mvn.sample(torch.Size([2])).shape == torch.Size([2, 2, 3]))
            self.assertTrue(mvn.sample(torch.Size([2, 4])).shape == torch.Size([2, 4, 2, 3]))
//  [python]
def safe_cholesky(A, epsilon=1.0e-8):
    """
    Equivalent of torch.linalg.cholesky that progressively adds diagonal jitter to avoid cholesky errors.
    """
    if A.shape == (1, 1):
        return A.sqrt()
    try:
        return torch.cholesky(A)
    except RuntimeError as e:
        Aprime = A.clone()
        jitter_prev = 0.0
        for i in range(5):
            jitter_new = epsilon * (10 ** i)
            Aprime.diagonal(dim1=-2, dim2=-1).add_(jitter_new - jitter_prev)
            jitter_prev = jitter_new
            try:
                return torch.cholesky(Aprime)
            except RuntimeError:
                continue
        raise e
//  [python]
def safe_cholesky(A, epsilon=1.0e-8):
    """
    Equivalent of torch.linalg.cholesky that progressively adds diagonal jitter to avoid cholesky errors.
    """
    if A.shape == (1, 1):
        return A.sqrt()
    try:
        return torch.linalg.cholesky(A)
    except RuntimeError as e:
        Aprime = A.clone()
        jitter_prev = 0.0
        for i in range(5):
            jitter_new = epsilon * (10 ** i)
            Aprime.diagonal(dim1=-2, dim2=-1).add_(jitter_new - jitter_prev)
            jitter_prev = jitter_new
            try:
                return torch.linalg.cholesky(Aprime)
            except RuntimeError:
                continue
        raise e
//  [python]
def test_cg_with_tridiag(self):
        size = 10
        matrix = torch.randn(size, size, dtype=torch.float64)
        matrix = matrix.matmul(matrix.mT)
        matrix.div_(matrix.norm())
        matrix.add_(torch.eye(matrix.size(-1), dtype=torch.float64).mul_(1e-1))

        rhs = torch.randn(size, 50, dtype=torch.float64)
        with warnings.catch_warnings(record=True) as ws:
            solves, t_mats = linear_cg(
                matrix.matmul,
                rhs=rhs,
                n_tridiag=5,
                max_tridiag_iter=10,
                max_iter=size,
                tolerance=0,
                eps=1e-15,
            )
            self.assertTrue(any(issubclass(w.category, NumericalWarning) for w in ws))

        # Check cg
        matrix_chol = torch.cholesky(matrix)
        actual = torch.cholesky_solve(rhs, matrix_chol)
        self.assertTrue(torch.allclose(solves, actual, atol=1e-3, rtol=1e-4))

        # Check tridiag
        eigs = torch.linalg.eigvalsh(matrix)
        for i in range(5):
            approx_eigs = torch.linalg.eigvalsh(t_mats[i])
            self.assertTrue(torch.allclose(eigs, approx_eigs, atol=1e-3, rtol=1e-4))
//  [python]
def test_cg_with_tridiag(self):
        size = 10
        matrix = torch.randn(size, size, dtype=torch.float64)
        matrix = matrix.matmul(matrix.mT)
        matrix.div_(matrix.norm())
        matrix.add_(torch.eye(matrix.size(-1), dtype=torch.float64).mul_(1e-1))

        rhs = torch.randn(size, 50, dtype=torch.float64)
        with warnings.catch_warnings(record=True) as ws:
            solves, t_mats = linear_cg(
                matrix.matmul,
                rhs=rhs,
                n_tridiag=5,
                max_tridiag_iter=10,
                max_iter=size,
                tolerance=0,
                eps=1e-15,
            )
            self.assertTrue(any(issubclass(w.category, NumericalWarning) for w in ws))

        # Check cg
        matrix_chol = torch.linalg.cholesky(matrix)
        actual = torch.cholesky_solve(rhs, matrix_chol)
        self.assertTrue(torch.allclose(solves, actual, atol=1e-3, rtol=1e-4))

        # Check tridiag
        eigs = torch.linalg.eigvalsh(matrix)
        for i in range(5):
            approx_eigs = torch.linalg.eigvalsh(t_mats[i])
            self.assertTrue(torch.allclose(eigs, approx_eigs, atol=1e-3, rtol=1e-4))
//  [python]
def matrix2ellipse(matrixes):
    Rinv = torch.inverse(torch.cholesky(matrixes, upper=True))
    t = torch.linspace(-np.pi, np.pi, 2000)
    xx = torch.cos(t)
    yy = torch.sin(t)
    circle = torch.stack((xx, yy))
    ellipses = Rinv @ circle
    return ellipses
//  [python]
def matrix2ellipse(matrixes):
    Rinv = torch.linalg.inv(torch.linalg.cholesky(matrixes, upper=True))
    t = torch.linspace(-np.pi, np.pi, 2000)
    xx = torch.cos(t)
    yy = torch.sin(t)
    circle = torch.stack((xx, yy))
    ellipses = Rinv @ circle
    return ellipses
//  [python]
def test_batch_cg_with_tridiag(self):
        batch = 5
        size = 10
        matrix = torch.randn(batch, size, size, dtype=torch.float64)
        matrix = matrix.matmul(matrix.mT)
        matrix.div_(matrix.norm())
        matrix.add_(torch.eye(matrix.size(-1), dtype=torch.float64).mul_(1e-1))

        rhs = torch.randn(batch, size, 10, dtype=torch.float64)
        with warnings.catch_warnings(record=True) as ws:
            solves, t_mats = linear_cg(
                matrix.matmul,
                rhs=rhs,
                n_tridiag=8,
                max_iter=size,
                max_tridiag_iter=10,
                tolerance=0,
                eps=1e-30,
            )
            self.assertTrue(any(issubclass(w.category, NumericalWarning) for w in ws))

        # Check cg
        matrix_chol = torch.cholesky(matrix)
        actual = torch.cholesky_solve(rhs, matrix_chol)
        self.assertTrue(torch.allclose(solves, actual, atol=1e-3, rtol=1e-4))

        # Check tridiag
        for i in range(5):
            eigs = torch.linalg.eigvalsh(matrix[i])
            for j in range(8):
                approx_eigs = torch.linalg.eigvalsh(t_mats[j, i])
                self.assertTrue(torch.allclose(eigs, approx_eigs, atol=1e-3, rtol=1e-4))
//  [python]
def test_batch_cg_with_tridiag(self):
        batch = 5
        size = 10
        matrix = torch.randn(batch, size, size, dtype=torch.float64)
        matrix = matrix.matmul(matrix.mT)
        matrix.div_(matrix.norm())
        matrix.add_(torch.eye(matrix.size(-1), dtype=torch.float64).mul_(1e-1))

        rhs = torch.randn(batch, size, 10, dtype=torch.float64)
        with warnings.catch_warnings(record=True) as ws:
            solves, t_mats = linear_cg(
                matrix.matmul,
                rhs=rhs,
                n_tridiag=8,
                max_iter=size,
                max_tridiag_iter=10,
                tolerance=0,
                eps=1e-30,
            )
            self.assertTrue(any(issubclass(w.category, NumericalWarning) for w in ws))

        # Check cg
        matrix_chol = torch.linalg.cholesky(matrix)
        actual = torch.cholesky_solve(rhs, matrix_chol)
        self.assertTrue(torch.allclose(solves, actual, atol=1e-3, rtol=1e-4))

        # Check tridiag
        for i in range(5):
            eigs = torch.linalg.eigvalsh(matrix[i])
            for j in range(8):
                approx_eigs = torch.linalg.eigvalsh(t_mats[j, i])
                self.assertTrue(torch.allclose(eigs, approx_eigs, atol=1e-3, rtol=1e-4))
//  [python]
def marginal_log_likelihood(
    X: torch.Tensor,  # [len(trials), len(params)]
    Y: torch.Tensor,  # [len(trials)]
    is_categorical: torch.Tensor,  # [len(params)]
    kernel_params: KernelParamsTensor,
) -> torch.Tensor:  # Scalar
    # -0.5 * log((2pi)^n |C|) - 0.5 * Y^T C^-1 Y, where C^-1 = cov_Y_Y_inv
    # We apply the cholesky decomposition to efficiently compute log(|C|) and C^-1.

    cov_fX_fX = kernel(is_categorical, kernel_params, X, X)

    cov_Y_Y_chol = torch.cholesky(
        cov_fX_fX + kernel_params.noise_var * torch.eye(X.shape[0], dtype=torch.float64)
    )
    # log |L| = 0.5 * log|L^T L| = 0.5 * log|C|
    logdet = 2 * torch.log(torch.diag(cov_Y_Y_chol)).sum()
    # cov_Y_Y_chol @ cov_Y_Y_chol_inv_Y = Y --> cov_Y_Y_chol_inv_Y = inv(cov_Y_Y_chol) @ Y
    cov_Y_Y_chol_inv_Y = torch.triangular_solve(Y[:, None], cov_Y_Y_chol, upper=False)[0][:, 0]
    return -0.5 * (
        logdet
        + X.shape[0] * math.log(2 * math.pi)
        # Y^T C^-1 Y = Y^T inv(L^T L) Y --> cov_Y_Y_chol_inv_Y @ cov_Y_Y_chol_inv_Y
        + (cov_Y_Y_chol_inv_Y @ cov_Y_Y_chol_inv_Y)
    )
//  [python]
def marginal_log_likelihood(
    X: torch.Tensor,  # [len(trials), len(params)]
    Y: torch.Tensor,  # [len(trials)]
    is_categorical: torch.Tensor,  # [len(params)]
    kernel_params: KernelParamsTensor,
) -> torch.Tensor:  # Scalar
    # -0.5 * log((2pi)^n |C|) - 0.5 * Y^T C^-1 Y, where C^-1 = cov_Y_Y_inv
    # We apply the cholesky decomposition to efficiently compute log(|C|) and C^-1.

    cov_fX_fX = kernel(is_categorical, kernel_params, X, X)

    cov_Y_Y_chol = torch.linalg.cholesky(
        cov_fX_fX + kernel_params.noise_var * torch.eye(X.shape[0], dtype=torch.float64)
    )
    # log |L| = 0.5 * log|L^T L| = 0.5 * log|C|
    logdet = 2 * torch.log(torch.diag(cov_Y_Y_chol)).sum()
    # cov_Y_Y_chol @ cov_Y_Y_chol_inv_Y = Y --> cov_Y_Y_chol_inv_Y = inv(cov_Y_Y_chol) @ Y
    cov_Y_Y_chol_inv_Y = torch.linalg.solve_triangular(cov_Y_Y_chol, Y[:, None], upper=False)[:, 0]
    return -0.5 * (
        logdet
        + X.shape[0] * math.log(2 * math.pi)
        # Y^T C^-1 Y = Y^T inv(L^T L) Y --> cov_Y_Y_chol_inv_Y @ cov_Y_Y_chol_inv_Y
        + (cov_Y_Y_chol_inv_Y @ cov_Y_Y_chol_inv_Y)
    )
//  [python]
def compute_loss(self, inputs, outputs):
        self.inputs_tr = inputs
        self.outputs_tr = outputs
        self.already_set_train_data = True

        _noise = self._get_noise_according_exp_format()
        x = inputs[0]
        y = outputs[0]
        l1, h1, l2, h2 = self.l1, self.h1, self.l2, self.h2

        sigma = self.kernel_list[0](x, x, l1, h1, l2, h2) + _noise.pow(-1) * torch.eye(x.size(0), device=list(self.parameters())[0].device)
        sigma = sigma + JITTER * torch.eye(x.size(0), device=list(self.parameters())[0].device)

        L = torch.cholesky(sigma)
        y_num, y_dimension = y.shape
        Gamma,_ = torch.triangular_solve(y, L, upper = False)
        nll =  0.5 * (Gamma ** 2).sum() +  L.diag().log().sum() * y_dimension  \
            + 0.5 * y_num * torch.log(2 * torch.tensor(PI)) * y_dimension

        return nll
//  [python]
def compute_loss(self, inputs, outputs):
        self.inputs_tr = inputs
        self.outputs_tr = outputs
        self.already_set_train_data = True

        _noise = self._get_noise_according_exp_format()
        x = inputs[0]
        y = outputs[0]
        l1, h1, l2, h2 = self.l1, self.h1, self.l2, self.h2

        sigma = self.kernel_list[0](x, x, l1, h1, l2, h2) + _noise.pow(-1) * torch.eye(x.size(0), device=list(self.parameters())[0].device)
        sigma = sigma + JITTER * torch.eye(x.size(0), device=list(self.parameters())[0].device)

        L = torch.linalg.cholesky(sigma)
        y_num, y_dimension = y.shape
        Gamma,_ = torch.triangular_solve(y, L, upper = False)
        nll =  0.5 * (Gamma ** 2).sum() +  L.diag().log().sum() * y_dimension  \
            + 0.5 * y_num * torch.log(2 * torch.tensor(PI)) * y_dimension

        return nll
//  [python]
def _cholesky(x):
    """
    Like :func:`torch.linalg.cholesky` but uses sqrt for scalar matrices.
    Works around https://github.com/pytorch/pytorch/issues/24403 often.
    """
    if x.size(-1) == 1:
        return x.sqrt()
    return torch.cholesky(x)
//  [python]
def _cholesky(x):
    """
    Like :func:`torch.linalg.cholesky` but uses sqrt for scalar matrices.
    Works around https://github.com/pytorch/pytorch/issues/24403 often.
    """
    if x.size(-1) == 1:
        return x.sqrt()
    return torch.linalg.cholesky(x)