jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def fasterquant(
        self, blocksize=128, percdamp=.1, groupsize=-1, clip=False, baseline=False
    ):
        W = self.layer.weight.data.clone()
        W = W.float()

        tick = time.time()

        if self.stable:
            self.H /= self.nsamples
            self.H += self.mean.matmul(self.mean.t())
            self.H *= 2
        H = self.H
        del self.H

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

        if not baseline:
            try:
                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
            except:
                print('Singularity.')
                baseline = True
        if baseline:
            del H
            Hinv = torch.eye(self.columns, device=self.dev)

        if groupsize == -1:
            self.quantizer.find_params(W)
        groups = []

        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 groupsize != -1:
                self.quantizer.find_params(W1, solve=Hinv1 if clip else None)
                groups.append(copy.deepcopy(self.quantizer))

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

                q = quantize(
                    w.unsqueeze(1), self.quantizer.scale, self.quantizer.zero, self.quantizer.maxq
                ).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()
        print('time %.2f' % (time.time() - tick))
        print('error', torch.sum(Losses).item())

        self.layer.weight.data = Q.reshape(self.layer.weight.shape).to(self.layer.weight.data.dtype)
        if DEBUG:
            print(torch.sum((self.layer(self.inp1) - self.out1) ** 2))

        if groups:
            scale = torch.cat([q.scale for q in groups], dim=1)
            zero = torch.cat([q.zero for q in groups], dim=1)
            return scale, zero
        return self.quantizer.scale, self.quantizer.zero
//  [python]
def fasterquant(
        self, blocksize=128, percdamp=.1, groupsize=-1, clip=False, baseline=False
    ):
        W = self.layer.weight.data.clone()
        W = W.float()

        tick = time.time()

        if self.stable:
            self.H /= self.nsamples
            self.H += self.mean.matmul(self.mean.t())
            self.H *= 2
        H = self.H
        del self.H

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

        if not baseline:
            try:
                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
            except:
                print('Singularity.')
                baseline = True
        if baseline:
            del H
            Hinv = torch.eye(self.columns, device=self.dev)

        if groupsize == -1:
            self.quantizer.find_params(W)
        groups = []

        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 groupsize != -1:
                self.quantizer.find_params(W1, solve=Hinv1 if clip else None)
                groups.append(copy.deepcopy(self.quantizer))

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

                q = quantize(
                    w.unsqueeze(1), self.quantizer.scale, self.quantizer.zero, self.quantizer.maxq
                ).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()
        print('time %.2f' % (time.time() - tick))
        print('error', torch.sum(Losses).item())

        self.layer.weight.data = Q.reshape(self.layer.weight.shape).to(self.layer.weight.data.dtype)
        if DEBUG:
            print(torch.sum((self.layer(self.inp1) - self.out1) ** 2))

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

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

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

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

#        S = A_hat

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

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

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

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

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

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

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

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

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

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

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

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

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

#        S = A_hat

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

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

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

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

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

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

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

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

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

        return L_3_sum_term, KL_term
//  [python]
def __call__(self) -> Tensor:
        coalescent = self.coalescent.distribution()
        gamma = self.gmrf.field.tensor
        sufficient_statistics, coalescent_counts = coalescent.sufficient_statistics(
            self.coalescent.tree_model.node_heights
        )
        self.gmrf.precision.tensor = self.propose_precision()
        proposed_precision_matrix = self.gmrf.precision_matrix()

        if gamma.dim() > 1:
            mode_forward = []
            for i in range(gamma.shape[-2]):
                mode = self.newton_raphson(
                    coalescent_counts[i],
                    sufficient_statistics[i],
                    gamma[i],
                    proposed_precision_matrix[i],
                )
                mode_forward.append(mode)
            mode_forward = torch.stack(mode_forward)
        else:
            mode_forward = self.newton_raphson(
                coalescent_counts,
                sufficient_statistics,
                gamma,
                proposed_precision_matrix,
            )

        forwardQW = proposed_precision_matrix.clone()

        dim = gamma.shape[-1]

        diagonal1 = sufficient_statistics * torch.exp(-mode_forward)
        forwardQW[..., range(dim), range(dim)] += diagonal1
        diagonal1 = diagonal1 * (mode_forward + 1) - coalescent_counts

        z = torch.randn(gamma.shape)

        cholesky = torch.cholesky(forwardQW, upper=True)

        v = torch.triangular_solve(diagonal1, cholesky.transpose(-1, -2), upper=True).solution
        mu = torch.triangular_solve(v, cholesky, upper=True).solution
        u = torch.triangular_solve(z, cholesky, upper=True).solution
        proposed_gamma = mu + u

        self.gmrf.field.tensor = proposed_gamma.clone()

        diagonal = cholesky[..., range(dim), range(dim)]
        diagonal = torch.where(diagonal > 0.0000001, diagonal.log(), torch.zeros([1]))

        log_q_forward = diagonal.sum(-1) - 0.5 * (z * z).sum(-1)
        return log_q_forward
