jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_cg(self):
        size = 100
        matrix = torch.randn(size, size, dtype=torch.float64)
        matrix = matrix.matmul(matrix.mT)
        matrix.div_(matrix.norm())
        matrix.add_(torch.eye(matrix.size(-1), dtype=torch.float64).mul_(1e-1))

        # set up vector rhs
        rhs = torch.randn(size, dtype=torch.float64)

        # basic solve
        solves = linear_cg(matrix.matmul, rhs=rhs, max_iter=size)

        # solve with init value
        init = torch.randn(size, dtype=torch.float64)
        solves_with_init = linear_cg(matrix.matmul, rhs=rhs, max_iter=size, initial_guess=init)

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

        # set up matrix rhs
        numcols = 50
        rhs = torch.randn(size, numcols, dtype=torch.float64)

        # basic solve
        solves = linear_cg(matrix.matmul, rhs=rhs, max_iter=size)

        # solve with init value
        init = torch.randn(size, numcols, dtype=torch.float64)
        solves_with_init = linear_cg(matrix.matmul, rhs=rhs, max_iter=size, initial_guess=init)

        # Check cg
        actual = torch.cholesky_solve(rhs, matrix_chol)
        self.assertTrue(torch.allclose(solves, actual, atol=1e-3, rtol=1e-4))
        self.assertTrue(torch.allclose(solves_with_init, actual, atol=1e-3, rtol=1e-4))
//  [python]
def test_cg(self):
        size = 100
        matrix = torch.randn(size, size, dtype=torch.float64)
        matrix = matrix.matmul(matrix.mT)
        matrix.div_(matrix.norm())
        matrix.add_(torch.eye(matrix.size(-1), dtype=torch.float64).mul_(1e-1))

        # set up vector rhs
        rhs = torch.randn(size, dtype=torch.float64)

        # basic solve
        solves = linear_cg(matrix.matmul, rhs=rhs, max_iter=size)

        # solve with init value
        init = torch.randn(size, dtype=torch.float64)
        solves_with_init = linear_cg(matrix.matmul, rhs=rhs, max_iter=size, initial_guess=init)

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

        # set up matrix rhs
        numcols = 50
        rhs = torch.randn(size, numcols, dtype=torch.float64)

        # basic solve
        solves = linear_cg(matrix.matmul, rhs=rhs, max_iter=size)

        # solve with init value
        init = torch.randn(size, numcols, dtype=torch.float64)
        solves_with_init = linear_cg(matrix.matmul, rhs=rhs, max_iter=size, initial_guess=init)

        # Check cg
        actual = torch.cholesky_solve(rhs, matrix_chol)
        self.assertTrue(torch.allclose(solves, actual, atol=1e-3, rtol=1e-4))
        self.assertTrue(torch.allclose(solves_with_init, actual, atol=1e-3, rtol=1e-4))
//  [python]
def fasterprune(
        self, args, sparsity, mask=None, prune_n=0, prune_m=0, blocksize=128, percdamp=.01
    ):
        W = self.layer.weight.data.clone()
        if isinstance(self.layer, nn.Conv2d):
            W = W.flatten(1)
        if isinstance(self.layer, transformers.Conv1D):
            W = W.t()
        W = W.float()

        tick = time.time()

        H = self.H
        del self.H
        dead = torch.diag(H) == 0
        H[dead, dead] = 1
        W[:, dead] = 0

        Losses = torch.zeros(self.rows, device=self.dev)

        damp = percdamp * torch.mean(torch.diag(H))
        diag = torch.arange(self.columns, device=self.dev)
        H[diag, diag] += damp
        H = torch.cholesky(H)
        H = torch.cholesky_inverse(H)
        H = torch.cholesky(H, upper=True)
        Hinv = H

        for i1 in range(0, self.columns, blocksize):
            i2 = min(i1 + blocksize, self.columns)
            count = i2 - i1

            W1 = W[:, i1:i2].clone()
            Q1 = torch.zeros_like(W1)
            Err1 = torch.zeros_like(W1)
            Losses1 = torch.zeros_like(W1)
            Hinv1 = Hinv[i1:i2, i1:i2]

            if prune_n == 0 or mask is not None: 
                if mask is not None:
                    mask1 = mask[:, i1:i2]
                else:
                    # tmp = W1 ** 2 / (torch.diag(Hinv1).reshape((1, -1))) ** 2
                    if "wanda" in args.prune_method:
                        tmp = torch.abs(W1) * torch.sqrt(self.scaler_row[i1:i2].reshape((1,-1)))
                    elif "mag" in args.prune_method:
                        tmp = torch.abs(W1)
                    thresh = torch.sort(tmp.flatten())[0][int(tmp.numel() * sparsity)]
                    mask1 = tmp <= thresh
            else:
                mask1 = torch.zeros_like(W1) == 1

            for i in range(count):
                w = W1[:, i]
                d = Hinv1[i, i]

                if prune_n != 0 and i % prune_m == 0 and mask is None:
                    # tmp = W1[:, i:(i + prune_m)] ** 2 / (torch.diag(Hinv1)[i:(i + prune_m)].reshape((1, -1))) ** 2
                    if "wanda" in args.prune_method:
                        tmp = torch.abs(W1[:, i:(i+prune_m)]) * torch.sqrt(self.scaler_row[(i+i1):(i+i1+prune_m)].reshape((1,-1)))
                    elif "mag" in args.prune_method:
                        tmp = torch.abs(W1[:, i:(i+prune_m)])
                    mask1.scatter_(1, i + torch.topk(tmp, prune_n, dim=1, largest=False)[1], True)

                q = w.clone()
                q[mask1[:, i]] = 0

                Q1[:, i] = q
                Losses1[:, i] = (w - q) ** 2 / d ** 2

                err1 = (w - q) / d 
                W1[:, i:] -= err1.unsqueeze(1).matmul(Hinv1[i, i:].unsqueeze(0))
                Err1[:, i] = err1

            W[:, i1:i2] = Q1
            Losses += torch.sum(Losses1, 1) / 2

            W[:, i2:] -= Err1.matmul(Hinv[i1:i2, i2:])

        torch.cuda.synchronize()
        if isinstance(self.layer, transformers.Conv1D):
            W = W.t()
        self.layer.weight.data = W.reshape(self.layer.weight.shape).to(self.layer.weight.data.dtype)
