jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def estimate_variance_tensor(self, X, criterion, dist, var0, lik_prec, mb_size=128):
        tau = 1/var0
        opt = KFAC(self.net)
        U_halfs = {}
        V_halfs = {}
        
        lik_std = sqrt(1/lik_prec)

        # Populate parameters with the means
        self.sample(scale=0, require_grad=True)

        # for x, y in tqdm(train_loader):
        for i in range(0, X.shape[0], mb_size):
            x = X[i:i+mb_size]

            output = self(x)
            distribution = dist(output, lik_std)
            y = distribution.sample()
            
            loss = criterion(output, y)
            loss.backward()
            opt.step()

        with torch.no_grad():
            for group in opt.param_groups:
                if len(group['params']) == 2:
                    weight, bias = group['params']
                else:
                    weight = group['params'][0]
                    bias = None

                module = group['mod']
                state = opt.state[module]

                U = state['ggt']
                V = state['xxt']

                m, n = int(U.shape[0]), int(V.shape[0])

                # Add priors
                n_data = X.shape[0]
                U = sqrt(n_data)*U + sqrt(tau)*torch.eye(m, device=self.device)
                V = sqrt(n_data)*V + sqrt(tau)*torch.eye(n, device=self.device)

#                 try:
                    # Use cholesky_inverse for numerical stability since U, V SPD
                u = torch.linalg.cholesky(U)
                v = torch.linalg.cholesky(V)
#                 except Exception as e:
#                     u = torch.cholesky(U + 1*torch.eye(m, device=self.device))
#                     v = torch.cholesky(V)

                U_half = torch.linalg.cholesky(torch.linalg.cholesky_ex(u).inverse(), upper=False)
                V_half = torch.linalg.cholesky(torch.linalg.cholesky_ex(v).inverse(), upper=True)

                U_halfs[(module, 'weight')] = U_half
                V_halfs[(module, 'weight')] = V_half


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

            if name == 'bias':
                continue

            U_half = U_halfs[(module, name)]
            V_half = V_halfs[(module, name)]

            module.__getattr__(f'{name}_U_half').copy_(U_half)
            module.__getattr__(f'{name}_V_half').copy_(V_half)
//  [python]
def test_cholesky(self):
        def func(root, upper):
            x = 0.5 * (root + root.transpose(-1, -2).conj())
            return torch.cholesky(x, upper)

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

            gradcheck(func, [root, upper])
            # TODO: gradgradcheck does not work correctly yet for complex
            if not dtype.is_complex:
                gradgradcheck(func, [root, upper])

            root = torch.rand(*dims, dtype=dtype)
            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, dtype in product([True, False],
                                          [(3, 3), (4, 3, 2, 2)],
                                          [torch.double, torch.cdouble]):
            run_test(upper, dims, dtype)
//  [python]
def test_cholesky(self):
        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, dtype):
            root = torch.rand(*dims, dtype=dtype, requires_grad=True)
            root = root + torch.eye(dims[-1])

            gradcheck(func, [root, upper])
            # TODO: gradgradcheck does not work correctly yet for complex
            if not dtype.is_complex:
                gradgradcheck(func, [root, upper])

            root = torch.rand(*dims, dtype=dtype)
            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, dtype in product([True, False],
                                          [(3, 3), (4, 3, 2, 2)],
                                          [torch.double, torch.cdouble]):
            run_test(upper, dims, dtype)
//  [python]
def __init__(self,
                 kernel: Kernel,
                 n: int,
                 m: int,
                 n_samples: int,
                 z: InducingPoints,
                 likelihood: Likelihood,
                 whiten: Optional[bool] = True,
                 tied_samples: Optional[bool] = True):
        """
        __init__ method for Sparse GP Class
        Parameters
        ----------
        kernel : Kernel
            kernel used for sparse GP (e.g., QuadExp)
        n : int
            number of neurons
        m : int
            number of conditions
        n_samples : int
            number of samples 
        z : InducingPoints
            inducing points for sparse GP
        likelihood : Likelihood
            likleihood p(y | f) 
        whiten : Optional bool
            whiten q if true
        tied_samples : Optional bool

        Returns
        -------
        
        """
        # initalize q_sqrt^2 at the prior kzz
        n_inducing = z.n_z
        _z = self._expand_z(z.prms)
        e = torch.eye(n_inducing)
        kzz = kernel(_z, _z) + (e * jitter)
        l = torch.cholesky(kzz, upper=False)[None, ...]

        super().__init__(kernel,
                         n,
                         m,
                         n_samples,
                         n_inducing,
                         likelihood,
                         whiten=whiten,
                         tied_samples=tied_samples)
        self.z = z
//  [python]
def __init__(self,
                 kernel: Kernel,
                 n: int,
                 m: int,
                 n_samples: int,
                 z: InducingPoints,
                 likelihood: Likelihood,
                 whiten: Optional[bool] = True,
                 tied_samples: Optional[bool] = True):
        """
        __init__ method for Sparse GP Class
        Parameters
        ----------
        kernel : Kernel
            kernel used for sparse GP (e.g., QuadExp)
        n : int
            number of neurons
        m : int
            number of conditions
        n_samples : int
            number of samples 
        z : InducingPoints
            inducing points for sparse GP
        likelihood : Likelihood
            likleihood p(y | f) 
        whiten : Optional bool
            whiten q if true
        tied_samples : Optional bool

        Returns
        -------
        
        """
        # initalize q_sqrt^2 at the prior kzz
        n_inducing = z.n_z
        _z = self._expand_z(z.prms)
        e = torch.eye(n_inducing)
        kzz = kernel(_z, _z) + (e * jitter)
        l = torch.linalg.cholesky(kzz, upper=False)[None, ...]

        super().__init__(kernel,
                         n,
                         m,
                         n_samples,
                         n_inducing,
                         likelihood,
                         whiten=whiten,
                         tied_samples=tied_samples)
        self.z = z
//  [python]
def estimate_variance_batch(self, X, criterion, dist, var0):
        tau = 1/var0
        opt = KFAC(self.net)
        U_halfs = {}
        V_halfs = {}

        # Populate parameters with the means
        self.sample(scale=0, require_grad=True)

        output = self(X)