//  [python]
def __call__(self) -> Tensor:
        coalescent = self.coalescent.distribution()
        gamma = self.gmrf.field.tensor
        sufficient_statistics, coalescent_counts = coalescent.sufficient_statistics(
            self.coalescent.tree_model.node_heights
        )
        self.gmrf.precision.tensor = self.propose_precision()
        proposed_precision_matrix = self.gmrf.precision_matrix()

        if gamma.dim() > 1:
            mode_forward = []
            for i in range(gamma.shape[-2]):
                mode = self.newton_raphson(
                    coalescent_counts[i],
                    sufficient_statistics[i],
                    gamma[i],
                    proposed_precision_matrix[i],
                )
                mode_forward.append(mode)
            mode_forward = torch.stack(mode_forward)
        else:
            mode_forward = self.newton_raphson(
                coalescent_counts,
                sufficient_statistics,
                gamma,
                proposed_precision_matrix,
            )

        forwardQW = proposed_precision_matrix.clone()

        dim = gamma.shape[-1]

        diagonal1 = sufficient_statistics * torch.exp(-mode_forward)
        forwardQW[..., range(dim), range(dim)] += diagonal1
        diagonal1 = diagonal1 * (mode_forward + 1) - coalescent_counts

        z = torch.randn(gamma.shape)

        cholesky = torch.linalg.cholesky(forwardQW, upper=True)

        v = torch.linalg.solve(cholesky.transpose(-1, -2), diagonal1)
        mu = torch.linalg.solve(cholesky, v)
        u = torch.linalg.solve(cholesky, z)
        proposed_gamma = mu + u

        self.gmrf.field.tensor = proposed_gamma.clone()

        diagonal = cholesky[..., range(dim), range(dim)]
        diagonal = torch.where(diagonal > 0.0000001, diagonal.log(), torch.zeros([1]))

        log_q_forward = diagonal.sum(-1) - 0.5 * (z * z).sum(-1)
        return log_q_forward
//  [python]
def cholesky(self):
        """
        Computes the Cholesky decomposition of a given Kronecker-factorized matrix.

        self: `TTMatrix` containing a matrix or a
            batch of matrices of size N x N, factorized into a Kronecker product of
            square matrices (all tt-ranks are 1 and all tt-cores are square). All the
            cores must be symmetric positive-definite.

        :return: `TTMatrix` of size N x N
        """
        self._check_kron_properties()

        cho_cores = []
        for core_idx in range(self.d):
            if self.batch:
                core_cho = torch.cholesky(self.cores[core_idx][:, 0, :, :, 0])
                core_cho = torch.unsqueeze(core_cho, 1)
            else:
                core_cho = torch.cholesky(self.cores[core_idx][0, :, :, 0])
                core_cho = torch.unsqueeze(core_cho, 0)
            cho_cores.append(torch.unsqueeze(core_cho, -1))

        # NOTE: ranks will be computed based on cores shape
        return TTMatrix(cho_cores, None, self.input_dims, self.output_dims)
//  [python]
def cholesky(self):
        """
        Computes the Cholesky decomposition of a given Kronecker-factorized matrix.

        self: `TTMatrix` containing a matrix or a
            batch of matrices of size N x N, factorized into a Kronecker product of
            square matrices (all tt-ranks are 1 and all tt-cores are square). All the
            cores must be symmetric positive-definite.

        :return: `TTMatrix` of size N x N
        """
        self._check_kron_properties()

        cho_cores = []
        for core_idx in range(self.d):
            if self.batch:
                core_cho = torch.linalg.cholesky(self.cores[core_idx][:, 0, :, :, 0])
                core_cho = torch.unsqueeze(core_cho, 1)
            else:
                core_cho = torch.linalg.cholesky(self.cores[core_idx][0, :, :, 0])
                core_cho = torch.unsqueeze(core_cho, 0)
            core_cho.append(torch.unsqueeze(core_cho, -1))

        # NOTE: ranks will be computed based on cores shape
        return TTMatrix(cho_cores, None, self.input_dims, self.output_dims)
//  [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)
        # print('k', k)
        # 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)
        # print('c', c)
        # torch.clamp_(self.local_ls[:, dim], min=10**-4)
        # print(self.local_ls[dim].device, self.local_ls[dim].shape)
        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)
        # print('L0', l[0])

        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
            # k_post_det = torch.det(k_post)
            # k_post_det = torch.clamp(k_post_det, min=10**-20)
            # if k_post_det<=0:
            #     k_post_det = torch.tensor(10**-20)
            # B = torch.log(k_post_det).reshape(-1,1)

            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)
        # print('k', k)
        # 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)
        # print('c', c)
        # torch.clamp_(self.local_ls[:, dim], min=10**-4)
        # print(self.local_ls[dim].device, self.local_ls[dim].shape)
        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)
        # print('L0', l[0])

        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
            # k_post_det = torch.det(k_post)
            # k_post_det = torch.clamp(k_post_det, min=10**-20)
            # if k_post_det<=0:
            #     k_post_det = torch.tensor(10**-20)
            # B = torch.log(k_post_det).reshape(-1,1)

            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 setup_user_vals(self):
        self.N_x,self.m = self.X_U.shape
        self.reg = self.N_x*self.lamb
        self.eye = torch.eye(self.N_x).to(self.device)*self.reg
        L = torch.cholesky(self.k_U(self.X_U)+ self.eye)
        self.precond = torch.cholesky_inverse(L)

        self.setup_the_rest()