//  [python]
def fasterprune(
        self, args, sparsity, mask=None, prune_n=0, prune_m=0, blocksize=128, percdamp=.01
    ):
        W = self.layer.weight.data.clone()
        if isinstance(self.layer, nn.Conv2d):
            W = W.flatten(1)
        if isinstance(self.layer, transformers.Conv1D):
            W = W.t()
        W = W.float()

        tick = time.time()

        H = self.H
        del self.H
        dead = torch.diag(H) == 0
        H[dead, dead] = 1
        W[:, dead] = 0

        Losses = torch.zeros(self.rows, device=self.dev)

        damp = percdamp * torch.mean(torch.diag(H))
        diag = torch.arange(self.columns, device=self.dev)
        H[diag, diag] += damp
        H = torch.linalg.cholesky(H)
        H = torch.cholesky_inverse(H)
        H = torch.linalg.cholesky(H, upper=True)
        Hinv = H

        for i1 in range(0, self.columns, blocksize):
            i2 = min(i1 + blocksize, self.columns)
            count = i2 - i1

            W1 = W[:, i1:i2].clone()
            Q1 = torch.zeros_like(W1)
            Err1 = torch.zeros_like(W1)
            Losses1 = torch.zeros_like(W1)
            Hinv1 = Hinv[i1:i2, i1:i2]

            if prune_n == 0 or mask is not None: 
                if mask is not None:
                    mask1 = mask[:, i1:i2]
                else:
                    # tmp = W1 ** 2 / (torch.diag(Hinv1).reshape((1, -1))) ** 2
                    if "wanda" in args.prune_method:
                        tmp = torch.abs(W1) * torch.sqrt(self.scaler_row[i1:i2].reshape((1,-1)))
                    elif "mag" in args.prune_method:
                        tmp = torch.abs(W1)
                    thresh = torch.sort(tmp.flatten())[0][int(tmp.numel() * sparsity)]
                    mask1 = tmp <= thresh
            else:
                mask1 = torch.zeros_like(W1) == 1

            for i in range(count):
                w = W1[:, i]
                d = Hinv1[i, i]

                if prune_n != 0 and i % prune_m == 0 and mask is None:
                    # tmp = W1[:, i:(i + prune_m)] ** 2 / (torch.diag(Hinv1)[i:(i + prune_m)].reshape((1, -1))) ** 2
                    if "wanda" in args.prune_method:
                        tmp = torch.abs(W1[:, i:(i+prune_m)]) * torch.sqrt(self.scaler_row[(i+i1):(i+i1+prune_m)].reshape((1,-1)))
                    elif "mag" in args.prune_method:
                        tmp = torch.abs(W1[:, i:(i+prune_m)])
                    mask1.scatter_(1, i + torch.topk(tmp, prune_n, dim=1, largest=False)[1], True)

                q = w.clone()
                q[mask1[:, i]] = 0

                Q1[:, i] = q
                Losses1[:, i] = (w - q) ** 2 / d ** 2

                err1 = (w - q) / d 
                W1[:, i:] -= err1.unsqueeze(1).matmul(Hinv1[i, i:].unsqueeze(0))
                Err1[:, i] = err1

            W[:, i1:i2] = Q1
            Losses += torch.sum(Losses1, 1) / 2

            W[:, i2:] -= Err1.matmul(Hinv[i1:i2, i2:])

        torch.cuda.synchronize()
        if isinstance(self.layer, transformers.Conv1D):
            W = W.t()
        self.layer.weight.data = W.reshape(self.layer.weight.shape).to(self.layer.weight.data.dtype)
//  [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 select_sigma_points(
        self,
        input_mean: torch.Tensor,
        input_covariance: types.CovarianceTorch,
    ) -> torch.Tensor:
        """Select sigma points.

        Args:
            input_mean (torch.Tensor): Distribution mean. Shape should be `(N, dim)`.
            input_covariance (torch.Tensor): Distribution covariance. Shape should be
                `(N, dim, dim)`.
        Returns:
            torch.Tensor: Selected sigma points, with shape `(N, 2 * dim + 1, dim)`.
        """

        N, dim = input_mean.shape
        assert input_covariance.shape == (N, dim, dim)
        assert dim == self._dim
        return self.select_sigma_points_square_root(
            input_mean, input_scale_tril=torch.cholesky(input_covariance)
        )
//  [python]
def select_sigma_points(
        self,
        input_mean: torch.Tensor,
        input_covariance: types.CovarianceTorch,
    ) -> torch.Tensor:
        """Select sigma points.

        Args:
            input_mean (torch.Tensor): Distribution mean. Shape should be `(N, dim)`.
            input_covariance (torch.Tensor): Distribution covariance. Shape should be
                `(N, dim, dim)`.
        Returns:
            torch.Tensor: Selected sigma points, with shape `(N, 2 * dim + 1, dim)`.
        """

        N, dim = input_mean.shape
        assert input_covariance.shape == (N, dim, dim)
        assert dim == self._dim
        return self.select_sigma_points_square_root(
            input_mean, input_scale_tril=torch.linalg.cholesky(input_covariance)
        )
//  [python]
def _calculate_log_det(self, var):
        """
        Calculate log determinant in log space, to prevent overflow errors.
        args:
            var:            torch.Tensor (1, k, d, d)
        """
        log_det = torch.empty(size=(self.n_components,)).to(var.device)

        for k in range(self.n_components):
            log_det[k] = (
                2 * torch.log(torch.diagonal(torch.cholesky(var[0, k]))).sum()
            )

        return log_det.unsqueeze(-1)
//  [python]
def _calculate_log_det(self, var):
        """
        Calculate log determinant in log space, to prevent overflow errors.
        args:
            var:            torch.Tensor (1, k, d, d)
        """
        log_det = torch.empty(size=(self.n_components,)).to(var.device)

        for k in range(self.n_components):
            log_det[k] = (
                2 * torch.log(torch.diagonal(torch.linalg.cholesky(var[0, k]))).sum()
            )

        return log_det.unsqueeze(-1)
