jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_old_cholesky(self, device, dtype):
        from torch.testing._internal.common_utils import random_hermitian_pd_matrix

        A = random_hermitian_pd_matrix(10, dtype=dtype, device=device)

        # default Case
        C = torch.linalg.cholesky(A)
        B = torch.mm(C, C.t().conj())
        self.assertEqual(A, B, atol=1e-14, rtol=0)

        # test Upper Triangular
        U = torch.linalg.cholesky(A, upper=True)
        B = torch.mm(U.t().conj(), U)
        self.assertEqual(A, B, atol=1e-14, rtol=0, msg='cholesky (upper) did not allow rebuilding the original matrix')

        # test Lower Triangular
        L = torch.linalg.cholesky(A, upper=False)
        B = torch.mm(L, L.t().conj())
        self.assertEqual(A, B, atol=1e-14, rtol=0, msg='cholesky (lower) did not allow rebuilding the original matrix')
//  [python]
def cholesky_solve_test_helper(self, A_dims, b_dims, upper, device, dtype):
        from torch.testing._internal.common_utils import random_hermitian_pd_matrix

        b = torch.randn(*b_dims, dtype=dtype, device=device)
        A = random_hermitian_pd_matrix(*A_dims, dtype=dtype, device=device)
        L = torch.cholesky(A, upper=upper)
        return b, A, L
//  [python]
def cholesky_solve_test_helper(self, A_dims, b_dims, upper, device, dtype):
        from torch.testing._internal.common_utils import random_hermitian_pd_matrix

        b = torch.randn(*b_dims, dtype=dtype, device=device)
        A = random_hermitian_pd_matrix(*A_dims, dtype=dtype, device=device)
        L = torch.linalg.cholesky(A, upper=upper)
        return b, A, L
//  [python]
def test_old_cholesky_autograd(self, device, dtype):
        def func(root, upper):
            x = 0.5 * (root + root.transpose(-1, -2).conj())
            return torch.cholesky(x, upper)

        def run_test(upper, dims):
            root = torch.rand(*dims, dtype=dtype, device=device, requires_grad=True)
            root = root + torch.eye(dims[-1])

            gradcheck(func, [root, upper])
            gradgradcheck(func, [root, upper])

            root = torch.rand(*dims, dtype=dtype, device=device)
            root = torch.matmul(root, root.transpose(-1, -2).conj())
            root.requires_grad_()
            chol = root.cholesky().sum().backward()
            self.assertEqual(root.grad, root.grad.transpose(-1, -2).conj())  # Check the gradient is hermitian

        for upper, dims in itertools.product([True, False], [(3, 3), (4, 3, 2, 2)]):
            run_test(upper, dims)
//  [python]
def test_old_cholesky_autograd(self, device, dtype):
        def func(root, upper):
            x = 0.5 * (root + root.transpose(-1, -2).conj())
            return torch.linalg.cholesky(x, upper=upper)

        def run_test(upper, dims):
            root = torch.rand(*dims, dtype=dtype, device=device, requires_grad=True)
            root = root + torch.eye(dims[-1])

            gradcheck(func, [root, upper])
            gradgradcheck(func, [root, upper])

            root = torch.rand(*dims, dtype=dtype, device=device)
            root = torch.matmul(root, root.transpose(-1, -2).conj())
            root.requires_grad_()
            chol = root.cholesky().sum().backward()
            self.assertEqual(root.grad, root.grad.transpose(-1, -2).conj())  # Check the gradient is hermitian

        for upper, dims in itertools.product([True, False], [(3, 3), (4, 3, 2, 2)]):
            run_test(upper, dims)
//  [python]
def test_cholesky_solve_batched_non_contiguous(self, device, dtype):
        from numpy.linalg import solve
        from torch.testing._internal.common_utils import random_hermitian_pd_matrix

        for upper in [True, False]:
            A = random_hermitian_pd_matrix(2, 2, dtype=dtype, device='cpu')
            b = torch.randn(2, 2, 2, dtype=dtype, device='cpu')
            x_exp = solve(A.permute(0, 2, 1).numpy(), b.permute(2, 1, 0).numpy())
            A = A.to(device).permute(0, 2, 1)
            b = b.to(device).permute(2, 1, 0)
            assert not A.is_contiguous() and not b.is_contiguous(), "contiguous inputs"
            L = torch.cholesky(A, upper)
            x = torch.cholesky_solve(b, L, upper=upper)
            self.assertEqual(x, x_exp)
//  [python]
def test_cholesky_solve_batched_non_contiguous(self, device, dtype):
        from numpy.linalg import solve
        from torch.testing._internal.common_utils import random_hermitian_pd_matrix

        for upper in [True, False]:
            A = random_hermitian_pd_matrix(2, 2, dtype=dtype, device='cpu')
            b = torch.randn(2, 2, 2, dtype=dtype, device='cpu')
            x_exp = solve(A.permute(0, 2, 1).numpy(), b.permute(2, 1, 0).numpy())
            A = A.to(device).permute(0, 2, 1)
            b = b.to(device).permute(2, 1, 0)
            assert not A.is_contiguous() and not b.is_contiguous(), "contiguous inputs"
            L = torch.linalg.cholesky(A, upper=upper)
            x = torch.cholesky_solve(b, L, upper=upper)
            self.assertEqual(x, x_exp)
//  [python]
def forward(ctx, H, b):
        U = torch.cholesky(H)

        xs = torch.cholesky_solve(b, U)
        ctx.save_for_backward(U, xs)
        return xs
//  [python]
def forward(ctx, H, b):
        U = torch.linalg.cholesky(H)

        xs = torch.cholesky_solve(b, U)
        ctx.save_for_backward(U, xs)
        return xs
//  [python]
def batch_capacitance_tril(W, D):
    r"""
    Copied from pytorch source code
    Computes Cholesky of :math:`I + W.T @ inv(D) @ W` for a batch of matrices :math:`W`
    and a batch of vectors :math:`D`.
    """
    m = W.size(-1)
    Wt_Dinv = W.transpose(-1, -2) / D.unsqueeze(-2)
    K = torch.matmul(Wt_Dinv, W).contiguous()
    K.view(-1, m * m)[:, ::m + 1] += 1
    return torch.cholesky(K)
//  [python]
def batch_capacitance_tril(W, D):
    r"""
    Copied from pytorch source code
    Computes Cholesky of :math:`I + W.T @ inv(D) @ W` for a batch of matrices :math:`W`
    and a batch of vectors :math:`D`.
    """
    m = W.size(-1)
    Wt_Dinv = W.transpose(-1, -2) / D.unsqueeze(-2)
    K = torch.matmul(Wt_Dinv, W).contiguous()
    K.view(-1, m * m)[:, ::m + 1] += 1
    return torch.linalg.cholesky(K)