//  [python]
def setup_user_vals(self):
        self.N_x,self.m = self.X_U.shape
        self.reg = self.N_x*self.lamb
        self.eye = torch.eye(self.N_x).to(self.device)*self.reg
        L = torch.linalg.cholesky(self.k_U(self.X_U)+ self.eye)
        self.precond = torch.cholesky_inverse(L)

        self.setup_the_rest()
//  [python]
def OLS_solve(Y_target,Z_in,weights,big_weight=1e5):
    mask = Z_in.sum(0)>0
    shapley_vals = torch.zeros(Z_in.shape[1],Y_target.shape[1])
    weights[0] = big_weight
    weights[-1] = big_weight
    Z = Z_in[:,mask]
    b = Z.t()@Y_target
    A = Z.t()@(Z*weights)
    L = torch.cholesky(A)
    sol = torch.cholesky_solve(b,L)
    shapley_vals[mask,:] = sol
    return shapley_vals
//  [python]
def OLS_solve(Y_target,Z_in,weights,big_weight=1e5):
    mask = Z_in.sum(0)>0
    shapley_vals = torch.zeros(Z_in.shape[1],Y_target.shape[1])
    weights[0] = big_weight
    weights[-1] = big_weight
    Z = Z_in[:,mask]
    b = Z.t()@Y_target
    A = Z.t()@(Z*weights)
    L = torch.linalg.cholesky(A)
    sol = torch.cholesky_solve(b,L)
    shapley_vals[mask,:] = sol
    return shapley_vals
//  [python]
def _exact_log_det_jac_and_reconstruction(self, latent):
        eps = 1e-6
        EPS_FACTOR = 10
        MAX_ATTEMPTS = 6

        batch_size = latent.shape[0]
        latent_dim = latent.shape[1]

        jacobian_t_jacobian, reconstructed_x = self._get_full_jac_transpose_jac(latent, self.training)

        jac_attempts = 1
        eye = torch.eye(
            latent_dim,
            device=jacobian_t_jacobian.device,
            dtype=jacobian_t_jacobian.dtype
            ).view(1, latent_dim, latent_dim)
        eye_repeat = eye.repeat((batch_size, 1, 1))

        while True:
            try:
                cholesky_factor = torch.cholesky(jacobian_t_jacobian)
                break
            except RuntimeError:
                # HACK: If we end up running into non-invertibility, add eps*I to JtJ
                jacobian_t_jacobian = jacobian_t_jacobian + eps*eye_repeat
                jac_attempts += 1
                eps *= EPS_FACTOR

        if jac_attempts > 1:
            print(f"WARNING: Numerical non-invertibility in JtJ observed - {jac_attempts} attempts needed to fix")

        cholesky_diagonal = torch.diagonal(cholesky_factor, dim1=-2, dim2=-1)
        log_det_jacobian = 2 * torch.sum(torch.log(cholesky_diagonal), dim=1, keepdim=True)

        return log_det_jacobian, reconstructed_x
//  [python]
def _exact_log_det_jac_and_reconstruction(self, latent):
        eps = 1e-6
        EPS_FACTOR = 10
        MAX_ATTEMPTS = 6

        batch_size = latent.shape[0]
        latent_dim = latent.shape[1]

        jacobian_t_jacobian, reconstructed_x = self._get_full_jac_transpose_jac(latent, self.training)

        jac_attempts = 1
        eye = torch.eye(
            latent_dim,
            device=jacobian_t_jacobian.device,
            dtype=jacobian_t_jacobian.dtype
            ).view(1, latent_dim, latent_dim)
        eye_repeat = eye.repeat((batch_size, 1, 1))

        while True:
            try:
                cholesky_factor = torch.linalg.cholesky(jacobian_t_jacobian)
                break
            except RuntimeError:
                # HACK: If we end up running into non-invertibility, add eps*I to JtJ
                jacobian_t_jacobian = jacobian_t_jacobian + eps*eye_repeat
                jac_attempts += 1
                eps *= EPS_FACTOR

        if jac_attempts > 1:
            print(f"WARNING: Numerical non-invertibility in JtJ observed - {jac_attempts} attempts needed to fix")

        cholesky_diagonal = torch.diagonal(cholesky_factor, dim1=-2, dim2=-1)
        log_det_jacobian = 2 * torch.sum(torch.log(cholesky_diagonal), dim=1, keepdim=True)

        return log_det_jacobian, reconstructed_x
//  [python]
x = torch.tensor([[ 2.3999,  2.2497, -1.7825],
        [ 2.2497,  3.1641, -3.4624],
        [-1.7825, -3.4624,  4.5418]])