//  [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 logdet(self, X):

        if self.faster_logdet:
            return 2 * torch.sum(torch.log(torch.diag(torch.cholesky(X, upper=True))))
        else:
            return torch.logdet(X)
//  [python]
def logdet(self, X):

        if self.faster_logdet:
            return 2 * torch.sum(torch.log(torch.diag(torch.linalg.cholesky(X, upper=True))))
        else:
            return torch.logdet(X)
//  [python]
def forward(self, x, L, mass, evals, evecs):

        # project times to the positive halfspace
        # (and away from 0 in the incredibly rare chance that they get stuck)
        with torch.no_grad():
            self.diffusion_time.data = torch.clamp(self.diffusion_time, min=1e-8)

        if x.shape[-1] != self.C_inout:
            raise ValueError(
                "Tensor has wrong shape = {}. Last dim shape should have number of channels = {}".format(
                    x.shape, self.C_inout))

        if self.method == 'spectral':

            # Transform to spectral
            x_spec = to_basis(x, evecs, mass)

            # Diffuse
            time = self.diffusion_time
            diffusion_coefs = torch.exp(-evals.unsqueeze(-1) * time.unsqueeze(0))
            x_diffuse_spec = diffusion_coefs * x_spec

            # Transform back to per-vertex 
            x_diffuse = from_basis(x_diffuse_spec, evecs)
            
        elif self.method == 'implicit_dense':
            V = x.shape[-2]

            # Form the dense matrices (M + tL) with dims (B,C,V,V)
            mat_dense = L.to_dense().unsqueeze(1).expand(-1, self.C_inout, V, V).clone()
            mat_dense *= self.diffusion_time.unsqueeze(0).unsqueeze(-1).unsqueeze(-1)
            mat_dense += torch.diag_embed(mass).unsqueeze(1)

            # Factor the system
            cholesky_factors = torch.cholesky(mat_dense)
            
            # Solve the system
            rhs = x * mass.unsqueeze(-1)
            rhsT = torch.transpose(rhs, 1, 2).unsqueeze(-1)
            sols = torch.cholesky_solve(rhsT, cholesky_factors)
            x_diffuse = torch.transpose(sols.squeeze(-1), 1, 2)

        else:
            raise ValueError("unrecognized method")


        return x_diffuse
//  [python]
def forward(self, x, L, mass, evals, evecs):

        # project times to the positive halfspace
        # (and away from 0 in the incredibly rare chance that they get stuck)
        with torch.no_grad():
            self.diffusion_time.data = torch.clamp(self.diffusion_time, min=1e-8)

        if x.shape[-1] != self.C_inout:
            raise ValueError(
                "Tensor has wrong shape = {}. Last dim shape should have number of channels = {}".format(
                    x.shape, self.C_inout))

        if self.method == 'spectral':

            # Transform to spectral
            x_spec = to_basis(x, evecs, mass)

            # Diffuse
            time = self.diffusion_time
            diffusion_coefs = torch.exp(-evals.unsqueeze(-1) * time.unsqueeze(0))
            x_diffuse_spec = diffusion_coefs * x_spec

            # Transform back to per-vertex 
            x_diffuse = from_basis(x_diffuse_spec, evecs)
            
        elif self.method == 'implicit_dense':
            V = x.shape[-2]

            # Form the dense matrices (M + tL) with dims (B,C,V,V)
            mat_dense = L.to_dense().unsqueeze(1).expand(-1, self.C_inout, V, V).clone()
            mat_dense *= self.diffusion_time.unsqueeze(0).unsqueeze(-1).unsqueeze(-1)
            mat_dense += torch.diag_embed(mass).unsqueeze(1)

            # Factor the system
            cholesky_factors = torch.linalg.cholesky(mat_dense)
            
            # Solve the system
            rhs = x * mass.unsqueeze(-1)
            rhsT = torch.transpose(rhs, 1, 2).unsqueeze(-1)
            sols = torch.cholesky_solve(rhsT, cholesky_factors)
            x_diffuse = torch.transpose(sols.squeeze(-1), 1, 2)

        else:
            raise ValueError("unrecognized method")


        return x_diffuse
//  [python]
def predict(self, inputs: torch.Tensor) -> ModelOutput:
        """The input data is encoded and decoded without computing loss

        Args:
            inputs (torch.Tensor): The input data to be reconstructed, as well as to generate the embedding.

        Returns:
            ModelOutput: An instance of ModelOutput containing reconstruction, raw embedding (output of encoder), and the final embedding (output of metric)
        """
        encoder_output = self.encoder(inputs)
        mu, log_var = encoder_output.embedding, encoder_output.log_covariance

        std = torch.exp(0.5 * log_var)
        z0, _ = self._sample_gauss(mu, std)

        z = z0

        G = self.G(z)
        G_inv = self.G_inv(z)
        L = torch.cholesky(G)

        G_log_det = -torch.logdet(G_inv)

        gamma = torch.randn_like(z0, device=inputs.device)
        rho = gamma / self.beta_zero_sqrt
        beta_sqrt_old = self.beta_zero_sqrt

        # sample \rho from N(0, G)
        rho = (L @ rho.unsqueeze(-1)).squeeze(-1)

        recon_x = self.decoder(z)["reconstruction"]

        for k in range(self.n_lf):
            # perform leapfrog steps

            # step 1
            rho_ = self._leap_step_1(recon_x, inputs, z, rho, G_inv, G_log_det)

            # step 2
            z = self._leap_step_2(recon_x, inputs, z, rho_, G_inv, G_log_det)

            recon_x = self.decoder(z)["reconstruction"]

            # compute metric value on new z using final metric
            G = self.G(z)
            G_inv = self.G_inv(z)

            G_log_det = -torch.logdet(G_inv)

            # step 3
            rho__ = self._leap_step_3(recon_x, inputs, z, rho_, G_inv, G_log_det)

            # tempering
            beta_sqrt = self._tempering(k + 1, self.n_lf)
            rho = (beta_sqrt_old / beta_sqrt) * rho__
            beta_sqrt_old = beta_sqrt

        output = ModelOutput(
            recon_x=recon_x,
            raw_embedding=encoder_output.embedding,
            embedding=z if self.n_lf > 0 else encoder_output.embedding,
        )

        return output