#         distribution = dist(logits=output)
        distribution = dist(output, 1)
        y = distribution.sample()
        loss = criterion(output, y)
        loss.backward()
        opt.step()

        with torch.no_grad():
            for group in opt.param_groups:
                if len(group['params']) == 2:
                    weight, bias = group['params']
                else:
                    weight = group['params'][0]
                    bias = None

                module = group['mod']
                state = opt.state[module]

                U = state['ggt']
                V = state['xxt']

                m, n = int(U.shape[0]), int(V.shape[0])

                # Add priors
                n_data = len(X)
                U = sqrt(n_data)*U + sqrt(tau)*torch.eye(m, device=self.device)
                V = sqrt(n_data)*V + sqrt(tau)*torch.eye(n, device=self.device)
#                 U = U + sqrt(tau)*torch.eye(m, device=self.device)
#                 V = V + sqrt(tau)*torch.eye(n, device=self.device)
                
                # Use cholesky_inverse for numerical stability since U, V SPD
                u = torch.cholesky(U)
                v = torch.cholesky(V)

                U_half = torch.cholesky(torch.cholesky_inverse(u), upper=False)
                V_half = torch.cholesky(torch.cholesky_inverse(v), upper=True)

                U_halfs[(module, 'weight')] = U_half
                V_halfs[(module, 'weight')] = V_half


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

            if name == 'bias':
                continue

            U_half = U_halfs[(module, name)]
            V_half = V_halfs[(module, name)]

            module.__getattr__(f'{name}_U_half').copy_(U_half)
            module.__getattr__(f'{name}_V_half').copy_(V_half)
//  [python]
def estimate_variance_batch(self, X, criterion, dist, var0):
        tau = 1/var0
        opt = KFAC(self.net)
        U_halfs = {}
        V_halfs = {}

        # Populate parameters with the means
        self.sample(scale=0, require_grad=True)

        output = self(X)
#         distribution = dist(logits=output)
        distribution = dist(output, 1)
        y = distribution.sample()
        loss = criterion(output, y)
        loss.backward()
        opt.step()

        with torch.no_grad():
            for group in opt.param_groups:
                if len(group['params']) == 2:
                    weight, bias = group['params']
                else:
                    weight = group['params'][0]
                    bias = None

                module = group['mod']
                state = opt.state[module]

                U = state['ggt']
                V = state['xxt']

                m, n = int(U.shape[0]), int(V.shape[0])

                # Add priors
                n_data = len(X)
                U = sqrt(n_data)*U + sqrt(tau)*torch.eye(m, device=self.device)
                V = sqrt(n_data)*V + sqrt(tau)*torch.eye(n, device=self.device)
#                 U = U + sqrt(tau)*torch.eye(m, device=self.device)
#                 V = V + sqrt(tau)*torch.eye(n, device=self.device)
                
                # Use cholesky_inverse for numerical stability since U, V SPD
                u = torch.linalg.cholesky(U)
                v = torch.linalg.cholesky(V)

                U_half = torch.linalg.cholesky(torch.linalg.cholesky_inverse(u), upper=False)
                V_half = torch.linalg.cholesky(torch.linalg.cholesky_inverse(v), upper=True)

                U_halfs[(module, 'weight')] = U_half
                V_halfs[(module, 'weight')] = V_half


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

            if name == 'bias':
                continue

            U_half = U_halfs[(module, name)]
            V_half = V_halfs[(module, name)]

            module.__getattr__(f'{name}_U_half').copy_(U_half)
            module.__getattr__(f'{name}_V_half').copy_(V_half)
//  [python]
def potrf(a, upper=True, out=None):
    r"""Computes the Cholesky decomposition of a symmetric positive-definite
    matrix :math:`A`.

    For more information regarding :func:`torch.potrf`, please check :func:`torch.cholesky`.

    .. warning::
        :func:`torch.potrf` is deprecated in favour of :func:`torch.cholesky` and will be removed
        in the next release. Please use :func:`torch.cholesky` instead and note that the :attr:`upper`
        argument in :func:`torch.cholesky` defaults to ``False``.
    """
    warnings.warn("torch.potrf is deprecated in favour of torch.cholesky and will be removed in the next "
                  "release. Please use torch.cholesky instead and note that the :attr:`upper` argument in"
                  " torch.cholesky defaults to ``False``.", stacklevel=2)
    return torch.cholesky(a, upper=upper, out=out)
//  [python]
def potrf(a, upper=True, out=None):
    r"""Computes the Cholesky decomposition of a symmetric positive-definite
    matrix :math:`A`.

    For more information regarding :func:`torch.potrf`, please check :func:`torch.cholesky`.

    .. warning::
        :func:`torch.potrf` is deprecated in favour of :func:`torch.cholesky` and will be removed
        in the next release. Please use :func:`torch.cholesky` instead and note that the :attr:`upper`
        argument in :func:`torch.cholesky` defaults to ``False``.
    """
    warnings.warn("torch.potrf is deprecated in favour of torch.cholesky and will be removed in the next "
                  "release. Please use torch.cholesky instead and note that the :attr:`upper` argument in"
                  " torch.cholesky defaults to ``False``.", stacklevel=2)
    return torch.linalg.cholesky(a, upper=upper, out=out)
//  [python]
def get_cov(self, ix=None):

        if ix is None:
            ix = torch.arange(0, self.N)

        #return torch.potrf(self.kernel(self.X[ix])
        #                   + torch.eye(ix.numel())
        #                        *transform_forward(self.variance),
        #                   upper=False)
        ix = ix.to(self.X.device)
        return torch.cholesky(self.kernel(self.X[ix])
                           + torch.eye(ix.numel()).to(self.X.device)
                           #     *transform_forward(self.variance).to(self.device),
                          *transform_forward(self.variance),
                           upper=False)