#torch_model.eval()

#torch_out = torch_model(x)

a = torch.randn(3, 3, dtype=torch.float32)
a = a @ a.mT + 1e-3
print(a)
l = torch.cholesky(a)
//  [python]
x = torch.tensor([[ 2.3999,  2.2497, -1.7825],
        [ 2.2497,  3.1641, -3.4624],
        [-1.7825, -3.4624,  4.5418]])

#torch_model.eval()

#torch_out = torch_model(x)

a = torch.randn(3, 3, dtype=torch.float32)
a = a @ a.mT + 1e-3
print(a)
l = torch.linalg.cholesky(a)
//  [python]
def test_multivariate_normal_properties(self):
        loc = torch.randn(5)
        scale_tril = transform_to(constraints.lower_cholesky)(torch.randn(5, 5))
        m = MultivariateNormal(loc=loc, scale_tril=scale_tril)
        self.assertEqual(m.covariance_matrix, m.scale_tril.mm(m.scale_tril.t()))
        self.assertEqual(
            m.covariance_matrix.mm(m.precision_matrix), torch.eye(m.event_shape[0])
        )
        self.assertEqual(m.scale_tril, torch.cholesky(m.covariance_matrix))
//  [python]
def test_multivariate_normal_properties(self):
        loc = torch.randn(5)
        scale_tril = transform_to(constraints.lower_cholesky)(torch.randn(5, 5))
        m = MultivariateNormal(loc=loc, scale_tril=scale_tril)
        self.assertEqual(m.covariance_matrix, m.scale_tril.mm(m.scale_tril.t()))
        self.assertEqual(
            m.covariance_matrix.mm(m.precision_matrix), torch.eye(m.event_shape[0])
        )
        self.assertEqual(m.scale_tril, torch.linalg.cholesky(m.covariance_matrix))
//  [python]
def linear_gaussian(
    theta: Tensor,
    likelihood_shift: Tensor,
    likelihood_cov: Tensor,
    num_discarded_dims: int = 0,
) -> Tensor:
    """
    Simulator for linear Gaussian.

    Uses Cholesky decomposition to transform samples from standard Gaussian.

    When `num_discarded_dims>0`, return simulation outputs with as many last dimensions
    discarded. This is implemented by throwing away the last `num_discarded_dims`
    dimensions of theta and then running the linear Gaussian as always.

    Args:
        theta: Parameter sets to be simulated.
        likelihood_shift: The simulator will shift the value of theta by this value.
            Thus, the mean of the Gaussian likelihood will be likelihood_shift+theta.
        likelihood_cov: Covariance matrix of the likelihood.
        num_discarded_dims: Number of output dimensions to discard.

    Returns: Simulated data.
    """

    if num_discarded_dims:
        theta = theta[:, :-num_discarded_dims]

    chol_factor = torch.cholesky(likelihood_cov)

    return likelihood_shift + theta + torch.mm(chol_factor, torch.randn_like(theta).T).T
//  [python]
def linear_gaussian(
    theta: Tensor,
    likelihood_shift: Tensor,
    likelihood_cov: Tensor,
    num_discarded_dims: int = 0,
) -> Tensor:
    """
    Simulator for linear Gaussian.

    Uses Cholesky decomposition to transform samples from standard Gaussian.

    When `num_discarded_dims>0`, return simulation outputs with as many last dimensions
    discarded. This is implemented by throwing away the last `num_discarded_dims`
    dimensions of theta and then running the linear Gaussian as always.

    Args:
        theta: Parameter sets to be simulated.
        likelihood_shift: The simulator will shift the value of theta by this value.
            Thus, the mean of the Gaussian likelihood will be likelihood_shift+theta.
        likelihood_cov: Covariance matrix of the likelihood.
        num_discarded_dims: Number of output dimensions to discard.

    Returns: Simulated data.
    """

    if num_discarded_dims:
        theta = theta[:, :-num_discarded_dims]

    chol_factor = torch.linalg.cholesky(likelihood_cov)

    return likelihood_shift + theta + torch.mm(chol_factor, torch.randn_like(theta).T).T