//  [python]
def predict(self, inputs: torch.Tensor) -> ModelOutput:
        """The input data is encoded and decoded without computing loss

        Args:
            inputs (torch.Tensor): The input data to be reconstructed, as well as to generate the embedding.

        Returns:
            ModelOutput: An instance of ModelOutput containing reconstruction, raw embedding (output of encoder), and the final embedding (output of metric)
        """
        encoder_output = self.encoder(inputs)
        mu, log_var = encoder_output.embedding, encoder_output.log_covariance

        std = torch.exp(0.5 * log_var)
        z0, _ = self._sample_gauss(mu, std)

        z = z0

        G = self.G(z)
        G_inv = self.G_inv(z)
        L = torch.linalg.cholesky(G)

        G_log_det = -torch.logdet(G_inv)

        gamma = torch.randn_like(z0, device=inputs.device)
        rho = gamma / self.beta_zero_sqrt
        beta_sqrt_old = self.beta_zero_sqrt

        # sample \rho from N(0, G)
        rho = (L @ rho.unsqueeze(-1)).squeeze(-1)

        recon_x = self.decoder(z)["reconstruction"]

        for k in range(self.n_lf):
            # perform leapfrog steps

            # step 1
            rho_ = self._leap_step_1(recon_x, inputs, z, rho, G_inv, G_log_det)

            # step 2
            z = self._leap_step_2(recon_x, inputs, z, rho_, G_inv, G_log_det)

            recon_x = self.decoder(z)["reconstruction"]

            # compute metric value on new z using final metric
            G = self.G(z)
            G_inv = self.G_inv(z)

            G_log_det = -torch.logdet(G_inv)

            # step 3
            rho__ = self._leap_step_3(recon_x, inputs, z, rho_, G_inv, G_log_det)

            # tempering
            beta_sqrt = self._tempering(k + 1, self.n_lf)
            rho = (beta_sqrt_old / beta_sqrt) * rho__
            beta_sqrt_old = beta_sqrt

        output = ModelOutput(
            recon_x=recon_x,
            raw_embedding=encoder_output.embedding,
            embedding=z if self.n_lf > 0 else encoder_output.embedding,
        )

        return output
//  [python]
def belief_covariance(self) -> types.CovarianceTorch:
        """Posterior covariance. Shape should be `(N, state_dim, state_dim)`."""
        return fannypack.utils.cholesky_inverse(torch.cholesky(self.information_matrix))
//  [python]
def belief_covariance(self) -> types.CovarianceTorch:
        """Posterior covariance. Shape should be `(N, state_dim, state_dim)`."""
        return fannypack.utils.cholesky_inverse(torch.linalg.cholesky(self.information_matrix))
//  [python]
def belief_covariance(self, covariance: types.CovarianceTorch):
        self.information_matrix = fannypack.utils.cholesky_inverse(
            torch.cholesky(covariance)
        )
//  [python]
def belief_covariance(self, covariance: types.CovarianceTorch):
        self.information_matrix = fannypack.utils.cholesky_inverse(
            torch.linalg.cholesky(covariance)
        )
//  [python]
def _cholesky(x):
    return x.sqrt() if x.dim() == 1 else t.cholesky(x)
//  [python]
def _cholesky(x):
    return x.sqrt() if x.dim() == 1 else torch.linalg.cholesky(x)
//  [python]
def stable_cholesky(
    K: Float[Tensor, "... d d"],
    mult_eps: float = 1e-8,
    abs_eps: float = 1e-8,
    max_tries: int = 1000,
) -> Float[Tensor, "... d d"]:
    for i in range(max_tries):
        try:
            scaled_mult_eps = mult_eps * (1.1**i)
            scaled_abs_eps = abs_eps * (1.1**i)
            # L = t.linalg.cholesky(stabilise(K, i * mult_eps, i * abs_eps))
            L = t.cholesky(stabilise(K, scaled_mult_eps, scaled_abs_eps))
            return L
        except LinAlgError:
            logging.debug(f"Chokesky decomposition failed ({i})")
            continue
    raise ValueError(f"Could not calculate Cholesky decomposition of {K}")
//  [python]
def stable_cholesky(
    K: Float[Tensor, "... d d"],
    mult_eps: float = 1e-8,
    abs_eps: float = 1e-8,
    max_tries: int = 1000,
) -> Float[Tensor, "... d d"]:
    for i in range(max_tries):
        try:
            scaled_mult_eps = mult_eps * (1.1**i)
            scaled_abs_eps = abs_eps * (1.1**i)
            # L = t.linalg.cholesky(stabilise(K, i * mult_eps, i * abs_eps))
            L = t.linalg.cholesky(stabilise(K, scaled_mult_eps, scaled_abs_eps))
            return L
        except LinAlgError:
            logging.debug(f"Chokesky decomposition failed ({i})")
            continue
    raise ValueError(f"Could not calculate Cholesky decomposition of {K}")