//  [python]
def test_cholesky_inverse(self, device, dtype):
        from torch.testing._internal.common_utils import random_symmetric_pd_matrix
        a = random_symmetric_pd_matrix(5, dtype=dtype, device=device)

        # compute inverse directly
        inv0 = torch.inverse(a)

        # default case
        chol = torch.cholesky(a)
        inv1 = torch.cholesky_inverse(chol, False)
        self.assertLessEqual(inv0.dist(inv1), 1e-12)

        # upper Triangular Test
        chol = torch.cholesky(a, True)
        inv1 = torch.cholesky_inverse(chol, True)
        self.assertLessEqual(inv0.dist(inv1), 1e-12)

        # lower Triangular Test
        chol = torch.cholesky(a, False)
        inv1 = torch.cholesky_inverse(chol, False)
        self.assertLessEqual(inv0.dist(inv1), 1e-12)
//  [python]
def test_cholesky_inverse(self, device, dtype):
        from torch.testing._internal.common_utils import random_symmetric_pd_matrix
        a = random_symmetric_pd_matrix(5, dtype=dtype, device=device)

        # compute inverse directly
        inv0 = torch.inverse(a)

        # default case
        chol = torch.linalg.cholesky(a)
        inv1 = torch.cholesky_inverse(chol, False)
        self.assertLessEqual(inv0.dist(inv1), 1e-12)

        # upper Triangular Test
        chol = torch.linalg.cholesky(a, upper=True)
        inv1 = torch.cholesky_inverse(chol, True)
        self.assertLessEqual(inv0.dist(inv1), 1e-12)

        # lower Triangular Test
        chol = torch.linalg.cholesky(a, upper=False)
        inv1 = torch.cholesky_inverse(chol, False)
        self.assertLessEqual(inv0.dist(inv1), 1e-12)
//  [python]
def estimate_variance(self, var0, invert=True):
        tau = 1/var0

        U, V = self.hessians

        for module, name in self.params:
            mod_class = module.__class__.__name__
            if mod_class not in ['Linear', 'Conv2d']:
                continue

            if name == 'bias':
                continue

            U_ = U[(module, name)].clone()
            V_ = V[(module, name)].clone()

            if invert:
                m, n = int(U_.shape[0]), int(V_.shape[0])

                U_ += torch.sqrt(tau)*torch.eye(m, device='cuda')
                V_ += torch.sqrt(tau)*torch.eye(n, device='cuda')

                U_ = torch.cholesky(torch.inverse(U_), upper=False)
                V_ = torch.cholesky(torch.inverse(V_), upper=True)

            module.__getattr__(f'{name}_U_half').copy_(U_)
            module.__getattr__(f'{name}_V_half').copy_(V_)
//  [python]
def estimate_variance(self, var0, invert=True):
        tau = 1/var0

        U, V = self.hessians

        for module, name in self.params:
            mod_class = module.__class__.__name__
            if mod_class not in ['Linear', 'Conv2d']:
                continue

            if name == 'bias':
                continue

            U_ = U[(module, name)].clone()
            V_ = V[(module, name)].clone()

            if invert:
                m, n = int(U_.shape[0]), int(V_.shape[0])

                U_ += torch.sqrt(tau)*torch.eye(m, device='cuda')
                V_ += torch.sqrt(tau)*torch.eye(n, device='cuda')

                U_ = torch.linalg.cholesky(torch.inverse(U_), upper=False)
                V_ = torch.linalg.cholesky(torch.inverse(V_), upper=True)

            module.__getattr__(f'{name}_U_half').copy_(U_)
            module.__getattr__(f'{name}_V_half').copy_(V_)
//  [python]
def inverse(self, method='torch'):
        assert method in ['cholesky', 'torch', 'compress_inverse']
        if method == 'cholesky':
            u = torch.cholesky(-self.mat)
            matinv = - torch.cholesky_inverse(u)
        elif method == 'torch':
            matinv = torch.inverse(self.mat)
        else:
            # compress inverse
            matinv = self.smat.inverse()
        self.matinv = matinv
        return matinv
//  [python]
def inverse(self, method='torch'):
        assert method in ['cholesky', 'torch', 'compress_inverse']
        if method == 'cholesky':
            u = torch.linalg.cholesky(-self.mat)
            matinv = - torch.linalg.cholesky_inverse(u)
        elif method == 'torch':
            matinv = torch.inverse(self.mat)
        else:
            # compress inverse
            matinv = self.smat.inverse()
        self.matinv = matinv
        return matinv
//  [python]
def test_cholesky_solve_batched_broadcasting(self, device, dtype):
        from numpy.linalg import solve
        from torch.testing._internal.common_utils import random_hermitian_pd_matrix

        def run_test(A_dims, b_dims, upper):
            A_matrix_size = A_dims[-1]
            A_batch_dims = A_dims[:-2]
            A = random_hermitian_pd_matrix(A_matrix_size, *A_batch_dims,
                                           dtype=dtype, device='cpu')
            b = torch.randn(*b_dims, dtype=dtype, device='cpu')
            x_exp = torch.tensor(solve(A.numpy(), b.numpy()), dtype=dtype, device=device)
            A, b = A.to(dtype=dtype, device=device), b.to(dtype=dtype, device=device)
            L = torch.cholesky(A, upper)
            x = torch.cholesky_solve(b, L, upper=upper)
            self.assertEqual(x, x_exp)
            # https://github.com/pytorch/pytorch/issues/42695
            x = torch.cholesky_solve(b, L, upper=upper, out=x)
            self.assertEqual(x, x_exp)

        # test against numpy.linalg.solve
        for upper in [True, False]:
            run_test((2, 1, 3, 4, 4), (2, 1, 3, 4, 6), upper)  # no broadcasting
            run_test((2, 1, 3, 4, 4), (4, 6), upper)  # broadcasting b
            run_test((4, 4), (2, 1, 3, 4, 2), upper)  # broadcasting A
            run_test((1, 3, 1, 4, 4), (2, 1, 3, 4, 5), upper)