//  [python]
def batch_gptq(
  W, H, quantizer, blocksize=128, percdamp=.1, groupsize=-1, actorder=False
):
  dtype = W.dtype
  W = W.clone()
  W = W.float()

  rows, columns = W.shape[1:]
  dev = W.device

  quantizer.find_params(W)

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

  diag = torch.arange(columns, device=dev)
  damp = percdamp * torch.mean(H[:, diag, diag], axis=-1, keepdim=True)
  damp = torch.maximum(damp, 1e-6 * torch.ones_like(damp)) # catch all zeros
  H[:, diag, diag] += damp

  if actorder:
    perm = torch.argsort(H[:, diag, diag], dim=1, descending=True)
    for i in range(W.shape[0]):
      W[i] = W[i, :, perm[i]]
      H[i] = H[i][perm[i]][:, perm[i]]
    invperm = torch.argsort(perm, dim=1)

  err = True
  while err:
    # We need to loop as batch operations only return the first error
    try:
      H1 = torch.cholesky(H)
      H1 = torch.cholesky_inverse(H1)
      H1 = torch.cholesky(H1, upper=True)
      H = H1
      err = False
    except RuntimeError as ex:
      print('Skip due to singularity.')
      idx = int(str(ex).replace('linalg.cholesky: (Batch element ', '').split('):')[0])
      # Do RTN for failed Hessians by turning them into identity
      H[idx] = torch.eye(columns, device=dev)
  Hinv = H

  for i1 in range(0, columns, blocksize):
    i2 = min(i1 + blocksize, 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].unsqueeze(1)

      if groupsize != -1:
        if (i1 + i) % groupsize == 0:
          quantizer.find_params(W[:, :, (i1 + i):(i1 + i + groupsize)])

      q = quantize(
        w.unsqueeze(2), quantizer.scale, quantizer.zero, quantizer.maxq
      ).flatten(1)
      Q1[:, :, i] = q
      Losses1[:, :, i] = (w - q) ** 2 / d ** 2
      err1 = (w - q) / d
      W1[:, :, i:] -= torch.bmm(err1.unsqueeze(2), Hinv1[:, i, i:].unsqueeze(1))
      Err1[:, :, i] = err1

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

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

  torch.cuda.synchronize(device=dev)
  print('error', torch.sum(Losses.flatten(1), 1))
  print('Sparsity:', torch.mean((Q == 0).float()))

  if actorder:
    for i in range(W.shape[0]):
      Q[i] = Q[i, :, invperm[i]]

  return Q.to(dtype)
//  [python]
def batch_gptq(
  W, H, quantizer, blocksize=128, percdamp=.1, groupsize=-1, actorder=False
):
  dtype = W.dtype
  W = W.clone()
  W = W.float()

  rows, columns = W.shape[1:]
  dev = W.device

  quantizer.find_params(W)

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

  diag = torch.arange(columns, device=dev)
  damp = percdamp * torch.mean(H[:, diag, diag], axis=-1, keepdim=True)
  damp = torch.maximum(damp, 1e-6 * torch.ones_like(damp)) # catch all zeros
  H[:, diag, diag] += damp

  if actorder:
    perm = torch.argsort(H[:, diag, diag], dim=1, descending=True)
    for i in range(W.shape[0]):
      W[i] = W[i, :, perm[i]]
      H[i] = H[i][perm[i]][:, perm[i]]
    invperm = torch.argsort(perm, dim=1)

  err = True
  while err:
    # We need to loop as batch operations only return the first error
    try:
      H1 = torch.linalg.cholesky(H)
      H1 = torch.cholesky_inverse(H1)
      H1 = torch.linalg.cholesky(H1, upper=True)
      H = H1
      err = False
    except RuntimeError as ex:
      print('Skip due to singularity.')
      idx = int(str(ex).replace('linalg.cholesky: (Batch element ', '').split('):')[0])
      # Do RTN for failed Hessians by turning them into identity
      H[idx] = torch.eye(columns, device=dev)
  Hinv = H

  for i1 in range(0, columns, blocksize):
    i2 = min(i1 + blocksize, 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].unsqueeze(1)

      if groupsize != -1:
        if (i1 + i) % groupsize == 0:
          quantizer.find_params(W[:, :, (i1 + i):(i1 + i + groupsize)])

      q = quantize(
        w.unsqueeze(2), quantizer.scale, quantizer.zero, quantizer.maxq
      ).flatten(1)
      Q1[:, :, i] = q
      Losses1[:, :, i] = (w - q) ** 2 / d ** 2
      err1 = (w - q) / d
      W1[:, :, i:] -= torch.bmm(err1.unsqueeze(2), Hinv1[:, i, i:].unsqueeze(1))
      Err1[:, :, i] = err1

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

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

  torch.cuda.synchronize(device=dev)
  print('error', torch.sum(Losses.flatten(1), 1))
  print('Sparsity:', torch.mean((Q == 0).float()))

  if actorder:
    for i in range(W.shape[0]):
      Q[i] = Q[i, :, invperm[i]]

  return Q.to(dtype)
//  [python]
def forward(self, x_train, y_train, x_test, Kinv_method='cholesky3'):
        """
        Computes the mean and covariance of the Gaussian process.

        Args:
            x_train (torch.Tensor): The input training data.
            y_train (torch.Tensor): The target training data.
            x_test (torch.Tensor): The input test data.
            Kinv_method (str, optional): The method used for computing the inverse of the covariance matrix. 
                Defaults to 'cholesky3'.

        Returns:
            tuple: A tuple containing the mean and covariance of the Gaussian process.

        Raises:
            ValueError: If Kinv_method is not 'direct' or 'cholesky'.

        """
        K = self.kernel(x_train, x_train) + self.noise_variance.pow(2) * torch.eye(len(x_train))
        K_s = self.kernel(x_train, x_test)
        K_ss = self.kernel(x_test, x_test)
        
        if Kinv_method == 'cholesky1':
            # kernel inverse is not stable, use cholesky decomposition instead
            L = torch.cholesky(K)
            L_inv = torch.inverse(L)
            K_inv = L_inv.T @ L_inv
            alpha = K_inv @ y_train
            mu = K_s.T @ alpha
            v = L_inv @ K_s
            cov = K_ss - v.T @ v
        elif Kinv_method == 'cholesky3':
            # recommended implementation, fastest so far
            L = torch.cholesky(K)
            alpha = torch.cholesky_solve(y_train, L)
            mu = K_s.T @ alpha
            v = L.inverse() @ K_s
            cov = K_ss - v.T @ v
        elif Kinv_method == 'direct':
            K_inv = torch.inverse(K)
            mu = K_s.T @ K_inv @ y_train
            cov = K_ss - K_s.T @ K_inv @ K_s
        else:
            raise ValueError('Kinv_method should be either direct or cholesky')
        return mu.squeeze(), cov