//  [python]
def test_gaussian_tensordot(
    dot_dims, x_batch_shape, x_dim, x_rank, y_batch_shape, y_dim, y_rank
):
    x_rank = min(x_rank, x_dim)
    y_rank = min(y_rank, y_dim)
    x = random_gaussian(x_batch_shape, x_dim, x_rank)
    y = random_gaussian(y_batch_shape, y_dim, y_rank)
    na = x_dim - dot_dims
    nb = dot_dims
    nc = y_dim - dot_dims
    try:
        torch.cholesky(x.precision[..., na:, na:] + y.precision[..., :nb, :nb])
    except Exception:
        pytest.skip("Cannot marginalize the common variables of two Gaussians.")

    z = gaussian_tensordot(x, y, dot_dims)
    assert z.dim() == x_dim + y_dim - 2 * dot_dims

    # We make these precision matrices positive definite to test the math
    x.precision = x.precision + 1e-1 * torch.eye(x.dim())
    y.precision = y.precision + 1e-1 * torch.eye(y.dim())
    z = gaussian_tensordot(x, y, dot_dims)
    # compare against broadcasting, adding, and marginalizing
    precision = pad(x.precision, (0, nc, 0, nc)) + pad(y.precision, (na, 0, na, 0))
    info_vec = pad(x.info_vec, (0, nc)) + pad(y.info_vec, (na, 0))
    covariance = torch.inverse(precision)
    loc = (
        covariance.matmul(info_vec.unsqueeze(-1)).squeeze(-1)
        if info_vec.size(-1) > 0
        else info_vec
    )
    z_covariance = torch.inverse(z.precision)
    z_loc = z_covariance.matmul(
        z.info_vec.view(z.info_vec.shape + (int(z.dim() > 0),))
    ).sum(-1)
    assert_close(loc[..., :na], z_loc[..., :na])
    assert_close(loc[..., x_dim:], z_loc[..., na:])
    assert_close(covariance[..., :na, :na], z_covariance[..., :na, :na])
    assert_close(covariance[..., :na, x_dim:], z_covariance[..., :na, na:])
    assert_close(covariance[..., x_dim:, :na], z_covariance[..., na:, :na])
    assert_close(covariance[..., x_dim:, x_dim:], z_covariance[..., na:, na:])

    # Assume a = c = 0, integrate out b
    # FIXME: this might be not a stable way to compute integral
    num_samples = 200000
    scale = 20
    # generate samples in [-10, 10]
    value_b = torch.rand((num_samples,) + z.batch_shape + (nb,)) * scale - scale / 2
    value_x = pad(value_b, (na, 0))
    value_y = pad(value_b, (0, nc))
    expect = torch.logsumexp(x.log_density(value_x) + y.log_density(value_y), dim=0)
    expect += math.log(scale**nb / num_samples)
    actual = z.log_density(torch.zeros(z.batch_shape + (z.dim(),)))
    # TODO(fehiepsi): find some condition to make this test stable, so we can compare large value
    # log densities.
    assert_close(actual.clamp(max=10.0), expect.clamp(max=10.0), atol=0.1, rtol=0.1)
//  [python]
def test_gaussian_tensordot(
    dot_dims, x_batch_shape, x_dim, x_rank, y_batch_shape, y_dim, y_rank
):
    x_rank = min(x_rank, x_dim)
    y_rank = min(y_rank, y_dim)
    x = random_gaussian(x_batch_shape, x_dim, x_rank)
    y = random_gaussian(y_batch_shape, y_dim, y_rank)
    na = x_dim - dot_dims
    nb = dot_dims
    nc = y_dim - dot_dims
    try:
        torch.linalg.cholesky(x.precision[..., na:, na:] + y.precision[..., :nb, :nb])
    except Exception:
        pytest.skip("Cannot marginalize the common variables of two Gaussians.")

    z = gaussian_tensordot(x, y, dot_dims)
    assert z.dim() == x_dim + y_dim - 2 * dot_dims

    # We make these precision matrices positive definite to test the math
    x.precision = x.precision + 1e-1 * torch.eye(x.dim())
    y.precision = y.precision + 1e-1 * torch.eye(y.dim())
    z = gaussian_tensordot(x, y, dot_dims)
    # compare against broadcasting, adding, and marginalizing
    precision = pad(x.precision, (0, nc, 0, nc)) + pad(y.precision, (na, 0, na, 0))
    info_vec = pad(x.info_vec, (0, nc)) + pad(y.info_vec, (na, 0))
    covariance = torch.inverse(precision)
    loc = (
        covariance.matmul(info_vec.unsqueeze(-1)).squeeze(-1)
        if info_vec.size(-1) > 0
        else info_vec
    )
    z_covariance = torch.inverse(z.precision)
    z_loc = z_covariance.matmul(
        z.info_vec.view(z.info_vec.shape + (int(z.dim() > 0),))
    ).sum(-1)
    assert_close(loc[..., :na], z_loc[..., :na])
    assert_close(loc[..., x_dim:], z_loc[..., na:])
    assert_close(covariance[..., :na, :na], z_covariance[..., :na, :na])
    assert_close(covariance[..., :na, x_dim:], z_covariance[..., :na, na:])
    assert_close(covariance[..., x_dim:, :na], z_covariance[..., na:, :na])
    assert_close(covariance[..., x_dim:, x_dim:], z_covariance[..., na:, na:])

    # Assume a = c = 0, integrate out b
    # FIXME: this might be not a stable way to compute integral
    num_samples = 200000
    scale = 20
    # generate samples in [-10, 10]
    value_b = torch.rand((num_samples,) + z.batch_shape + (nb,)) * scale - scale / 2
    value_x = pad(value_b, (na, 0))
    value_y = pad(value_b, (0, nc))
    expect = torch.logsumexp(x.log_density(value_x) + y.log_density(value_y), dim=0)
    expect += math.log(scale**nb / num_samples)
    actual = z.log_density(torch.zeros(z.batch_shape + (z.dim(),)))
    # TODO(fehiepsi): find some condition to make this test stable, so we can compare large value
    # log densities.
    assert_close(actual.clamp(max=10.0), expect.clamp(max=10.0), atol=0.1, rtol=0.1)
//  [python]
def _predict_step(self, *, controls: types.ControlsTorch) -> None:
        # Get previous belief
        prev_mean = self._belief_mean
        prev_covariance = self.belief_covariance
        N, state_dim = prev_mean.shape

        # Compute mu_{t+1|t}, covariance, and Jacobian
        pred_mean, dynamics_tril = self.dynamics_model(
            initial_states=prev_mean, controls=controls
        )
        dynamics_covariance = dynamics_tril @ dynamics_tril.transpose(-1, -2)
        dynamics_A_matrix = self.dynamics_model.jacobian(
            initial_states=prev_mean, controls=controls
        )
        assert dynamics_covariance.shape == (N, state_dim, state_dim)
        assert dynamics_A_matrix.shape == (N, state_dim, state_dim)

        # Calculate Sigma_{t+1|t}
        pred_information_matrix = fannypack.utils.cholesky_inverse(
            torch.cholesky(
                dynamics_A_matrix
                @ prev_covariance
                @ dynamics_A_matrix.transpose(-1, -2)
                + dynamics_covariance
            )
        )
        pred_information_vector = (
            pred_information_matrix @ pred_mean[:, :, None]
        ).squeeze(-1)

        # Update internal state
        self._belief_mean = pred_mean
        self.information_matrix = pred_information_matrix
        self.information_vector = pred_information_vector