//  [python]
def test_cholesky_solve_batched_broadcasting(self, device, dtype):
        from numpy.linalg import solve
        from torch.testing._internal.common_utils import random_hermitian_pd_matrix

        def run_test(A_dims, b_dims, upper):
            A_matrix_size = A_dims[-1]
            A_batch_dims = A_dims[:-2]
            A = random_hermitian_pd_matrix(A_matrix_size, *A_batch_dims,
                                           dtype=dtype, device='cpu')
            b = torch.randn(*b_dims, dtype=dtype, device='cpu')
            x_exp = torch.tensor(solve(A.numpy(), b.numpy()), dtype=dtype, device=device)
            A, b = A.to(dtype=dtype, device=device), b.to(dtype=dtype, device=device)
            L = torch.linalg.cholesky(A, upper)
            x = torch.cholesky_solve(b, L, upper=upper)
            self.assertEqual(x, x_exp)
            # https://github.com/pytorch/pytorch/issues/42695
            x = torch.cholesky_solve(b, L, upper=upper, out=x)
            self.assertEqual(x, x_exp)

        # test against numpy.linalg.solve
        for upper in [True, False]:
            run_test((2, 1, 3, 4, 4), (2, 1, 3, 4, 6), upper)  # no broadcasting
            run_test((2, 1, 3, 4, 4), (4, 6), upper)  # broadcasting b
            run_test((4, 4), (2, 1, 3, 4, 2), upper)  # broadcasting A
            run_test((1, 3, 1, 4, 4), (2, 1, 3, 4, 5), upper)
//  [python]
def _posterior_w(self, x, y):
        "returns mean and lower triangular precision of whitened p(w | x,y)"
        f = self.net(x) * self.last_layer_std
        sig = prior.value_or_call(self.noise_std)**2
        # Precision matrix
        A = (f.t()@f)/sig + torch.eye(f.size(-1), dtype=f.dtype, device=f.device)
        # switch to float64
        L = torch.cholesky(A.to(torch.float64))
        FY = (f.t()@y).to(torch.float64)
        white_mean = FY.triangular_solve(L, upper=False).solution
        return white_mean/sig, L
//  [python]
def _posterior_w(self, x, y):
        "returns mean and lower triangular precision of whitened p(w | x,y)"
        f = self.net(x) * self.last_layer_std
        sig = prior.value_or_call(self.noise_std)**2
        # Precision matrix
        A = (f.t()@f)/sig + torch.eye(f.size(-1), dtype=f.dtype, device=f.device)
        # switch to float64
        L = torch.linalg.cholesky(A.to(torch.float64))
        FY = (f.t()@y).to(torch.float64)
        white_mean = FY.triangular_solve(L, upper=False).solution
        return white_mean/sig, L
//  [python]
def _precision_to_scale_tril(P):
    # Ref: https://nbviewer.jupyter.org/gist/fehiepsi/5ef8e09e61604f10607380467eb82006#Precision-to-scale_tril
    Lf = torch.cholesky(torch.flip(P, (-2, -1)))
    L_inv = torch.transpose(torch.flip(Lf, (-2, -1)), -2, -1)
    L = torch.triangular_solve(torch.eye(P.shape[-1], dtype=P.dtype, device=P.device),
                               L_inv, upper=False)[0]
    return L
//  [python]
def _precision_to_scale_tril(P):
    Lf = torch.linalg.cholesky(torch.flip(P, (-2, -1)))
    L_inv = torch.transpose(torch.flip(Lf, (-2, -1)), -2, -1)
    L = torch.triangular_solve(torch.eye(P.shape[-1], dtype=P.dtype, device=P.device),
                               L_inv, upper=False)[0]
    return L
//  [python]
def partial_kernel_mmd_qp(x, y, n_perm=None, kernel=GaussianKernel(),
                          fraction_to_match=1.0):
    """Partial Kernel MMD using quadratic programming.

    This is very slow and mainly intended for reference purposes.
    You need to install qpsolvers to use this function."""

    torchdrift.utils.check(
        n_perm is None,
        "Bootstrapping within partial MMD is not implemented, use bootstrap during fit",
        error_class=NotImplementedError
    )
    import qpsolvers
    n, d = x.shape
    m, d2 = y.shape
    torchdrift.utils.check(d == d2, "feature dimension mismatch")
    xy = torch.cat([x.detach(), y.detach()], dim=0)
    dists = torch.cdist(xy, xy, p=2.0)
    # we are a bit sloppy here as we just keep the diagonal and everything twice
    k = kernel(dists.double())
    k_x = k[:n, :n]
    k_y = k[n:, n:]
    k_xy = k[:n, n:]

    v = torch.full((m,), 1 / m, dtype=k_y.dtype, device=k_y.device)
    R = torch.cholesky(k_x, upper=True)
    d = torch.inverse(R.t()) @ (k_xy.sum(1) / m)
    lb = torch.zeros((n,), dtype=k_x.dtype, device=k_x.device)
    ub = torch.full((n,), 1.0 / (n * fraction_to_match), dtype=k_x.dtype, device=k_x.device)
    w = qpsolvers.solve_ls(R.cpu().numpy(), d.cpu().numpy(),
                           lb=lb.cpu().numpy(), ub=ub.cpu().numpy(),
                           A=torch.ones((1, n,), dtype=R.dtype).numpy(),
                           b=torch.ones((1,), dtype=R.dtype).numpy())
    torchdrift.utils.check(
        w is not None,
        'QP failed to find a solution (numerical accuracy with the bounds?)'
    )
    w = torch.as_tensor(w, device=k_x.device, dtype=k_x.dtype)
    mmd = (w @ k_x) @ w + k_y.sum() / (m * m) - 2 * (w @ k_xy).sum() / m
    return mmd
//  [python]
def partial_kernel_mmd_qp(x, y, n_perm=None, kernel=GaussianKernel(),
                          fraction_to_match=1.0):
    """Partial Kernel MMD using quadratic programming.

    This is very slow and mainly intended for reference purposes.
    You need to install qpsolvers to use this function."""

    torchdrift.utils.check(
        n_perm is None,
        "Bootstrapping within partial MMD is not implemented, use bootstrap during fit",
        error_class=NotImplementedError
    )
    import qpsolvers
    n, d = x.shape
    m, d2 = y.shape
    torchdrift.utils.check(d == d2, "feature dimension mismatch")
    xy = torch.cat([x.detach(), y.detach()], dim=0)
    dists = torch.cdist(xy, xy, p=2.0)
    # we are a bit sloppy here as we just keep the diagonal and everything twice
    k = kernel(dists.double())
    k_x = k[:n, :n]
    k_y = k[n:, n:]
    k_xy = k[:n, n:]

    v = torch.full((m,), 1 / m, dtype=k_y.dtype, device=k_y.device)
    R = torch.linalg.cholesky(k_x, upper=True)
    d = torch.inverse(R.t()) @ (k_xy.sum(1) / m)
    lb = torch.zeros((n,), dtype=k_x.dtype, device=k_x.device)
    ub = torch.full((n,), 1.0 / (n * fraction_to_match), dtype=k_x.dtype, device=k_x.device)
    w = qpsolvers.solve_ls(R.cpu().numpy(), d.cpu().numpy(),
                           lb=lb.cpu().numpy(), ub=ub.cpu().numpy(),
                           A=torch.ones((1, n,), dtype=R.dtype).numpy(),
                           b=torch.ones((1,), dtype=R.dtype).numpy())
    torchdrift.utils.check(
        w is not None,
        'QP failed to find a solution (numerical accuracy with the bounds?)'
    )
    w = torch.as_tensor(w, device=k_x.device, dtype=k_x.dtype)
    mmd = (w @ k_x) @ w + k_y.sum() / (m * m) - 2 * (w @ k_xy).sum() / m
    return mmd