//  [python]
def get_cov(self, ix=None):

        if ix is None:
            ix = torch.arange(0, self.N)

        #return torch.potrf(self.kernel(self.X[ix])
        #                   + torch.eye(ix.numel())
        #                        *transform_forward(self.variance),
        #                   upper=False)
        ix = ix.to(self.X.device)
        return torch.linalg.cholesky(self.kernel(self.X[ix])
                           + torch.eye(ix.numel()).to(self.X.device)
                           #     *transform_forward(self.variance).to(self.device),
                          *transform_forward(self.variance),
                           upper=False)
//  [python]
def forward(ctx, x):
        n = x.shape[0]
        mask = torch.triu_indices(n, n, 1)

        xpy = 0.5 * (x[mask[0]] + x[mask[1]])
        xpy_l = torch.cholesky(xpy)
        lhs = 2 * xpy_l.diagonal(dim1=1, dim2=2).log().sum(1)

        x_l = torch.cholesky(x)
        logdet_l = 2 * x_l.diagonal(dim1=1, dim2=2).log().sum(1)
        rhs = -0.5 * (logdet_l[mask[0]] + logdet_l[mask[1]])

        ctx.save_for_backward(xpy_l, x_l, mask)
        return lhs + rhs
//  [python]
def forward(ctx, x):
        n = x.shape[0]
        mask = torch.triu_indices(n, n, 1)

        xpy = 0.5 * (x[mask[0]] + x[mask[1]])
        xpy_l = torch.linalg.cholesky(xpy)
        lhs = 2 * xpy_l.diagonal(dim1=1, dim2=2).log().sum(1)

        x_l = torch.linalg.cholesky(x)
        logdet_l = 2 * x_l.diagonal(dim1=1, dim2=2).log().sum(1)
        rhs = -0.5 * (logdet_l[mask[0]] + logdet_l[mask[1]])

        ctx.save_for_backward(xpy_l, x_l, mask)
        return lhs + rhs
//  [python]
def manifold_sq_pdists(X):
    # We implement it manually rather than calling spd.dist() to avoid computing
    # the cholesky decomposition and its inverse several times.
    n, d, _ = X.shape
    mask = torch.triu_indices(n, n, 1)
    eye = torch.eye(d, out=X.new(d, d))

    # First, compute the matrix inside log for all pairs of points.
    # TODO(ccruceru): Check if manually computing the square root is beneficial.
    L = torch.cholesky(X)
    L_inv = torch.triangular_solve(eye, L, upper=False).solution
    A = multiAXAt(L_inv[mask[0]], X[mask[1]])

    # Then, compute the sum of squared logarithms of eigenvalues.
    # NOTE: This is equivalent to computing the matrix logarithm and then taking
    # its squared Frobenius norm:
    #       A_log = multilog(A)
    #       pdists = A_log.pow(2).sum((1, 2))
    W = eig(A)
    W.data.clamp_(min=X.manifold.wmin, max=X.manifold.wmax)
    pdists = W.log().pow(2).sum(-1)

    return pdists
//  [python]
def manifold_sq_pdists(X):
    # We implement it manually rather than calling spd.dist() to avoid computing
    # the cholesky decomposition and its inverse several times.
    n, d, _ = X.shape
    mask = torch.triu_indices(n, n, 1)
    eye = torch.eye(d, out=X.new(d, d))

    # First, compute the matrix inside log for all pairs of points.
    # TODO(ccruceru): Check if manually computing the square root is beneficial.
    L = torch.linalg.cholesky(X)
    L_inv = torch.triangular_solve(eye, L, upper=False).solution
    A = multiAXAt(L_inv[mask[0]], X[mask[1]])

    # Then, compute the sum of squared logarithms of eigenvalues.
    # NOTE: This is equivalent to computing the matrix logarithm and then taking
    # its squared Frobenius norm:
    #       A_log = multilog(A)
    #       pdists = A_log.pow(2).sum((1, 2))
    W = eig(A)
    W.data.clamp_(min=X.manifold.wmin, max=X.manifold.wmax)
    pdists = W.log().pow(2).sum(-1)

    return pdists
//  [python]
def test_cholesky_batched(self, device, dtype):
        from torch.testing._internal.common_utils import \
            (random_symmetric_pd_matrix,
             random_fullrank_matrix_distinct_singular_value)

        def cholesky_test_helper(n, batch_dims, upper):
            # This is a workaround while there is no support for complex random_symmetric_pd_matrix
            if dtype.is_complex:
                real_dtype = torch.float32 if dtype is torch.complex64 else torch.float64
                A_real = random_fullrank_matrix_distinct_singular_value(n, *batch_dims, dtype=real_dtype, device=device)
                A_imag = random_fullrank_matrix_distinct_singular_value(n, *batch_dims, dtype=real_dtype, device=device)
                A = A_real + 1j * A_imag
                # There is no support for complex batched matmul yet
                matmul_list = []
                for mat in A.contiguous().view(-1, n, n):
                    matmul_list.append(mat @ mat.t().conj())
                A = torch.stack(matmul_list).view(*batch_dims, n, n)
            else:
                A = random_symmetric_pd_matrix(n, *batch_dims, dtype=dtype, device=device)
            cholesky_exp = torch.stack([m.cholesky(upper=upper) for m in A.reshape(-1, n, n)])
            cholesky_exp = cholesky_exp.reshape_as(A)
            self.assertEqual(cholesky_exp, torch.cholesky(A, upper=upper))

        for upper, batchsize in product([True, False], [(3,), (3, 4), (2, 3, 4)]):
            cholesky_test_helper(3, batchsize, upper)