//  [python]
def _predict_step(self, *, controls: types.ControlsTorch) -> None:
        # Get previous belief
        prev_mean = self._belief_mean
        prev_covariance = self.belief_covariance
        N, state_dim = prev_mean.shape

        # Compute mu_{t+1|t}, covariance, and Jacobian
        pred_mean, dynamics_tril = self.dynamics_model(
            initial_states=prev_mean, controls=controls
        )
        dynamics_covariance = dynamics_tril @ dynamics_tril.transpose(-1, -2)
        dynamics_A_matrix = self.dynamics_model.jacobian(
            initial_states=prev_mean, controls=controls
        )
        assert dynamics_covariance.shape == (N, state_dim, state_dim)
        assert dynamics_A_matrix.shape == (N, state_dim, state_dim)

        # Calculate Sigma_{t+1|t}
        pred_information_matrix = fannypack.utils.cholesky_inverse(
            torch.linalg.cholesky(
                dynamics_A_matrix
                @ prev_covariance
                @ dynamics_A_matrix.transpose(-1, -2)
                + dynamics_covariance
            )
        )
        pred_information_vector = (
            pred_information_matrix @ pred_mean[:, :, None]
        ).squeeze(-1)

        # Update internal state
        self._belief_mean = pred_mean
        self.information_matrix = pred_information_matrix
        self.information_vector = pred_information_vector
//  [python]
def ldlp_admm(H, rho = 0.1, niters = 100):
    n = H.shape[0]
    L = torch.cholesky(2 * H + rho * torch.eye(n,device=H.device))
    Linv = torch.inverse(L)
    M = (torch.arange(n)[None,:] < torch.arange(n)[:,None]).to(H.device)
    MH = M * H
    X = torch.zeros(n,n,device=H.device)
    Z = torch.zeros(n,n,device=H.device)
    W = torch.zeros(n,n,device=H.device)
    for ii in range(100):
        X = (((rho * Z - rho * W - 2 * MH) @ Linv.T) * M) @ Linv
        C = torch.diag(1 / torch.max(torch.tensor(1.0,device=H.device),((X + W).T @ (X + W)).diag().sqrt()))
        Z = (X + W) @ C
        W = W + X - Z
        objective = ((Z + torch.eye(n,device=H.device)) @ H @ (Z + torch.eye(n,device=H.device)).T).trace()
        print(f"ldlp_admm iter {ii}: {objective}")
    return Z
//  [python]
def ldlp_admm(H, rho = 0.1, niters = 100):
    n = H.shape[0]
    L = torch.linalg.cholesky(2 * H + rho * torch.eye(n,device=H.device))
    Linv = torch.linalg.inv(L)
    M = (torch.arange(n)[None,:] < torch.arange(n)[:,None]).to(H.device)
    MH = M * H
    X = torch.zeros(n,n,device=H.device)
    Z = torch.zeros(n,n,device=H.device)
    W = torch.zeros(n,n,device=H.device)
    for ii in range(100):
        X = (((rho * Z - rho * W - 2 * MH) @ Linv.T) * M) @ Linv
        C = torch.diag(1 / torch.max(torch.tensor(1.0,device=H.device),((X + W).T @ (X + W)).diag().sqrt()))
        Z = (X + W) @ C
        W = W + X - Z
        objective = ((Z + torch.eye(n,device=H.device)) @ H @ (Z + torch.eye(n,device=H.device)).T).trace()
        print(f"ldlp_admm iter {ii}: {objective}")
    return Z
//  [python]
def forward(self, X, y):
        if self.use_Trans:
            Xm = self.F_transform(X)
            self.Xm_bar = self.F_transform(self.X_bar)
        else:
            Xm = X
            self.Xm_bar = self.X_bar
        K, B = self.GlobalKernel(Xm, Xm)

        dK = K.diagonal()
        dK += self.global_gp_noise_std**2 + self.jitter
        # print(K)
        L = torch.cholesky(K)
        alpha = torch.cholesky_solve(y, L)
        
        Apart1 = y.T@alpha
        Apart2 = torch.sum(torch.log(L.diagonal()))

        A = 0.5*( Apart1 + Apart2)[0, 0]
        
        return (A+B)/X.nelement()
//  [python]
def forward(self, X, y):
        if self.use_Trans:
            Xm = self.F_transform(X)
            self.Xm_bar = self.F_transform(self.X_bar)
        else:
            Xm = X
            self.Xm_bar = self.X_bar
        K, B = self.GlobalKernel(Xm, Xm)

        dK = K.diagonal()
        dK += self.global_gp_noise_std**2 + self.jitter
        # print(K)
        L = torch.linalg.cholesky(K)
        alpha = torch.cholesky_solve(y, L)
        
        Apart1 = y.T@alpha
        Apart2 = torch.sum(torch.log(L.diagonal()))

        A = 0.5*( Apart1 + Apart2)[0, 0]
        
        return (A+B)/X.nelement()
//  [python]
def round_ldl(
    w,
    H,
    nbits,
    n_greedy_passes=9,
    unbiased=False
):
    """
    w in R^{m,d}
    d: input_shape, m: output_shape
    note: this has been updated with a (hopefully) more efficient LDL pass
    """
    assert (not unbiased) or (n_greedy_passes == 0), "greedy passes are incompatible with unbiased LDL rounding"
    (d, d_) = H.shape
    assert (d == d_)
    (m, d) = w.shape
    L = torch.cholesky(H)
    L = L @ torch.diag(1/torch.diag(L))
    L = L - torch.eye(d, device=L.device)
    if unbiased:
        eta = torch.rand(w.shape).to(w.device)
    else:
        eta = 0.5 * torch.ones(w.shape).to(w.device)
    w_hat = w.clone()
    for i in reversed(range(d)):
        w_hat[:,i] = torch.clamp(torch.floor(w[:,i] + (w[:,i:] - w_hat[:,i:]) @ L[i:,i] + eta[:,i]), min=0, max=2**nbits-1)
    
    wr = w_hat.clone()
    s = w_hat - w
    H = H / H.diag().max()

    for igp in range(n_greedy_passes):
        for i in reversed(range(d)):
            Hs = s @ H[:, i]
            epsXTsj = wr[:, i] - torch.round(wr[:, i] - Hs / H[i,i])
            wr[:, i] -= epsXTsj
            s[:, i] -= epsXTsj
        wr = torch.clamp(wr, min=0, max=2**nbits - 1)
        if ((w_hat == wr).all()):
            sys.stderr.write(f"breaking after {igp+1} greedy passes found fixed point")
            break
        w_hat.copy_(wr)
    
    wr_counts = check_nbits(wr, nbits)
    return wr