//  [python]
def test_multivariate_normal_properties(self):
        loc = torch.randn(5)
        scale_tril = transform_to(constraints.lower_cholesky)(torch.randn(5, 5))
        m = MultivariateNormal(loc=loc, scale_tril=scale_tril)
        self.assertEqual(m.covariance_matrix, m.scale_tril.mm(m.scale_tril.t()))
        self.assertEqual(m.covariance_matrix.mm(m.precision_matrix), torch.eye(m.event_shape[0]))
        self.assertEqual(m.scale_tril, torch.cholesky(m.covariance_matrix, upper=False))
//  [python]
def test_multivariate_normal_properties(self):
        loc = torch.randn(5)
        scale_tril = transform_to(constraints.lower_cholesky)(torch.randn(5, 5))
        m = MultivariateNormal(loc=loc, scale_tril=scale_tril)
        self.assertEqual(m.covariance_matrix, m.scale_tril.mm(m.scale_tril.t()))
        self.assertEqual(m.covariance_matrix.mm(m.precision_matrix), torch.eye(m.event_shape[0]))
        self.assertEqual(m.scale_tril, torch.linalg.cholesky(m.covariance_matrix, upper=False))
//  [python]
def negELBO(self):
        K_u_u = self.kernel.matrix(self.X_concat, torch.exp(self.log_ls_h))
        K_u1 = self.kernel.cross(self.X_col, self.X_concat, torch.exp(self.log_ls_h))
        K_u2 = self.kernel.cross(self.X_col, self.X_col, torch.exp(self.log_ls_h))
        K_uxx_uxx = self.kernel.matrix_pp_pp(K_u2, self.X_col, 0, torch.exp(self.log_ls_h))
        K_uxx = self.kernel.cross_pp_0(K_u1, self.X_col, self.X_concat, 0, torch.exp(self.log_ls_h))
        K_ux_ux = self.kernel.matrix_p_p(K_u2, self.X_col, 1, torch.exp(self.log_ls_h))
        K_ux = self.kernel.cross_p_0(K_u1, self.X_col, self.X_concat, 1, torch.exp(self.log_ls_h))
        K_ux_uxx = self.kernel.cross_p1_pp0(K_u2, self.X_col, self.X_col, torch.exp(self.log_ls_h))
        K = torch.cat(
            (
                torch.cat((K_u_u, K_uxx.T, K_ux.T), 1),
                torch.cat((K_uxx, K_uxx_uxx, K_ux_uxx.T), 1),
                torch.cat((K_ux, K_ux_uxx, K_ux_ux), 1),
            ),
            0,
        )
        A = torch.cholesky(K)
        cov = torch.diag(self.L1.exp().view(-1)) + torch.matmul(self.L2, self.L2.T)
        self.L = torch.cholesky(cov)
        Ltril = torch.tril(self.L)
        kl = self.KL_term(self.mu, cov)
        eta = self.mu + torch.matmul(Ltril, torch.randn(self.M, 1, device=self.device))
        f = torch.matmul(A, eta)
        elbo = (-kl + 0.5 * self.N * self.log_tau - 0.5 * torch.exp(self.log_tau) * torch.sum(torch.square(f[0:self.N].view(-1, 1) - self.y)) + 0.5 * self.N_col * self.log_v - 0.5 * torch.exp(self.log_v) * torch.sum(
            torch.square(f[self.N_concat + self.N_col:self.N_concat + 2 * self.N_col].view(-1, 1) - 0.0001 * f[self.N_concat:self.N_concat + self.N_col].view(-1, 1) + 5 *
                         (f[self.N:self.N_concat]**3).view(-1, 1) - 5 * f[self.N:self.N_concat].view(-1, 1))))
        return -elbo
//  [python]
def negELBO(self):
        K_u_u = self.kernel.matrix(self.X_concat, torch.exp(self.log_ls_h))
        K_u1 = self.kernel.cross(self.X_col, self.X_concat, torch.exp(self.log_ls_h))
        K_u2 = self.kernel.cross(self.X_col, self.X_col, torch.exp(self.log_ls_h))
        K_uxx_uxx = self.kernel.matrix_pp_pp(K_u2, self.X_col, 0, torch.exp(self.log_ls_h))
        K_uxx = self.kernel.cross_pp_0(K_u1, self.X_col, self.X_concat, 0, torch.exp(self.log_ls_h))
        K_ux_ux = self.kernel.matrix_p_p(K_u2, self.X_col, 1, torch.exp(self.log_ls_h))
        K_ux = self.kernel.cross_p_0(K_u1, self.X_col, self.X_concat, 1, torch.exp(self.log_ls_h))
        K_ux_uxx = self.kernel.cross_p1_pp0(K_u2, self.X_col, self.X_col, torch.exp(self.log_ls_h))
        K = torch.cat(
            (
                torch.cat((K_u_u, K_uxx.T, K_ux.T), 1),
                torch.cat((K_uxx, K_uxx_uxx, K_ux_uxx.T), 1),
                torch.cat((K_ux, K_ux_uxx, K_ux_ux), 1),
            ),
            0,
        )
        A = torch.linalg.cholesky(K)
        cov = torch.diag(self.L1.exp().view(-1)) + torch.matmul(self.L2, self.L2.T)
        self.L = torch.linalg.cholesky(cov)
        Ltril = torch.tril(self.L)
        kl = self.KL_term(self.mu, cov)
        eta = self.mu + torch.matmul(Ltril, torch.randn(self.M, 1, device=self.device))
        f = torch.matmul(A, eta)
        elbo = (-kl + 0.5 * self.N * self.log_tau - 0.5 * torch.exp(self.log_tau) * torch.sum(torch.square(f[0:self.N].view(-1, 1) - self.y)) + 0.5 * self.N_col * self.log_v - 0.5 * torch.exp(self.log_v) * torch.sum(
            torch.square(f[self.N_concat + self.N_col:self.N_concat + 2 * self.N_col].view(-1, 1) - 0.0001 * f[self.N_concat:self.N_concat + self.N_col].view(-1, 1) + 5 *
                         (f[self.N:self.N_concat]**3).view(-1, 1) - 5 * f[self.N:self.N_concat].view(-1, 1))))
        return -elbo