//  [python]
def test_cholesky_batched(self, device, dtype):
        from torch.testing._internal.common_utils import \
            (random_symmetric_pd_matrix,
             random_fullrank_matrix_distinct_singular_value)

        def cholesky_test_helper(n, batch_dims, upper):
            # This is a workaround while there is no support for complex random_symmetric_pd_matrix
            if dtype.is_complex:
                real_dtype = torch.float32 if dtype is torch.complex64 else torch.float64
                A_real = random_fullrank_matrix_distinct_singular_value(n, *batch_dims, dtype=real_dtype, device=device)
                A_imag = random_fullrank_matrix_distinct_singular_value(n, *batch_dims, dtype=real_dtype, device=device)
                A = A_real + 1j * A_imag
                # There is no support for complex batched matmul yet
                matmul_list = []
                for mat in A.contiguous().view(-1, n, n):
                    matmul_list.append(mat @ mat.t().conj())
                A = torch.stack(matmul_list).view(*batch_dims, n, n)
            else:
                A = random_symmetric_pd_matrix(n, *batch_dims, dtype=dtype, device=device)
            cholesky_exp = torch.stack([torch.linalg.cholesky(m, upper=upper) for m in A.reshape(-1, n, n)])
            cholesky_exp = cholesky_exp.reshape_as(A)
            self.assertEqual(cholesky_exp, torch.linalg.cholesky(A, upper=upper))

        for upper, batchsize in product([True, False], [(3,), (3, 4), (2, 3, 4)]):
            cholesky_test_helper(3, batchsize, upper)
//  [python]
def predict_with_var(self, inputs, vars=None):
        x_test = inputs[0]
        l1, h1, l2, h2 = self.l1, self.h1, self.l2, self.h2
        _noise = self._get_noise_according_exp_format()

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

        kx = self.kernel_list[0](self.inputs_tr[0], x_test, l1, h1, l2, h2)
        L = torch.cholesky(sigma)
        LinvKx,_ = torch.triangular_solve(kx, L, upper = False)

        mean = kx.t() @ torch.cholesky_solve(self.outputs_tr[0], L)  # torch.linalg.cholesky()
        
        var_diag = self.kernel_list[0](x_test, x_test, l1, h1, l2, h2).diag().view(-1, 1) \
            - (LinvKx**2).sum(dim = 0).view(-1, 1)
        # add the noise uncertainty
        var_diag = var_diag + _noise.pow(-1)

        return mean, var_diag
//  [python]
def predict_with_var(self, inputs, vars=None):
        x_test = inputs[0]
        l1, h1, l2, h2 = self.l1, self.h1, self.l2, self.h2
        _noise = self._get_noise_according_exp_format()

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

        kx = self.kernel_list[0](self.inputs_tr[0], x_test, l1, h1, l2, h2)
        L = torch.linalg.cholesky(sigma)
        LinvKx,_ = torch.triangular_solve(kx, L, upper = False)

        mean = kx.t() @ torch.cholesky_solve(self.outputs_tr[0], L)  # torch.linalg.cholesky()
        
        var_diag = self.kernel_list[0](x_test, x_test, l1, h1, l2, h2).diag().view(-1, 1) \
            - (LinvKx**2).sum(dim = 0).view(-1, 1)
        # add the noise uncertainty
        var_diag = var_diag + _noise.pow(-1)

        return mean, var_diag