//  [python]
def round_ldl(
    w,
    H,
    nbits,
    n_greedy_passes=9,
    unbiased=False
):
    """
    w in R^{m,d}
    d: input_shape, m: output_shape
    note: this has been updated with a (hopefully) more efficient LDL pass
    """
    assert (not unbiased) or (n_greedy_passes == 0), "greedy passes are incompatible with unbiased LDL rounding"
    (d, d_) = H.shape
    assert (d == d_)
    (m, d) = w.shape
    L = torch.linalg.cholesky(H)
    L = L @ torch.diag(1/torch.diag(L))
    L = L - torch.eye(d, device=L.device)
    if unbiased:
        eta = torch.rand(w.shape).to(w.device)
    else:
        eta = 0.5 * torch.ones(w.shape).to(w.device)
    w_hat = w.clone()
    for i in reversed(range(d)):
        w_hat[:,i] = torch.clamp(torch.floor(w[:,i] + (w[:,i:] - w_hat[:,i:]) @ L[i:,i] + eta[:,i]), min=0, max=2**nbits-1)
    
    wr = w_hat.clone()
    s = w_hat - w
    H = H / H.diag().max()

    for igp in range(n_greedy_passes):
        for i in reversed(range(d)):
            Hs = s @ H[:, i]
            epsXTsj = wr[:, i] - torch.round(wr[:, i] - Hs / H[i,i])
            wr[:, i] -= epsXTsj
            s[:, i] -= epsXTsj
        wr = torch.clamp(wr, min=0, max=2**nbits - 1)
        if ((w_hat == wr).all()):
            sys.stderr.write(f"breaking after {igp+1} greedy passes found fixed point")
            break
        w_hat.copy_(wr)
    
    wr_counts = check_nbits(wr, nbits)
    return wr
//  [python]
def fasterquant(self, blocksize=128, percdamp=0.01):
        W = self.layer.weight.data.clone()
        W = W.float()

        tick = time.time()

        self.quantizer.calibration(W)

        H = self.H
        del self.H
        dead = torch.diag(H) == 0
        H[dead, dead] = 1
        W[:, dead] = 0

        Losses = torch.zeros_like(W)
        Q = torch.zeros_like(W)

        damp = percdamp * torch.mean(torch.diag(H))
        diag = torch.arange(self.columns, device=self.dev)
        H[diag, diag] += damp
        H = torch.cholesky(H)
        H = torch.cholesky_inverse(H)
        H = torch.cholesky(H, upper=True)
        Hinv = H

        for i1 in range(0, self.columns, blocksize):
            i2 = min(i1 + blocksize, self.columns)
            count = i2 - i1

            W1 = W[:, i1:i2].clone()
            Q1 = torch.zeros_like(W1)
            Err1 = torch.zeros_like(W1)
            Losses1 = torch.zeros_like(W1)
            Hinv1 = Hinv[i1:i2, i1:i2]

            for i in range(count):
                w = W1[:, i]
                d = Hinv1[i, i]

                dim = 1
                if self.quantizer.scale.size(dim) != 1:
                    scale = self.quantizer.scale[:, i1 + i].unsqueeze(1)
                    if self.quantizer.round_zero_point is not None:
                        round_zero_point = self.quantizer.round_zero_point[
                            :, i1 + i
                        ].unsqueeze(1)
                    else:
                        round_zero_point = None
                else:
                    scale = self.quantizer.scale
                    round_zero_point = self.quantizer.round_zero_point
                q = self.quantizer.fake_quant(
                    w.unsqueeze(1), scale, round_zero_point
                ).flatten()
                Q1[:, i] = q
                Losses1[:, i] = (w - q) ** 2 / d**2

                err1 = (w - q) / d
                W1[:, i:] -= err1.unsqueeze(1).matmul(Hinv1[i, i:].unsqueeze(0))
                Err1[:, i] = err1

            Q[:, i1:i2] = Q1
            Losses[:, i1:i2] = Losses1 / 2

            W[:, i2:] -= Err1.matmul(Hinv[i1:i2, i2:])

            if DEBUG:
                self.layer.weight.data[:, :i2] = Q[:, :i2]
                self.layer.weight.data[:, i2:] = W[:, i2:]
                print(torch.sum((self.layer(self.inp1) - self.out1) ** 2))
                print(torch.sum(Losses))

        torch.cuda.synchronize()

        return (
            Q.reshape(self.layer.weight.shape),
            torch.sum(Losses).item(),
        )