//  [python]
def forward(self, x_train, y_train, x_test, Kinv_method='cholesky3'):
        """
        Computes the mean and covariance of the Gaussian process.

        Args:
            x_train (torch.Tensor): The input training data.
            y_train (torch.Tensor): The target training data.
            x_test (torch.Tensor): The input test data.
            Kinv_method (str, optional): The method used for computing the inverse of the covariance matrix. 
                Defaults to 'cholesky3'.

        Returns:
            tuple: A tuple containing the mean and covariance of the Gaussian process.

        Raises:
            ValueError: If Kinv_method is not 'direct' or 'cholesky'.

        """
        K = self.kernel(x_train, x_train) + self.noise_variance.pow(2) * torch.eye(len(x_train))
        K_s = self.kernel(x_train, x_test)
        K_ss = self.kernel(x_test, x_test)
        
        if Kinv_method == 'cholesky1':
            # kernel inverse is not stable, use cholesky decomposition instead
            L = torch.linalg.cholesky(K)
            L_inv = torch.inverse(L)
            K_inv = L_inv.T @ L_inv
            alpha = K_inv @ y_train
            mu = K_s.T @ alpha
            v = L_inv @ K_s
            cov = K_ss - v.T @ v
        elif Kinv_method == 'cholesky3':
            # recommended implementation, fastest so far
            L = torch.linalg.cholesky(K)
            alpha = torch.cholesky_solve(y_train, L)
            mu = K_s.T @ alpha
            v = L.inverse() @ K_s
            cov = K_ss - v.T @ v
        elif Kinv_method == 'direct':
            K_inv = torch.inverse(K)
            mu = K_s.T @ K_inv @ y_train
            cov = K_ss - K_s.T @ K_inv @ K_s
        else:
            raise ValueError('Kinv_method should be either direct or cholesky')
        return mu.squeeze(), cov
//  [python]
def log_likelihood(self, x_train, y_train, Kinv_method='cholesky3'):
        """
        Computes the log-likelihood of the Gaussian process.

        Args:
            x_train (torch.Tensor): The input training data.
            y_train (torch.Tensor): The target training data.
            Kinv_method (str, optional): The method used for computing the inverse of the covariance matrix. 
                Defaults to 'cholesky3'.

        Returns:
            torch.Tensor: The log-likelihood of the Gaussian process.

        Raises:
            ValueError: If Kinv_method is not 'direct' or 'cholesky'.

        """
        K = self.kernel(x_train, x_train) + self.noise_variance.pow(2) * torch.eye(len(x_train))
        
        if Kinv_method == 'cholesky1':
            L = torch.cholesky(K)
            L_inv = torch.inverse(L)
            K_inv = L_inv.T @ L_inv
            return -0.5 * (y_train.T @ K_inv @ y_train + torch.logdet(K) + len(x_train) * np.log(2 * np.pi))
        elif Kinv_method == 'cholesky2':
            L = torch.cholesky(K)
            gamma = torch.cholesky_solve(y_train, L)
            return -0.5 * (gamma.T @ gamma + torch.logdet(K) + len(x_train) * np.log(2 * np.pi))
        elif Kinv_method == 'cholesky3':
            L = torch.cholesky(K)
            gamma = torch.cholesky_solve(y_train, L)
            return -0.5 * (gamma.T @ gamma + L.diag().log().sum() + len(x_train) * np.log(2 * np.pi))
        elif Kinv_method == 'direct':
            K_inv = torch.inverse(K)
            return -0.5 * (y_train.T @ K_inv @ y_train + torch.logdet(K) + len(x_train) * np.log(2 * np.pi))
        elif Kinv_method == 'torch_distribution_MN1':
            L = torch.cholesky(K)
            return torch.distributions.MultivariateNormal(y_train, scale_tril=L).log_prob(y_train)
        elif Kinv_method == 'torch_distribution_MN2':
            return torch.distributions.MultivariateNormal(y_train, K).log_prob(y_train)
        else:
            raise ValueError('Kinv_method should be either direct or cholesky')