//  [python]
def test_multivariate_normal_sample(self):
        set_rng_seed(0)  # see Note [Randomized statistical tests]
        mean = torch.randn(3, requires_grad=True)
        tmp = torch.randn(3, 10)
        cov = (torch.matmul(tmp, tmp.t()) / tmp.size(-1)).requires_grad_()
        prec = cov.inverse().requires_grad_()
        scale_tril = torch.cholesky(cov, upper=False).requires_grad_()

        self._check_sampler_sampler(MultivariateNormal(mean, cov),
                                    scipy.stats.multivariate_normal(mean.detach().numpy(), cov.detach().numpy()),
                                    'MultivariateNormal(loc={}, cov={})'.format(mean, cov),
                                    multivariate=True)
        self._check_sampler_sampler(MultivariateNormal(mean, precision_matrix=prec),
                                    scipy.stats.multivariate_normal(mean.detach().numpy(), cov.detach().numpy()),
                                    'MultivariateNormal(loc={}, atol={})'.format(mean, prec),
                                    multivariate=True)
        self._check_sampler_sampler(MultivariateNormal(mean, scale_tril=scale_tril),
                                    scipy.stats.multivariate_normal(mean.detach().numpy(), cov.detach().numpy()),
                                    'MultivariateNormal(loc={}, scale_tril={})'.format(mean, scale_tril),
                                    multivariate=True)
//  [python]
def test_multivariate_normal_sample(self):
        set_rng_seed(0)  # see Note [Randomized statistical tests]
        mean = torch.randn(3, requires_grad=True)
        tmp = torch.randn(3, 10)
        cov = (torch.matmul(tmp, tmp.t()) / tmp.size(-1)).requires_grad_()
        prec = cov.inverse().requires_grad_()
        scale_tril = torch.linalg.cholesky(cov, upper=False).requires_grad_()

        self._check_sampler_sampler(MultivariateNormal(mean, cov),
                                    scipy.stats.multivariate_normal(mean.detach().numpy(), cov.detach().numpy()),
                                    'MultivariateNormal(loc={}, cov={})'.format(mean, cov),
                                    multivariate=True)
        self._check_sampler_sampler(MultivariateNormal(mean, precision_matrix=prec),
                                    scipy.stats.multivariate_normal(mean.detach().numpy(), cov.detach().numpy()),
                                    'MultivariateNormal(loc={}, atol={})'.format(mean, prec),
                                    multivariate=True)
        self._check_sampler_sampler(MultivariateNormal(mean, scale_tril=scale_tril),
                                    scipy.stats.multivariate_normal(mean.detach().numpy(), cov.detach().numpy()),
                                    'MultivariateNormal(loc={}, scale_tril={})'.format(mean, scale_tril),
                                    multivariate=True)
//  [python]
def log_det_by_cholesky(matrix):
    """
	Args:
		matrix: matrix must be a positive define matrix.
				shape [N, C, D, D].
	Ref:
		https://github.com/tensorflow/tensorflow/blob/r1.13/tensorflow/python/ops/linalg/linalg_impl.py
	"""
    # This uses the property that the log det(A) = 2 * sum(log(real(diag(C))))
    # where C is the cholesky decomposition of A.
    chol = torch.cholesky(matrix)
    # return 2.0 * torch.sum(torch.log(torch.diagonal(chol, dim1=-2, dim2=-1) + 1e-6), dim=-1)
    return 2.0 * torch.sum(torch.log(torch.diagonal(chol, dim1=-2, dim2=-1) + 1e-8), dim=-1)
//  [python]
def log_det_by_cholesky(matrix):
    """
	Args:
		matrix: matrix must be a positive define matrix.
				shape [N, C, D, D].
	Ref:
		https://github.com/tensorflow/tensorflow/blob/r1.13/tensorflow/python/ops/linalg/linalg_impl.py
	"""
    # This uses the property that the log det(A) = 2 * sum(log(real(diag(C))))
    # where C is the cholesky decomposition of A.
    chol = torch.linalg.cholesky(matrix)
    # return 2.0 * torch.sum(torch.log(torch.diagonal(chol, dim1=-2, dim2=-1) + 1e-6), dim=-1)
    return 2.0 * torch.sum(torch.log(torch.diagonal(chol, dim1=-2, dim2=-1) + 1e-8), dim=-1)
//  [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.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 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, 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 solve_linear_system(XtX, XtY, sub_batch_size=None):
    """
    Solve linear system of equations. use sub batches to avoid MAGMA bug
    :param XtX: matrix of the coefficients
    :param XtY: vector of the
    :param sub_batch_size: size of mini mini batch to avoid MAGMA error, if None - uses the entire batch
    :return:
    """
    if sub_batch_size is None:
        sub_batch_size = XtX.size(0)
    n_iterations = int(XtX.size(0) / sub_batch_size)
    assert sub_batch_size%sub_batch_size == 0, "batch size should be a factor of {}".format(sub_batch_size)
    beta = torch.zeros_like(XtY)
    n_elements = XtX.shape[2]
    for i in range(n_iterations):
        try:
            L = torch.cholesky(XtX[sub_batch_size * i:sub_batch_size * (i + 1), ...], upper=False)
            beta[sub_batch_size * i:sub_batch_size * (i + 1), ...] = \
                torch.cholesky_solve(XtY[sub_batch_size * i:sub_batch_size * (i + 1), ...], L, upper=False)
        except:
            # # add noise to diagonal for cases where XtX is low rank
            eps = torch.normal(torch.zeros(sub_batch_size, n_elements, device=XtX.device),
                               0.01 * torch.ones(sub_batch_size, n_elements, device=XtX.device))
            eps = torch.diag_embed(torch.abs(eps))
            XtX[sub_batch_size * i:sub_batch_size * (i + 1), ...] = \
                XtX[sub_batch_size * i:sub_batch_size * (i + 1), ...] + \
                eps * XtX[sub_batch_size * i:sub_batch_size * (i + 1), ...]
            try:
                L = torch.cholesky(XtX[sub_batch_size * i:sub_batch_size * (i + 1), ...], upper=False)
                beta[sub_batch_size * i:sub_batch_size * (i + 1), ...] = \
                    torch.cholesky_solve(XtY[sub_batch_size * i:sub_batch_size * (i + 1), ...], L, upper=False)
            except:
                beta[sub_batch_size * i:sub_batch_size * (i + 1), ...], _ =\
                    torch.solve(XtY[sub_batch_size * i:sub_batch_size * (i + 1), ...], XtX[sub_batch_size * i:sub_batch_size * (i + 1), ...])
    return beta