//  [python]
def test_cholesky(self, device, dtype):
        from torch.testing._internal.common_utils import \
            (random_symmetric_pd_matrix,
             random_fullrank_matrix_distinct_singular_value)

        # This is a workaround while there is no support for complex random_symmetric_pd_matrix
        if dtype.is_complex:
            real_dtype = torch.float32 if dtype is torch.complex64 else torch.float64
            A_real = random_fullrank_matrix_distinct_singular_value(10, dtype=real_dtype, device=device)
            A_imag = random_fullrank_matrix_distinct_singular_value(10, dtype=real_dtype, device=device)
            A = A_real + 1j * A_imag
            A = A @ A.t().conj()
        else:
            A = random_symmetric_pd_matrix(10, dtype=dtype, device=device)

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

        # test Upper Triangular
        U = torch.cholesky(A, 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.cholesky(A, 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 test_cholesky(self, device, dtype):
        from torch.testing._internal.common_utils import \
            (random_symmetric_pd_matrix,
             random_fullrank_matrix_distinct_singular_value)

        # This is a workaround while there is no support for complex random_symmetric_pd_matrix
        if dtype.is_complex:
            real_dtype = torch.float32 if dtype is torch.complex64 else torch.float64
            A_real = random_fullrank_matrix_distinct_singular_value(10, dtype=real_dtype, device=device)
            A_imag = random_fullrank_matrix_distinct_singular_value(10, dtype=real_dtype, device=device)
            A = A_real + 1j * A_imag
            A = A @ A.t().conj()
        else:
            A = random_symmetric_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 __init__(self, covariance_matrix=None):
		self.use_chol_factor = False
		if covariance_matrix is not None:
			self.covariance_matrix_col = torch.cholesky(covariance_matrix)
			self.use_chol_factor = True
		# Call parent constructor
		super(ESS, self).__init__()
//  [python]
def __init__(self, covariance_matrix=None):
		self.use_chol_factor = False
		if covariance_matrix is not None:
			self.covariance_matrix_col = torch.linalg.cholesky(covariance_matrix)
			self.use_chol_factor = True
		# Call parent constructor
		super(ESS, self).__init__()
//  [python]
def _precision_to_scale_tril(P):
    Lf = torch.cholesky(torch.flip(P, (-2, -1)))
    L_inv = torch.transpose(torch.flip(Lf, (-2, -1)), -2, -1)
    Id = torch.eye(P.shape[-1], dtype=P.dtype, device=P.device)
    L = torch.solve(Id, L_inv)
    return L
//  [python]
def _precision_to_scale_tril(P):
    # Ref: https://nbviewer.jupyter.org/gist/fehiepsi/5ef8e09e61604f10607380467eb82006#Precision-to-scale_tril
    Lf = torch.linalg.cholesky(torch.flip(P, (-2, -1)))
    L_inv = torch.transpose(torch.flip(Lf, (-2, -1)), -2, -1)
    Id = torch.eye(P.shape[-1], dtype=P.dtype, device=P.device)
    L = torch.linalg.solve_triangular(L_inv, Id, upper=False)
    return L
//  [python]
def __init__(self, mean, covariance):
        assert (
            mean.ndim == 3
            and covariance.D() == 2
            and mean.shape[1] == 1
            and mean.shape[2] == 1
        )
        self.mean = mean
        self.covariance = covariance
        self.precision = torch.inverse(covariance)
        self.cholesky_lower_covariance = torch.cholesky(covariance)
        self.d = mean.shape[0]
        self.log_det_covariance = torch.logdet(covariance)
//  [python]
def __init__(self, mean, covariance):
        assert (
            mean.ndim == 3
            and covariance.D() == 2
            and mean.shape[1] == 1
            and mean.shape[2] == 1
        )
        self.mean = mean
        self.covariance = covariance
        self.precision = torch.inverse(covariance)
        self.cholesky_lower_covariance = torch.linalg.cholesky(covariance)
        self.d = mean.shape[0]
        self.log_det_covariance = torch.logdet(covariance)
//  [python]
def covariance_cholesky(self, t: TensorType[..., 'seq_len', 1]) -> TensorType[..., 'seq_len', 'seq_len']:
        return torch.cholesky(self.covariance(t))
//  [python]
def covariance_cholesky(self, t: TensorType[..., 'seq_len', 1]) -> TensorType[..., 'seq_len', 'seq_len']:
        return torch.linalg.cholesky(self.covariance(t))
//  [python]
def forward(self, x):
        v1 = self.conv1(x)
        v2 = self.conv2(x)
        v3 = v1 + v2
        v4 = self.bn1(v3)
        v5 = self.bn2(v3)
        v6 = self.bn3(v3)
        p1 = v4.permute(0, 3, 1, 2)
        p2 = v6.permute(0, 3, 1, 2)
        p3 = torch.cat((p1, p2), dim=3)
        p4 = p3.flatten(1)
        l1 = p4.size()[0]
        p5 = p4.reshape(l1, 2, 16)
        l2 = p5.size()[0]
        p6 = p5.reshape(l2, 16, 2)
        c1 = torch.cholesky(p6)
        l3 = c1.size()[0]
        c2 = c1.reshape(l3, 2, 2, 4)
        (p1, p2, p3, g1, g2) = c2.contiguous().split([2, 2, 2, 4], dim=1)
        g1 = g1.squeeze(1)
        g2 = g2.squeeze(1)
        g3 = [g1, g2]
        g4 = torch.cat(g3, dim=-1)
        x1 = g4.unsqueeze(0)
        p1 = p1.permute(0, 1, 3, 2)
        p2 = p2.permute(0, 1, 3, 2)
        g5 = g4.permute(0, 2, 1)
        g6 = g5.unsqueeze(0).expand_as(x)
        s1 = torch.bmm(x, p2)
        s2 = torch.bmm(x1, g4)
        s3 = torch.bmm(s1, g5).squeeze(0).permute(2, 0, 1)
        s4 = torch.bmm(s2, g6).squeeze(0)
        s5 = s3 * s4
        g7 = s5.squeeze(1).reshape(1, 1, *x.size())
        bn2d = self.bn3(g7)
        soft = self.softmax(bn2d)
        return (g4, soft)
//  [python]
def forward(self, x):
        v1 = self.conv1(x)
        v2 = self.conv2(x)
        v3 = v1 + v2
        v4 = self.bn1(v3)
        v5 = self.bn2(v3)
        v6 = self.bn3(v3)
        p1 = v4.permute(0, 3, 1, 2)
        p2 = v6.permute(0, 3, 1, 2)
        p3 = torch.cat((p1, p2), dim=3)
        p4 = p3.flatten(1)
        l1 = p4.size()[0]
        p5 = p4.reshape(l1, 2, 16)
        l2 = p5.size()[0]
        p6 = p5.reshape(l2, 16, 2)
        c1 = torch.linalg.cholesky(p6)
        l3 = c1.size()[0]
        c2 = c1.reshape(l3, 2, 2, 4)
        (p1, p2, p3, g1, g2) = c2.contiguous().split([2, 2, 2, 4], dim=1)
        g1 = g1.squeeze(1)
        g2 = g2.squeeze(1)
        g3 = [g1, g2]
        g4 = torch.cat(g3, dim=-1)
        x1 = g4.unsqueeze(0)
        p1 = p1.permute(0, 1, 3, 2)
        p2 = p2.permute(0, 1, 3, 2)
        g5 = g4.permute(0, 2, 1)
        g6 = g5.unsqueeze(0).expand_as(x)
        s1 = torch.bmm(x, p2)
        s2 = torch.bmm(x1, g4)
        s3 = torch.bmm(s1, g5).squeeze(0).permute(2, 0, 1)
        s4 = torch.bmm(s2, g6).squeeze(0)
        s5 = s3 * s4
        g7 = s5.squeeze(1).reshape(1, 1, *x.size())
        bn2d = self.bn3(g7)
        soft = self.softmax(bn2d)
        return (g4, soft)
//  [python]
def covariance_cholesky(self, t: TensorType[..., 'N', 1]) -> TensorType[..., 'N', 'N']:
        return torch.cholesky(self.covariance(t))
//  [python]
def covariance_cholesky(self, t: TensorType[..., 'N', 1]) -> TensorType[..., 'N', 'N']:
        return torch.linalg.cholesky(self.covariance(t))
//  [python]
def __init__(
        self,
        dim: int = 1600,
        data_path: str | Path = DATA_DIR / "pines.csv",
        use_whitened: bool = False,
    ):
        # `log_norm_const` is taken from DDS paper, computed using long run SMC chain
        # (1000 temperatures, 30 seeds with 2000 samples each)
        super().__init__(dim=dim, log_norm_const=512.6)
        self.use_whitened = use_whitened
        self.path = Path(data_path)
        self.num_bins_per_dim = int(np.sqrt(dim))
        assert self.num_bins_per_dim**2 == self.dim

        self.signal_variance = 1.91
        self._poisson_a = 1.0 / self.dim
        self._beta = 1.0 / 33
        self.white_gaussian_log_normalizer = (
            -0.5 * self.dim * np.log(2.0 * np.pi)
        ).item()
        self.mu_zero = (np.log(126.0) - 0.5 * self.signal_variance).item()

        bin_counts = np.array(
            get_bin_counts(read_points(self.path), self.num_bins_per_dim)
        )
        bin_vals = torch.from_numpy(get_bin_vals(self.num_bins_per_dim)).float()
        gram_matrix = th_batch_kernel_fn(
            bin_vals, bin_vals, self.signal_variance, self.num_bins_per_dim, self._beta
        )

        self.register_buffer(
            "cholesky_gram", torch.cholesky(gram_matrix), persistent=False
        )
        self.register_buffer(
            "flat_bin_counts",
            torch.from_numpy(bin_counts.flatten()).float(),
            persistent=False,
        )

        half_log_det_gram = torch.sum(
            torch.log(torch.abs(torch.diag(self.cholesky_gram)))
        )
        self.unwhitened_gaussian_log_normalizer = (
            -0.5 * self.dim * np.log(2.0 * np.pi) - half_log_det_gram
        ).item()
//  [python]
def __init__(
        self,
        dim: int = 1600,
        data_path: str | Path = DATA_DIR / "pines.csv",
        use_whitened: bool = False,
    ):
        # `log_norm_const` is taken from DDS paper, computed using long run SMC chain
        # (1000 temperatures, 30 seeds with 2000 samples each)
        super().__init__(dim=dim, log_norm_const=512.6)
        self.use_whitened = use_whitened
        self.path = Path(data_path)
        self.num_bins_per_dim = int(np.sqrt(dim))
        assert self.num_bins_per_dim**2 == self.dim

        self.signal_variance = 1.91
        self._poisson_a = 1.0 / self.dim
        self._beta = 1.0 / 33
        self.white_gaussian_log_normalizer = (
            -0.5 * self.dim * np.log(2.0 * np.pi)
        ).item()
        self.mu_zero = (np.log(126.0) - 0.5 * self.signal_variance).item()

        bin_counts = np.array(
            get_bin_counts(read_points(self.path), self.num_bins_per_dim)
        )
        bin_vals = torch.from_numpy(get_bin_vals(self.num_bins_per_dim)).float()
        gram_matrix = th_batch_kernel_fn(
            bin_vals, bin_vals, self.signal_variance, self.num_bins_per_dim, self._beta
        )

        self.register_buffer(
            "cholesky_gram", torch.linalg.cholesky(gram_matrix), persistent=False
        )
        self.register_buffer(
            "flat_bin_counts",
            torch.from_numpy(bin_counts.flatten()).float(),
            persistent=False,
        )

        half_log_det_gram = torch.sum(
            torch.log(torch.abs(torch.diag(self.cholesky_gram)))
        )
        self.unwhitened_gaussian_log_normalizer = (
            -0.5 * self.dim * np.log(2.0 * np.pi) - half_log_det_gram
        ).item()
//  [python]
def gmm(cfg,all_fea, pi, mu, all_output):    
    Cov = []
    dist = []
    log_probs = []
    
    for i in range(len(mu)):
        temp = all_fea - mu[i]
        predi = all_output[:,i].unsqueeze(dim=-1)
        Covi = torch.matmul(temp.t(), temp * predi.expand_as(temp)) / (predi.sum()) + cfg.COWA.EPSILON * torch.eye(temp.shape[1]).cuda()
        try:
            chol = torch.cholesky(Covi)
        except RuntimeError:
            Covi += cfg.COWA.EPSILON * torch.eye(temp.shape[1]).cuda() * 100
            chol = torch.cholesky(Covi)
        chol_inv = torch.inverse(chol)
        Covi_inv = torch.matmul(chol_inv.t(), chol_inv)
        logdet = torch.logdet(Covi)
        mah_dist = (torch.matmul(temp, Covi_inv) * temp).sum(dim=1)
        log_prob = -0.5*(Covi.shape[0] * np.log(2*math.pi) + logdet + mah_dist) + torch.log(pi)[i]
        Cov.append(Covi)
        log_probs.append(log_prob)
        dist.append(mah_dist)
    Cov = torch.stack(Cov, dim=0)
    dist = torch.stack(dist, dim=0).t()
    log_probs = torch.stack(log_probs, dim=0).t()
    zz = log_probs - torch.logsumexp(log_probs, dim=1, keepdim=True).expand_as(log_probs)
    gamma = torch.exp(zz)
    
    return zz, gamma
//  [python]
def gmm(cfg,all_fea, pi, mu, all_output):    
    Cov = []
    dist = []
    log_probs = []
    
    for i in range(len(mu)):
        temp = all_fea - mu[i]
        predi = all_output[:,i].unsqueeze(dim=-1)
        Covi = torch.matmul(temp.t(), temp * predi.expand_as(temp)) / (predi.sum()) + cfg.COWA.EPSILON * torch.eye(temp.shape[1]).cuda()
        try:
            chol = torch.linalg.cholesky(Covi)
        except RuntimeError:
            Covi += cfg.COWA.EPSILON * torch.eye(temp.shape[1]).cuda() * 100
            chol = torch.linalg.cholesky(Covi)
        chol_inv = torch.inverse(chol)
        Covi_inv = torch.matmul(chol_inv.t(), chol_inv)
        logdet = torch.logdet(Covi)
        mah_dist = (torch.matmul(temp, Covi_inv) * temp).sum(dim=1)
        log_prob = -0.5*(Covi.shape[0] * np.log(2*math.pi) + logdet + mah_dist) + torch.log(pi)[i]
        Cov.append(Covi)
        log_probs.append(log_prob)
        dist.append(mah_dist)
    Cov = torch.stack(Cov, dim=0)
    dist = torch.stack(dist, dim=0).t()
    log_probs = torch.stack(log_probs, dim=0).t()
    zz = log_probs - torch.logsumexp(log_probs, dim=1, keepdim=True).expand_as(log_probs)
    gamma = torch.exp(zz)
    
    return zz, gamma
//  [python]
def _batch_capacitance_tril(W, D):
    r"""
    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.mT / D.unsqueeze(-2)
    K = torch.matmul(Wt_Dinv, W).contiguous()
    K.view(-1, m * m)[:, ::m + 1] += 1  # add identity matrix to K
    return torch.cholesky(K)
//  [python]
def _batch_capacitance_tril(W, D):
    r"""
    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.mT / D.unsqueeze(-2)
    K = torch.matmul(Wt_Dinv, W).contiguous()
    K.view(-1, m * m)[:, ::m + 1] += 1  # add identity matrix to K
    return torch.linalg.cholesky(K)
//  [python]
def event_logsumexp(self):
        """
        Integrates out all latent state (i.e. operating on event dimensions) of Gaussian component.
        """
        n = self.dim()
        chol_P = torch.cholesky(self.precision)
        chol_P_u = torch.triangular_solve(
            self.info_vec.unsqueeze(-1), chol_P, upper=False
        ).solution.squeeze(-1)
        u_P_u = chol_P_u.pow(2).sum(-1)
        # considering GammaGaussian as a Gaussian with precision = s * precision, info_vec = s * info_vec,
        # marginalize x variable, we get
        #   logsumexp(s) = alpha * log(s) - s * beta + 0.5 n * log(2 pi) + \
        #       0.5 s * uPu - 0.5 * log|P| - 0.5 n * log(s)
        # use the original parameterization of Gamma, we get
        #   logsumexp(s) = (concentration - 1) * log(s) - s * rate + 0.5 n * log(2 pi) - 0.5 * |P|
        # Note that `(concentration - 1) * log(s) - s * rate` is unnormalized log_prob of
        #   Gamma(concentration, rate)
        concentration = self.alpha - 0.5 * n + 1
        rate = self.beta - 0.5 * u_P_u
        log_normalizer_tmp = 0.5 * n * math.log(2 * math.pi) - chol_P.diagonal(
            dim1=-2, dim2=-1
        ).log().sum(-1)
        return Gamma(self.log_normalizer + log_normalizer_tmp, concentration, rate)
//  [python]
def event_logsumexp(self):
        """
        Integrates out all latent state (i.e. operating on event dimensions) of Gaussian component.
        """
        n = self.dim()
        chol_P = torch.linalg.cholesky(self.precision)
        chol_P_u = torch.linalg.solve_triangular(
            chol_P, self.info_vec.unsqueeze(-1), upper=False
        ).squeeze(-1)
        u_P_u = chol_P_u.pow(2).sum(-1)
        # considering GammaGaussian as a Gaussian with precision = s * precision, info_vec = s * info_vec,
        # marginalize x variable, we get
        #   logsumexp(s) = alpha * log(s) - s * beta + 0.5 n * log(2 pi) + \
        #       0.5 s * uPu - 0.5 * log|P| - 0.5 n * log(s)
        # use the original parameterization of Gamma, we get
        #   logsumexp(s) = (concentration - 1) * log(s) - s * rate + 0.5 n * log(2 pi) - 0.5 * |P|
        # Note that `(concentration - 1) * log(s) - s * rate` is unnormalized log_prob of
        #   Gamma(concentration, rate)
        concentration = self.alpha - 0.5 * n + 1
        rate = self.beta - 0.5 * u_P_u
        log_normalizer_tmp = 0.5 * n * math.log(2 * math.pi) - chol_P.diagonal(
            dim1=-2, dim2=-1
        ).log().sum(-1)
        return Gamma(self.log_normalizer + log_normalizer_tmp, concentration, rate)
//  [python]
def scale_tril(self):
        # The following identity is used to increase the numerically computation stability
        # for Cholesky decomposition (see http://www.gaussianprocess.org/gpml/, Section 3.4.3):
        #     W @ W.T + D = D1/2 @ (I + D-1/2 @ W @ W.T @ D-1/2) @ D1/2
        # The matrix "I + D-1/2 @ W @ W.T @ D-1/2" has eigenvalues bounded from below by 1,
        # hence it is well-conditioned and safe to take Cholesky decomposition.
        n = self._event_shape[0]
        cov_diag_sqrt_unsqueeze = self._unbroadcasted_cov_diag.sqrt().unsqueeze(-1)
        Dinvsqrt_W = self._unbroadcasted_cov_factor / cov_diag_sqrt_unsqueeze
        K = torch.matmul(Dinvsqrt_W, Dinvsqrt_W.mT).contiguous()
        K.view(-1, n * n)[:, ::n + 1] += 1  # add identity matrix to K
        scale_tril = cov_diag_sqrt_unsqueeze * torch.cholesky(K)
        return scale_tril.expand(self._batch_shape + self._event_shape + self._event_shape)
//  [python]
def scale_tril(self):
        # The following identity is used to increase the numerically computation stability
        # for Cholesky decomposition (see http://www.gaussianprocess.org/gpml/, Section 3.4.3):
        #     W @ W.T + D = D1/2 @ (I + D-1/2 @ W @ W.T @ D-1/2) @ D1/2
        # The matrix "I + D-1/2 @ W @ W.T @ D-1/2" has eigenvalues bounded from below by 1,
        # hence it is well-conditioned and safe to take Cholesky decomposition.
        n = self._event_shape[0]
        cov_diag_sqrt_unsqueeze = self._unbroadcasted_cov_diag.sqrt().unsqueeze(-1)
        Dinvsqrt_W = self._unbroadcasted_cov_factor / cov_diag_sqrt_unsqueeze
        K = torch.matmul(Dinvsqrt_W, Dinvsqrt_W.mT).contiguous()
        K.view(-1, n * n)[:, ::n + 1] += 1  # add identity matrix to K
        scale_tril = cov_diag_sqrt_unsqueeze * torch.linalg.cholesky(K)
        return scale_tril.expand(self._batch_shape + self._event_shape + self._event_shape)
//  [python]
def marginalize(self, left=0, right=0):
        """
        Marginalizing out variables on either side of the event dimension::

            g.marginalize(left=n).event_logsumexp() = g.event_logsumexp()
            g.marginalize(right=n).event_logsumexp() = g.event_logsumexp()

        and for data ``x``:

            g.condition(x).event_logsumexp().log_density(s)
              = g.marginalize(left=g.dim() - x.size(-1)).log_density(x, s)
        """
        # NB: the easiest way to think about this process is to consider GammaGaussian
        # as a Gaussian with precision and info_vec scaled by `s`.
        if left == 0 and right == 0:
            return self
        if left > 0 and right > 0:
            raise NotImplementedError
        n = self.dim()
        n_b = left + right
        a = slice(left, n - right)  # preserved
        b = slice(None, left) if left else slice(n - right, None)

        P_aa = self.precision[..., a, a]
        P_ba = self.precision[..., b, a]
        P_bb = self.precision[..., b, b]
        P_b = torch.cholesky(P_bb)
        P_a = torch.triangular_solve(P_ba, P_b, upper=False)[0]
        P_at = P_a.transpose(-1, -2)
        precision = P_aa - P_at.matmul(P_a)

        info_a = self.info_vec[..., a]
        info_b = self.info_vec[..., b]
        b_tmp = torch.triangular_solve(info_b.unsqueeze(-1), P_b, upper=False)[0]
        info_vec = info_a
        if n_b < n:
            info_vec = info_vec - P_at.matmul(b_tmp).squeeze(-1)

        alpha = self.alpha - 0.5 * n_b
        beta = self.beta - 0.5 * b_tmp.squeeze(-1).pow(2).sum(-1)
        log_normalizer = (
            self.log_normalizer
            + 0.5 * n_b * math.log(2 * math.pi)
            - P_b.diagonal(dim1=-2, dim2=-1).log().sum(-1)
        )
        return GammaGaussian(log_normalizer, info_vec, precision, alpha, beta)
//  [python]
def marginalize(self, left=0, right=0):
        """
        Marginalizing out variables on either side of the event dimension::

            g.marginalize(left=n).event_logsumexp() = g.event_logsumexp()
            g.marginalize(right=n).event_logsumexp() = g.event_logsumexp()

        and for data ``x``:

            g.condition(x).event_logsumexp().log_density(s)
              = g.marginalize(left=g.dim() - x.size(-1)).log_density(x, s)
        """
        # NB: the easiest way to think about this process is to consider GammaGaussian
        # as a Gaussian with precision and info_vec scaled by `s`.
        if left == 0 and right == 0:
            return self
        if left > 0 and right > 0:
            raise NotImplementedError
        n = self.dim()
        n_b = left + right
        a = slice(left, n - right)  # preserved
        b = slice(None, left) if left else slice(n - right, None)

        P_aa = self.precision[..., a, a]
        P_ba = self.precision[..., b, a]
        P_bb = self.precision[..., b, b]
        P_b = torch.linalg.cholesky(P_bb)
        P_a = torch.linalg.solve_triangular(P_b, P_ba, upper=False)
        P_at = P_a.transpose(-1, -2)
        precision = P_aa - P_at.matmul(P_a)

        info_a = self.info_vec[..., a]
        info_b = self.info_vec[..., b]
        b_tmp = torch.linalg.solve_triangular(P_b, info_b.unsqueeze(-1), upper=False)
        info_vec = info_a
        if n_b < n:
            info_vec = info_vec - P_at.matmul(b_tmp).squeeze(-1)

        alpha = self.alpha - 0.5 * n_b
        beta = self.beta - 0.5 * b_tmp.squeeze(-1).pow(2).sum(-1)
        log_normalizer = (
            self.log_normalizer
            + 0.5 * n_b * math.log(2 * math.pi)
            - P_b.diagonal(dim1=-2, dim2=-1).log().sum(-1)
        )
        return GammaGaussian(log_normalizer, info_vec, precision, alpha, beta)
//  [python]
def encode_SPD(A, eps=1e-6):
        '''Apply Cholesky decomposition to A and return the log-diagonal entires [n] and below [n*(n-1)/2].'''
        n = A.shape[0]
        A = torch.cholesky(A+eps*torch.eye(n))
        D = A.diag()
        D = D.log()
        ids = torch.tril_indices(n,n,-1)
        L = A[ids[0,:],ids[1,:]]
        return D, L
//  [python]
def encode_SPD(A, eps=1e-6):
        '''Apply Cholesky decomposition to A and return the log-diagonal entires [n] and below [n*(n-1)/2].'''
        n = A.shape[0]
        A = torch.linalg.cholesky(A+eps*torch.eye(n))
        D = A.diag()
        D = D.log()
        ids = torch.tril_indices(n,n,-1)
        L = A[ids[0,:],ids[1,:]]
        return D, L
//  [python]
def sample(self, n_samples):
        samples = torch.randn(self.n_params, n_samples)
        d = self.prior_precision_diag
        Vs = self.V * d.sqrt().reshape(-1, 1)
        VtV = Vs.T @ Vs
        Ik = torch.eye(len(VtV))
        A = torch.cholesky(VtV)
        B = torch.cholesky(VtV + Ik)
        A_inv = torch.inverse(A)
        C = torch.inverse(A_inv.T @ (B - Ik) @ A_inv)
        Kern_inv = torch.inverse(torch.inverse(C) + Vs.T @ Vs)
        dinv_sqrt = (d).sqrt().reshape(-1, 1)
        prior_sample = dinv_sqrt * samples
        gain_sample = dinv_sqrt * Vs @ Kern_inv @ (Vs.T @ samples)
        return self.mean + (prior_sample - gain_sample).T
//  [python]
def sample(self, n_samples):
        samples = torch.randn(self.n_params, n_samples)
        d = self.prior_precision_diag
        Vs = self.V * d.sqrt().reshape(-1, 1)
        VtV = Vs.T @ Vs
        Ik = torch.eye(len(VtV))
        A = torch.linalg.cholesky(VtV)
        B = torch.linalg.cholesky(VtV + Ik)
        A_inv = torch.inverse(A)
        C = torch.inverse(A_inv.T @ (B - Ik) @ A_inv)
        Kern_inv = torch.inverse(torch.inverse(C) + Vs.T @ Vs)
        dinv_sqrt = (d).sqrt().reshape(-1, 1)
        prior_sample = dinv_sqrt * samples
        gain_sample = dinv_sqrt * Vs @ Kern_inv @ (Vs.T @ samples)
        return self.mean + (prior_sample - gain_sample).T