//  [python]
def fasterquant(self, blocksize=128, percdamp=0.01):
        W = self.layer.weight.data.clone()
        W = W.float()

        tick = time.time()

        self.quantizer.calibration(W)

        H = self.H
        del self.H
        dead = torch.diag(H) == 0
        H[dead, dead] = 1
        W[:, dead] = 0

        Losses = torch.zeros_like(W)
        Q = torch.zeros_like(W)

        damp = percdamp * torch.mean(torch.diag(H))
        diag = torch.arange(self.columns, device=self.dev)
        H[diag, diag] += damp
        H = torch.linalg.cholesky(H)
        H = torch.cholesky_inverse(H)
        H = torch.linalg.cholesky(H, upper=True)
        Hinv = H

        for i1 in range(0, self.columns, blocksize):
            i2 = min(i1 + blocksize, self.columns)
            count = i2 - i1

            W1 = W[:, i1:i2].clone()
            Q1 = torch.zeros_like(W1)
            Err1 = torch.zeros_like(W1)
            Losses1 = torch.zeros_like(W1)
            Hinv1 = Hinv[i1:i2, i1:i2]

            for i in range(count):
                w = W1[:, i]
                d = Hinv1[i, i]

                dim = 1
                if self.quantizer.scale.size(dim) != 1:
                    scale = self.quantizer.scale[:, i1 + i].unsqueeze(1)
                    if self.quantizer.round_zero_point is not None:
                        round_zero_point = self.quantizer.round_zero_point[
                            :, i1 + i
                        ].unsqueeze(1)
                    else:
                        round_zero_point = None
                else:
                    scale = self.quantizer.scale
                    round_zero_point = self.quantizer.round_zero_point
                q = self.quantizer.fake_quant(
                    w.unsqueeze(1), scale, round_zero_point
                ).flatten()
                Q1[:, i] = q
                Losses1[:, i] = (w - q) ** 2 / d**2

                err1 = (w - q) / d
                W1[:, i:] -= err1.unsqueeze(1).matmul(Hinv1[i, i:].unsqueeze(0))
                Err1[:, i] = err1

            Q[:, i1:i2] = Q1
            Losses[:, i1:i2] = Losses1 / 2

            W[:, i2:] -= Err1.matmul(Hinv[i1:i2, i2:])

            if DEBUG:
                self.layer.weight.data[:, :i2] = Q[:, :i2]
                self.layer.weight.data[:, i2:] = W[:, i2:]
                print(torch.sum((self.layer(self.inp1) - self.out1) ** 2))
                print(torch.sum(Losses))

        torch.cuda.synchronize()

        return (
            Q.reshape(self.layer.weight.shape),
            torch.sum(Losses).item(),
        )
//  [python]
def get_LS(self, X, dim):  # Infer lengthscales for train_X (self.X)
        # print('Log: dim', dim)
        k = self.LocalKernel(
            self.Xm_bar[dim], self.Xm_bar[dim], dim)
        # Diagonal Solution from https://stackoverflow.com/a/48170846/13330701
        dk = k.diagonal()
        dk += self.local_gp_noise_std[dim]**2 + self.jitter
        c = torch.cholesky(k)
        alpha = torch.cholesky_solve(
            torch.log(self.local_ls[dim].reshape(-1,1)), c)
        k_star = self.LocalKernel(
            X[:, dim, None], self.Xm_bar[dim], dim)
        l = torch.exp(k_star@alpha)

        if self.training:
            k_star = self.LocalKernel(
                X[:, dim, None], self.Xm_bar[dim], dim)
            k_star_star = self.LocalKernel(
                X[:, dim, None], X[:, dim, None], dim)

            chol = torch.cholesky(k)
            v = torch.cholesky_solve(k_star.T, chol)

            k_post = k_star_star - k_star@v

            dk_post = k_post.diagonal()
            dk_post += self.jitter*10
            post_chol = torch.cholesky(k_post)
            B = torch.sum(torch.log(post_chol.diagonal()))
            return l, B
        else:
            return l
//  [python]
def get_LS(self, X, dim):  # Infer lengthscales for train_X (self.X)
        # print('Log: dim', dim)
        k = self.LocalKernel(
            self.Xm_bar[dim], self.Xm_bar[dim], dim)
        # Diagonal Solution from https://stackoverflow.com/a/48170846/13330701
        dk = k.diagonal()
        dk += self.local_gp_noise_std[dim]**2 + self.jitter
        c = torch.linalg.cholesky(k)
        alpha = torch.cholesky_solve(
            torch.log(self.local_ls[dim].reshape(-1,1)), c)
        k_star = self.LocalKernel(
            X[:, dim, None], self.Xm_bar[dim], dim)
        l = torch.exp(k_star@alpha)

        if self.training:
            k_star = self.LocalKernel(
                X[:, dim, None], self.Xm_bar[dim], dim)
            k_star_star = self.LocalKernel(
                X[:, dim, None], X[:, dim, None], dim)

            chol = torch.linalg.cholesky(k)
            v = torch.cholesky_solve(k_star.T, chol)

            k_post = k_star_star - k_star@v

            dk_post = k_post.diagonal()
            dk_post += self.jitter*10
            post_chol = torch.linalg.cholesky(k_post)
            B = torch.sum(torch.log(post_chol.diagonal()))
            return l, B
        else:
            return l
//  [python]
def predict(self, X, y, X_new):  # Predict at new locations
        if self.use_Trans:
            Xm = self.F_transform(X)
            Xm_new = self.F_transform(X_new)
            self.Xm_bar = self.F_transform(self.X_bar)
        else:
            Xm = X
            Xm_new = X_new
            self.Xm_bar = self.X_bar

        K = self.GlobalKernel(Xm, Xm)
        K_star = self.GlobalKernel(Xm_new, Xm)
        K_star_star = self.GlobalKernel(Xm_new, Xm_new)

        dK = K.diagonal()
        dK += self.global_gp_noise_std**2
        L = torch.cholesky(K)
        alpha = torch.cholesky_solve(y, L)

        pred_mean = K_star@alpha

        v = torch.cholesky_solve(K_star.T, L)
        pred_var = K_star_star - K_star@v

        dpred_var = pred_var.diagonal()
        dpred_var += self.global_gp_noise_std**2
        return pred_mean, pred_var
//  [python]
def predict(self, X, y, X_new):  # Predict at new locations
        if self.use_Trans:
            Xm = self.F_transform(X)
            Xm_new = self.F_transform(X_new)
            self.Xm_bar = self.F_transform(self.X_bar)
        else:
            Xm = X
            Xm_new = X_new
            self.Xm_bar = self.X_bar

        K = self.GlobalKernel(Xm, Xm)
        K_star = self.GlobalKernel(Xm_new, Xm)
        K_star_star = self.GlobalKernel(Xm_new, Xm_new)

        dK = K.diagonal()
        dK += self.global_gp_noise_std**2
        L = torch.linalg.cholesky(K)
        alpha = torch.cholesky_solve(y, L)

        pred_mean = K_star@alpha

        v = torch.cholesky_solve(K_star.T, L)
        pred_var = K_star_star - K_star@v

        dpred_var = pred_var.diagonal()
        dpred_var += self.global_gp_noise_std**2
        return pred_mean, pred_var