//  [python]
def solve_linear_system(XtX, XtY, sub_batch_size=None):
    """
    Solve linear system of equations. use sub batches to avoid MAGMA bug
    :param XtX: matrix of the coefficients
    :param XtY: vector of the
    :param sub_batch_size: size of mini mini batch to avoid MAGMA error, if None - uses the entire batch
    :return:
    """
    if sub_batch_size is None:
        sub_batch_size = XtX.size(0)
    n_iterations = int(XtX.size(0) / sub_batch_size)
    assert sub_batch_size%sub_batch_size == 0, "batch size should be a factor of {}".format(sub_batch_size)
    beta = torch.zeros_like(XtY)
    n_elements = XtX.shape[2]
    for i in range(n_iterations):
        try:
            L = torch.linalg.cholesky(XtX[sub_batch_size * i:sub_batch_size * (i + 1), ...], upper=False)
            beta[sub_batch_size * i:sub_batch_size * (i + 1), ...] = \
                torch.linalg.cholesky_solve(XtY[sub_batch_size * i:sub_batch_size * (i + 1), ...], L, upper=False)
        except:
            # # add noise to diagonal for cases where XtX is low rank
            eps = torch.normal(torch.zeros(sub_batch_size, n_elements, device=XtX.device),
                               0.01 * torch.ones(sub_batch_size, n_elements, device=XtX.device))
            eps = torch.diag_embed(torch.abs(eps))
            XtX[sub_batch_size * i:sub_batch_size * (i + 1), ...] = \
                XtX[sub_batch_size * i:sub_batch_size * (i + 1), ...] + \
                eps * XtX[sub_batch_size * i:sub_batch_size * (i + 1), ...]
            try:
                L = torch.linalg.cholesky(XtX[sub_batch_size * i:sub_batch_size * (i + 1), ...], upper=False)
                beta[sub_batch_size * i:sub_batch_size * (i + 1), ...] = \
                    torch.linalg.cholesky_solve(XtY[sub_batch_size * i:sub_batch_size * (i + 1), ...], L, upper=False)
            except:
                beta[sub_batch_size * i:sub_batch_size * (i + 1), ...], _ =\
                    torch.solve(XtY[sub_batch_size * i:sub_batch_size * (i + 1), ...], XtX[sub_batch_size * i:sub_batch_size * (i + 1), ...])
    return beta