//  [python]
def log_likelihood(self, x_train, y_train, Kinv_method='cholesky3'):
        """
        Computes the log-likelihood of the Gaussian process.

        Args:
            x_train (torch.Tensor): The input training data.
            y_train (torch.Tensor): The target training data.
            Kinv_method (str, optional): The method used for computing the inverse of the covariance matrix. 
                Defaults to 'cholesky3'.

        Returns:
            torch.Tensor: The log-likelihood of the Gaussian process.

        Raises:
            ValueError: If Kinv_method is not 'direct' or 'cholesky'.

        """
        K = self.kernel(x_train, x_train) + self.noise_variance.pow(2) * torch.eye(len(x_train))
        
        if Kinv_method == 'cholesky1':
            L = torch.linalg.cholesky(K)
            L_inv = torch.inverse(L)
            K_inv = L_inv.T @ L_inv
            return -0.5 * (y_train.T @ K_inv @ y_train + torch.logdet(K) + len(x_train) * np.log(2 * np.pi))
        elif Kinv_method == 'cholesky2':
            L = torch.linalg.cholesky(K)
            gamma = torch.cholesky_solve(y_train, L)
            return -0.5 * (gamma.T @ gamma + torch.logdet(K) + len(x_train) * np.log(2 * np.pi))
        elif Kinv_method == 'cholesky3':
            L = torch.linalg.cholesky(K)
            gamma = torch.cholesky_solve(y_train, L)
            return -0.5 * (gamma.T @ gamma + L.diag().log().sum() + len(x_train) * np.log(2 * np.pi))
        elif Kinv_method == 'direct':
            K_inv = torch.inverse(K)
            return -0.5 * (y_train.T @ K_inv @ y_train + torch.logdet(K) + len(x_train) * np.log(2 * np.pi))
        elif Kinv_method == 'torch_distribution_MN1':
            L = torch.cholesky(K)
            return torch.distributions.MultivariateNormal(y_train, scale_tril=L).log_prob(y_train)
        elif Kinv_method == 'torch_distribution_MN2':
            return torch.distributions.MultivariateNormal(y_train, K).log_prob(y_train)
        else:
            raise ValueError('Kinv_method should be either direct or cholesky')
//  [python]
def test_wishart_properties(self):
        set_rng_seed(0)  # see Note [Randomized statistical tests]
        ndim = 5
        df = torch.rand([]) + ndim - 1
        scale_tril = transform_to(constraints.lower_cholesky)(torch.randn(ndim, ndim))
        m = Wishart(df=df, scale_tril=scale_tril)
        self.assertEqual(m.covariance_matrix, m.scale_tril.mm(m.scale_tril.t()))
        self.assertEqual(
            m.covariance_matrix.mm(m.precision_matrix), torch.eye(m.event_shape[0])
        )
        self.assertEqual(m.scale_tril, torch.cholesky(m.covariance_matrix))
//  [python]
def test_wishart_properties(self):
        set_rng_seed(0)  # see Note [Randomized statistical tests]
        ndim = 5
        df = torch.rand([]) + ndim - 1
        scale_tril = transform_to(constraints.lower_cholesky)(torch.randn(ndim, ndim))
        m = Wishart(df=df, scale_tril=scale_tril)
        self.assertEqual(m.covariance_matrix, m.scale_tril.mm(m.scale_tril.t()))
        self.assertEqual(
            m.covariance_matrix.mm(m.precision_matrix), torch.eye(m.event_shape[0])
        )
        self.assertEqual(m.scale_tril, torch.linalg.cholesky(m.covariance_matrix))
//  [python]
def test_wishart_sample(self):
        set_rng_seed(0)  # see Note [Randomized statistical tests]
        ndim = 3
        df = torch.rand([], requires_grad=True) + ndim - 1
        # SciPy allowed ndim -1 < df < ndim for Wishar distribution after version 1.7.0
        if version.parse(scipy.__version__) < version.parse("1.7.0"):
            df += 1.0
        tmp = torch.randn(ndim, 10)
        cov = (torch.matmul(tmp, tmp.t()) / tmp.size(-1)).requires_grad_()
        prec = cov.inverse().requires_grad_()
        scale_tril = torch.cholesky(cov).requires_grad_()

        ref_dist = scipy.stats.wishart(df.item(), cov.detach().numpy())

        self._check_sampler_sampler(
            Wishart(df, cov),
            ref_dist,
            f"Wishart(df={df}, covariance_matrix={cov})",
            multivariate=True,
        )
        self._check_sampler_sampler(
            Wishart(df, precision_matrix=prec),
            ref_dist,
            f"Wishart(df={df}, precision_matrix={prec})",
            multivariate=True,
        )
        self._check_sampler_sampler(
            Wishart(df, scale_tril=scale_tril),
            ref_dist,
            f"Wishart(df={df}, scale_tril={scale_tril})",
            multivariate=True,
        )