//  [python]
def test_multivariate_normal_log_prob(self):
        mean = torch.randn(3, requires_grad=True)
        tmp = torch.randn(3, 10)
        cov = (torch.matmul(tmp, tmp.t()) / tmp.size(-1)).requires_grad_()
        prec = cov.inverse().requires_grad_()
        scale_tril = torch.cholesky(cov, upper=False).requires_grad_()

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

        x = dist1.sample((10,))
        expected = ref_dist.logpdf(x.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
        mean = torch.randn(5, 3, requires_grad=True)
        tmp = torch.randn(5, 3, 10)
        cov = (tmp.unsqueeze(-2) * tmp.unsqueeze(-3)).mean(-1).requires_grad_()

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

        x = dist_batched.sample((10,))
        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(0.0, (batched_prob - unbatched_prob).abs().max(), atol=1e-3, rtol=0)
//  [python]
def test_multivariate_normal_log_prob(self):
        mean = torch.randn(3, requires_grad=True)
        tmp = torch.randn(3, 10)
        cov = (torch.matmul(tmp, tmp.t()) / tmp.size(-1)).requires_grad_()
        prec = cov.inverse().requires_grad_()
        scale_tril = torch.linalg.cholesky(cov, upper=False).requires_grad_()

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

        x = dist1.sample((10,))
        expected = ref_dist.logpdf(x.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
        mean = torch.randn(5, 3, requires_grad=True)
        tmp = torch.randn(5, 3, 10)
        cov = (tmp.unsqueeze(-2) * tmp.unsqueeze(-3)).mean(-1).requires_grad_()

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

        x = dist_batched.sample((10,))
        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(0.0, (batched_prob - unbatched_prob).abs().max(), atol=1e-3, rtol=0)
//  [python]
def plot_so3_helper(T_est, P_est, color, i1, i2, i_max, path):
    r = torch.randn(i_max, 3)
    P_est_chol = torch.cholesky(P_est[6:9, 6:9]+torch.eye(3)*1e-16)
    p_est = bmv(P_est_chol.expand(i_max, 3, 3), r) + T_est[-1, :3, 4]
    plt.scatter(p_est[:, i1], p_est[:, i2], color=color, s=2, alpha=0.5)
//  [python]
def plot_so3_helper(T_est, P_est, color, i1, i2, i_max, path):
    r = torch.randn(i_max, 3)
    P_est_chol = torch.linalg.cholesky(P_est[6:9, 6:9]+torch.eye(3)*1e-16)
    p_est = bmv(P_est_chol.expand(i_max, 3, 3), r) + T_est[-1, :3, 4]
    plt.scatter(p_est[:, i1], p_est[:, i2], color=color, s=2, alpha=0.5)
//  [python]
def plot_se23_helper(T_est, P_est, color, i1, i2, i_max, path):
    P_est_chol = torch.cholesky(P_est + torch.eye(9)*1e-16)
    r = bmv(P_est_chol.expand(i_max, 9, 9), torch.randn(i_max, 9))
    Ttemp = T_est[-1].expand(i_max, 5, 5).bmm(SE3_2.exp(r.cuda()).cpu())
    p_est = Ttemp[:, :3, 4]
    plt.scatter(p_est[:, i1], p_est[:, i2], color=color, s=3)
//  [python]
def plot_se23_helper(T_est, P_est, color, i1, i2, i_max, path):
    P_est_chol = torch.linalg.cholesky(P_est + torch.eye(9)*1e-16)
    r = bmv(P_est_chol.expand(i_max, 9, 9), torch.randn(i_max, 9))
    Ttemp = T_est[-1].expand(i_max, 5, 5).bmm(SE3_2.exp(r.cuda()).cpu())
    p_est = Ttemp[:, :3, 4]
    plt.scatter(p_est[:, i1], p_est[:, i2], color=color, s=3)
//  [python]
def _cholesky_inv(X):
    u = torch.cholesky(X)
    return torch.cholesky_inverse(u)
//  [python]
def _cholesky_inv(X):
    u = torch.linalg.cholesky(X)
    return torch.cholesky_inverse(u)
//  [python]
def _callback(self, Xte, yte, epoch):
        with torch.no_grad():
            K_u_u = self.kernel.matrix(self.X_concat, torch.exp(self.log_ls_h))
            K_u1 = self.kernel.cross(self.X_col, self.X_concat, torch.exp(self.log_ls_h))
            K_u2 = self.kernel.cross(self.X_col, self.X_col, torch.exp(self.log_ls_h))
            K_uxx_uxx = self.kernel.matrix_pp_pp(K_u2, self.X_col, 0, torch.exp(self.log_ls_h))
            K_uxx = self.kernel.cross_pp_0(K_u1, self.X_col, self.X_concat, 0, torch.exp(self.log_ls_h))
            K_ux_ux = self.kernel.matrix_p_p(K_u2, self.X_col, 1, torch.exp(self.log_ls_h))
            K_ux = self.kernel.cross_p_0(K_u1, self.X_col, self.X_concat, 1, torch.exp(self.log_ls_h))
            K_ux_uxx = self.kernel.cross_p1_pp0(K_u2, self.X_col, self.X_col, torch.exp(self.log_ls_h))
            K = torch.cat(
                (
                    torch.cat((K_u_u, K_uxx.T, K_ux.T), 1),
                    torch.cat((K_uxx, K_uxx_uxx, K_ux_uxx.T), 1),
                    torch.cat((K_ux, K_ux_uxx, K_ux_ux), 1),
                ),
                0,
            )
            A = torch.cholesky(K)
            eta = self.mu
            f = torch.matmul(A, eta)
            K_u1_te = self.kernel.cross(Xte, self.X_concat, torch.exp(self.log_ls_h))
            K_u2_te = self.kernel.cross(Xte, self.X_col, torch.exp(self.log_ls_h))
            K_utt_te = self.kernel.cross_pp_0(K_u2_te.T, self.X_col, Xte, 0, torch.exp(self.log_ls_h))
            K_ut_te = self.kernel.cross_p_0(K_u2_te.T, self.X_col, Xte, 1, torch.exp(self.log_ls_h))
            K_te = torch.cat((K_u1_te, K_utt_te.T, K_ut_te.T), 1)
            pred_mean = torch.matmul(K_te, torch.linalg.solve(K, f[0:self.N_concat + 2 * self.N_col]))
            err_te = torch.sqrt(torch.mean(torch.square(pred_mean - yte)))
            if self.best_RMSE is None or self.best_RMSE > err_te:
                self.best_RMSE = err_te
                cov = torch.diag(self.L1.exp().view(-1)) + torch.matmul(self.L2, self.L2.T)
                self.L = torch.cholesky(cov)
                Ltril = torch.tril(self.L)
                A1 = torch.cholesky(K)
                t1 = torch.linalg.solve(A1.T, Ltril[0:self.N_concat + 2 * self.N_col])
                t2 = torch.matmul(t1, t1.T)
                pred_stds = (1 + self.jitter - (K_te * torch.linalg.solve(K, K_te.T).T).sum(1) + (K_te * torch.matmul(t2, K_te.T).T).sum(1))
                pred_stds = pred_stds.view(pred_mean.shape)
                print("Epoch", self.epoch, "RMSE: ", err_te)
                pred_mean = (pred_mean.cpu().reshape((67, 128)))
                pred_stds = pred_stds.cpu().reshape((67, 128))
                plt.imshow(pred_mean, cmap="hot")
                ti = "RMSE " + str(err_te.cpu().numpy())
                plt.title(ti)
                plt.savefig("allen_cahn")
                plt.clf()
                plt.imshow(pred_stds, cmap="hot")
                ti = "RMSE " + str(err_te.cpu().numpy())
                plt.title(ti)
                plt.savefig("allen_cahn_var")
                plt.clf()
//  [python]
def _callback(self, Xte, yte, epoch):
        with torch.no_grad():
            K_u_u = self.kernel.matrix(self.X_concat, torch.exp(self.log_ls_h))
            K_u1 = self.kernel.cross(self.X_col, self.X_concat, torch.exp(self.log_ls_h))
            K_u2 = self.kernel.cross(self.X_col, self.X_col, torch.exp(self.log_ls_h))
            K_uxx_uxx = self.kernel.matrix_pp_pp(K_u2, self.X_col, 0, torch.exp(self.log_ls_h))
            K_uxx = self.kernel.cross_pp_0(K_u1, self.X_col, self.X_concat, 0, torch.exp(self.log_ls_h))
            K_ux_ux = self.kernel.matrix_p_p(K_u2, self.X_col, 1, torch.exp(self.log_ls_h))
            K_ux = self.kernel.cross_p_0(K_u1, self.X_col, self.X_concat, 1, torch.exp(self.log_ls_h))
            K_ux_uxx = self.kernel.cross_p1_pp0(K_u2, self.X_col, self.X_col, torch.exp(self.log_ls_h))
            K = torch.cat(
                (
                    torch.cat((K_u_u, K_uxx.T, K_ux.T), 1),
                    torch.cat((K_uxx, K_uxx_uxx, K_ux_uxx.T), 1),
                    torch.cat((K_ux, K_ux_uxx, K_ux_ux), 1),
                ),
                0,
            )
            A = torch.linalg.cholesky(K)
            eta = self.mu
            f = torch.matmul(A, eta)
            K_u1_te = self.kernel.cross(Xte, self.X_concat, torch.exp(self.log_ls_h))
            K_u2_te = self.kernel.cross(Xte, self.X_col, torch.exp(self.log_ls_h))
            K_utt_te = self.kernel.cross_pp_0(K_u2_te.T, self.X_col, Xte, 0, torch.exp(self.log_ls_h))
            K_ut_te = self.kernel.cross_p_0(K_u2_te.T, self.X_col, Xte, 1, torch.exp(self.log_ls_h))
            K_te = torch.cat((K_u1_te, K_utt_te.T, K_ut_te.T), 1)
            pred_mean = torch.matmul(K_te, torch.linalg.solve(K, f[0:self.N_concat + 2 * self.N_col]))
            err_te = torch.sqrt(torch.mean(torch.square(pred_mean - yte)))
            if self.best_RMSE is None or self.best_RMSE > err_te:
                self.best_RMSE = err_te
                cov = torch.diag(self.L1.exp().view(-1)) + torch.matmul(self.L2, self.L2.T)
                self.L = torch.linalg.cholesky(cov)
                Ltril = torch.tril(self.L)
                A1 = torch.linalg.cholesky(K)
                t1 = torch.linalg.solve(A1.T, Ltril[0:self.N_concat + 2 * self.N_col])
                t2 = torch.matmul(t1, t1.T)
                pred_stds = (1 + self.jitter - (K_te * torch.linalg.solve(K, K_te.T).T).sum(1) + (K_te * torch.matmul(t2, K_te.T).T).sum(1))
                pred_stds = pred_stds.view(pred_mean.shape)
                print("Epoch", self.epoch, "RMSE: ", err_te)
                pred_mean = (pred_mean.cpu().reshape((67, 128)))
                pred_stds = pred_stds.cpu().reshape((67, 128))
                plt.imshow(pred_mean, cmap="hot")
                ti = "RMSE " + str(err_te.cpu().numpy())
                plt.title(ti)
                plt.savefig("allen_cahn")
                plt.clf()
                plt.imshow(pred_stds, cmap="hot")
                ti = "RMSE " + str(err_te.cpu().numpy())
                plt.title(ti)
                plt.savefig("allen_cahn_var")
                plt.clf()
//  [python]
def test_old_cholesky_batched_upper(self, device, dtype):
        from torch.testing._internal.common_utils import random_hermitian_pd_matrix

        batchsize = 2
        A = random_hermitian_pd_matrix(3, batchsize, dtype=dtype, device=device)
        A_triu = A.triu()  # fill the lower triangular part with zero

        U = torch.cholesky(A_triu, upper=True)

        reconstruct_A = U.conj().transpose(-2, -1) @ U
        self.assertEqual(A, reconstruct_A)
//  [python]
def test_old_cholesky_batched_upper(self, device, dtype):
        from torch.testing._internal.common_utils import random_hermitian_pd_matrix

        batchsize = 2
        A = random_hermitian_pd_matrix(3, batchsize, dtype=dtype, device=device)
        A_triu = A.triu()  # fill the lower triangular part with zero

        U = torch.linalg.cholesky(A_triu, upper=True)

        reconstruct_A = U.conj().transpose(-2, -1) @ U
        self.assertEqual(A, reconstruct_A)
//  [python]
def update(self, g, prec, iteration, *, ForceUpdate = False):

        if not ForceUpdate:
            raise RuntimeError

        g = g.to(dtype=torch.double)
        prec = prec.to(dtype=torch.double)

        # dirty solution
        self._Gamma = self._querry.Gamma.t()
        self._GammaTransposed = self.querry.GammaTransposed.t()
        self._alpha = self.querry.alpha

        cov = 1/prec
        Lambda = torch.einsum('im, m, sm -> is', [self._GammaTransposed, cov, self._GammaTransposed])   # checked this; seems to be okay
        Lambda += torch.diag(self.vo_variances)
        L = torch.cholesky(Lambda)
        LambdaInv = torch.cholesky_inverse(L)

        solvec = LambdaInv @ (self._GammaTransposed @ g - self._alpha)
        mean = g - torch.einsum('i, mi, m -> i', [cov, self._GammaTransposed, solvec])

        A = self._GammaTransposed * cov
        postcov_diag_subtractor = torch.einsum('si, sm, mi -> i', [A, LambdaInv, A])

        self._mean = mean
        self._vars = cov - postcov_diag_subtractor
//  [python]
def update(self, g, prec, iteration, *, ForceUpdate = False):

        if not ForceUpdate:
            raise RuntimeError

        g = g.to(dtype=torch.double)
        prec = prec.to(dtype=torch.double)

        # dirty solution
        self._Gamma = self._querry.Gamma.t()
        self._GammaTransposed = self.querry.GammaTransposed.t()
        self._alpha = self.querry.alpha

        cov = 1/prec
        Lambda = torch.einsum('im, m, sm -> is', [self._GammaTransposed, cov, self._GammaTransposed])   # checked this; seems to be okay
        Lambda += torch.diag(self.vo_variances)
        L = torch.linalg.cholesky(Lambda)
        LambdaInv = torch.linalg.solve_triangular(L, torch.eye(L.shape[0]), upper=False)

        solvec = LambdaInv @ (self._GammaTransposed @ g - self._alpha)
        mean = g - torch.einsum('i, mi, m -> i', [cov, self._GammaTransposed, solvec])

        A = self._GammaTransposed * cov
        postcov_diag_subtractor = torch.einsum('si, sm, mi -> i', [A, LambdaInv, A])

        self._mean = mean
        self._vars = cov - postcov_diag_subtractor
//  [python]
def reparameterize(self, Lambda, Psi_diag, z):
        """Performs the reparameterization trick for a Gaussian random variable.
        For details, see:

            http://blog.shakirm.com/2015/10/
            machine-learning-trick-of-the-day-4-reparameterisation-tricks/

        :param Lambda:   Current value for Lambda parameter.
        :param Psi_diag: Current value for Psi parameter.
        :param z:        Latent variable.
        :return:         Samples of y from estimated parameters Lambda and Psi.
        """
        n = z.shape[1]
        p = Psi_diag.shape[0]
        eps = torch.randn(p, n, device=cuda.device())
        # For numerical stability. For Psi to be PD, all elements must be
        # positive: https://math.stackexchange.com/a/927916/159872.
        Psi_diag = LA.to_positive(Psi_diag)
        R = torch.cholesky(diag(Psi_diag), upper=False)
        return Lambda @ z + R @ eps
//  [python]
def reparameterize(self, Lambda, Psi_diag, z):
        """Performs the reparameterization trick for a Gaussian random variable.
        For details, see:

            http://blog.shakirm.com/2015/10/
            machine-learning-trick-of-the-day-4-reparameterisation-tricks/

        :param Lambda:   Current value for Lambda parameter.
        :param Psi_diag: Current value for Psi parameter.
        :param z:        Latent variable.
        :return:         Samples of y from estimated parameters Lambda and Psi.
        """
        n = z.shape[1]
        p = Psi_diag.shape[0]
        eps = torch.randn(p, n, device=cuda.device())
        # For numerical stability. For Psi to be PD, all elements must be
        # positive: https://math.stackexchange.com/a/927916/159872.
        Psi_diag = LA.to_positive(Psi_diag)
        R = torch.linalg.cholesky(diag(Psi_diag), upper=False)
        return Lambda @ z + R @ eps
//  [python]
def set_cholesky(self):
        self.L = torch.cholesky(self.K, upper=True)
//  [python]
def set_cholesky(self):
        self.L = torch.linalg.cholesky(self.K, upper=True)