//  [python]
def test_wishart_sample(self):
        set_rng_seed(0)  # see Note [Randomized statistical tests]
        ndim = 3
        df = torch.rand([], requires_grad=True) + ndim - 1
        # SciPy allowed ndim -1 < df < ndim for Wishar distribution after version 1.7.0
        if version.parse(scipy.__version__) < version.parse("1.7.0"):
            df += 1.0
        tmp = torch.randn(ndim, 10)
        cov = (torch.matmul(tmp, tmp.t()) / tmp.size(-1)).requires_grad_()
        prec = cov.inverse().requires_grad_()
        scale_tril = torch.linalg.cholesky(cov).requires_grad_()

        ref_dist = scipy.stats.wishart(df.item(), cov.detach().numpy())

        self._check_sampler_sampler(
            Wishart(df, cov),
            ref_dist,
            f"Wishart(df={df}, covariance_matrix={cov})",
            multivariate=True,
        )
        self._check_sampler_sampler(
            Wishart(df, precision_matrix=prec),
            ref_dist,
            f"Wishart(df={df}, precision_matrix={prec})",
            multivariate=True,
        )
        self._check_sampler_sampler(
            Wishart(df, scale_tril=scale_tril),
            ref_dist,
            f"Wishart(df={df}, scale_tril={scale_tril})",
            multivariate=True,
        )
//  [python]
def _fit_params(self) -> None:
        self.mus = {}
        self.invs = {}
        self.precision_chols = {}
        device = next(self.model.parameters()).device
        for layer_name, layer_features in self.train_features.items():
            self.mus[layer_name], cov, self.invs[layer_name] = self.mu_cov_inv_est_fn(
                layer_features, self.train_targets, self.cov_mat_method, device=device
            )
            cov_chol = torch.cholesky(cov.to(device) + self.cov_reg * torch.eye(cov.shape[1], device=device))
            self.precision_chols[layer_name] = torch.triangular_solve(
                torch.eye(cov_chol.shape[1], device=device), cov_chol, upper=False
            ).solution.T
//  [python]
def _fit_params(self) -> None:
        self.mus = {}
        self.invs = {}
        self.precision_chols = {}
        device = next(self.model.parameters()).device
        for layer_name, layer_features in self.train_features.items():
            self.mus[layer_name], cov, self.invs[layer_name] = self.mu_cov_inv_est_fn(
                layer_features, self.train_targets, self.cov_mat_method, device=device
            )
            cov_chol = torch.linalg.cholesky(cov.to(device) + self.cov_reg * torch.eye(cov.shape[1], device=device))
            self.precision_chols[layer_name] = torch.linalg.solve_triangular(
                cov_chol, torch.eye(cov_chol.shape[1], device=device), upper=False
            ).T
//  [python]
def test_lkj_prior_log_prob(self, cuda=False):
        device = torch.device("cuda") if cuda else torch.device("cpu")
        prior = LKJPrior(2, torch.tensor(0.5, device=device))
        dist = LKJCholesky(2, torch.tensor(0.5, device=device))

        S = torch.eye(2, device=device)
        S_chol = torch.cholesky(S)
        self.assertAlmostEqual(prior.log_prob(S), dist.log_prob(S_chol), places=4)
        S = torch.stack([S, torch.tensor([[1.0, 0.5], [0.5, 1]], device=S.device)])
        S_chol = torch.cholesky(S)
        self.assertTrue(approx_equal(prior.log_prob(S), dist.log_prob(S_chol)))
        with self.assertRaises(ValueError):
            prior.log_prob(torch.eye(3, device=device))
//  [python]
def test_lkj_prior_log_prob(self, cuda=False):
        device = torch.device("cuda") if cuda else torch.device("cpu")
        prior = LKJPrior(2, torch.tensor(0.5, device=device))
        dist = LKJCholesky(2, torch.tensor(0.5, device=device))

        S = torch.eye(2, device=device)
        S_chol = torch.linalg.cholesky(S)
        self.assertAlmostEqual(prior.log_prob(S), dist.log_prob(S_chol), places=4)
        S = torch.stack([S, torch.tensor([[1.0, 0.5], [0.5, 1]], device=S.device)])
        S_chol = torch.linalg.cholesky(S)
        self.assertTrue(approx_equal(prior.log_prob(S), dist.log_prob(S_chol)))
        with self.assertRaises(ValueError):
            prior.log_prob(torch.eye(3, device=device))
//  [python]
def test_lkj_prior_batch_log_prob(self, cuda=False):
        device = torch.device("cuda") if cuda else torch.device("cpu")
        prior = LKJPrior(2, torch.tensor([0.5, 1.5], device=device))
        dist = LKJCholesky(2, torch.tensor([0.5, 1.5], device=device))

        S = torch.eye(2, device=device)
        S_chol = torch.cholesky(S)
        self.assertTrue(approx_equal(prior.log_prob(S), dist.log_prob(S_chol)))
        S = torch.stack([S, torch.tensor([[1.0, 0.5], [0.5, 1]], device=S.device)])
        S_chol = torch.cholesky(S)
        self.assertTrue(approx_equal(prior.log_prob(S), dist.log_prob(S_chol)))
        with self.assertRaises(ValueError):
            prior.log_prob(torch.eye(3, device=device))