jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def __init__(self, x, v, S, M, dt) :
        super(ProjectiveDynamics, self).__init__()

        self.x = x                              #initial position
        self.v = v                              #initial velocity
        self.M = M                              #mass matrix 
        self.S = S                              #selection matrix   
        self.e = S.shape[0]//4                  #number of tet elements 
        self.n = x.shape[0]                     #number of tet elements
        self.dt = dt                            #time step 

        Li, Ji_T = self.__compute_L_and_J()     #initialize constant variables

        self.Li = Li 
        self.Ji_T = Ji_T

        A = M + dt * dt * torch.sum(Li)         #Laplacian matrix: M + h*h*L
        self.U = torch.cholesky(A)
//  [python]
def __init__(self, x, v, S, M, dt) :
        super(ProjectiveDynamics, self).__init__()

        self.x = x                              #initial position
        self.v = v                              #initial velocity
        self.M = M                              #mass matrix 
        self.S = S                              #selection matrix   
        self.e = S.shape[0]//4                  #number of tet elements 
        self.n = x.shape[0]                     #number of tet elements
        self.dt = dt                            #time step 

        Li, Ji_T = self.__compute_L_and_J()     #initialize constant variables

        self.Li = Li 
        self.Ji_T = Ji_T

        A = M + dt * dt * torch.sum(Li)         #Laplacian matrix: M + h*h*L
        self.U = torch.linalg.cholesky(A)
//  [python]
def cholesky(matrix, rhs):
    chol = torch.cholesky(matrix)
    return torch.cholesky_solve(chol, rhs)
//  [python]
def cholesky(matrix, rhs):
    chol = torch.linalg.cholesky(matrix)
    return torch.linalg.cholesky_solve(chol, rhs)
//  [python]
def _cholesky_solve(A, b, eps=1e-8):
    A = _add_value_to_diagonal(A, eps)
    if A.ndim > b.ndim:
        b = b.unsqueeze(dim=-1)
    u = torch.cholesky(A)
    return torch.cholesky_solve(b, u).squeeze(dim=-1)
//  [python]
def _cholesky_solve(A, b, eps=1e-8):
    A = _add_value_to_diagonal(A, eps)
    if A.ndim > b.ndim:
        b = b.unsqueeze(dim=-1)
    u = torch.linalg.cholesky(A)
    return torch.linalg.cholesky_solve(b, u).squeeze(dim=-1)
//  [python]
def forward(ctx, q: torch.Tensor, p: torch.Tensor):
        l = torch.cholesky(q)
        d = torch.cholesky_solve(p, l)  # D = Q-1 @ P
        ctx.save_for_backward(l, d)
        return d
//  [python]
def forward(ctx, q: torch.Tensor, p: torch.Tensor):
        l = torch.linalg.cholesky(q)
        d = torch.cholesky_solve(p, l)  # D = Q-1 @ P
        ctx.save_for_backward(l, d)
        return d
//  [python]
def test_fast_cholesky2x2_against_torch(seed, rand_spd, n):
    x = rand_spd(n, 2)
    torch_time = timeit(lambda: torch.cholesky(x), number=10)
    fast_time = timeit(lambda: fast.cholesky2x2(x), number=10)
    logging.info('Fast cholesky2x2 speedup (n={}, d=2): {:.2f}x'.format(
            n, torch_time / fast_time))
    assert fast_time < torch_time
//  [python]
def test_fast_cholesky2x2_against_torch(seed, rand_spd, n):
    x = rand_spd(n, 2)
    torch_time = timeit(lambda: torch.linalg.cholesky(x), number=10)
    fast_time = timeit(lambda: fast.cholesky2x2(x), number=10)
    logging.info('Fast cholesky2x2 speedup (n={}, d=2): {:.2f}x'.format(
            n, torch_time / fast_time))
    assert fast_time < torch_time
//  [python]
def log_det_by_cholesky(matrix):
    """
    Args:
        matrix: matrix must be a positive define matrix.
                shape [N, C, D, D].
    Ref:
        https://github.com/tensorflow/tensorflow/blob/r1.13/tensorflow/python/ops/linalg/linalg_impl.py
    """
    # This uses the property that the log det(A) = 2 * sum(log(real(diag(C))))
    # where C is the cholesky decomposition of A.
    chol = torch.cholesky(matrix)
    # return 2.0 * torch.sum(torch.log(torch.diagonal(chol, dim1=-2, dim2=-1) + 1e-6), dim=-1)
    return 2.0 * torch.sum(torch.log(torch.diagonal(chol, dim1=-2, dim2=-1) + 1e-8), dim=-1)
//  [python]
def log_det_by_cholesky(matrix):
    """
    Args:
        matrix: matrix must be a positive define matrix.
                shape [N, C, D, D].
    Ref:
        https://github.com/tensorflow/tensorflow/blob/r1.13/tensorflow/python/ops/linalg/linalg_impl.py
    """
    # This uses the property that the log det(A) = 2 * sum(log(real(diag(C))))
    # where C is the cholesky decomposition of A.
    chol = torch.linalg.cholesky(matrix)
    # return 2.0 * torch.sum(torch.log(torch.diagonal(chol, dim1=-2, dim2=-1) + 1e-6), dim=-1)
    return 2.0 * torch.sum(torch.log(torch.diagonal(chol, dim1=-2, dim2=-1) + 1e-8), dim=-1)
//  [python]
def forward(self, x, x_var=0.):
        """
        Forward pass of the module.

        Args:
            x (torch.Tensor or list): Input tensor or list of tensors.
            x_var (float, optional): Variance of the input tensor. Defaults to 0.

        Returns:
            tuple: Tuple containing the mean and variance of the output.
        """
        x = self.check_single_tensor(x)

        if self.train_x is None:
            MFGP_LOG.e("gp model model hasn't been trained. predict failed")
            return None
        
        l1, h1, l2, h2 = self.l1, self.h1, self.l2, self.h2
        _noise = self.noise_box.get()

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

        kx = self.kernel(self.train_x, x, l1, h1, l2, h2)
        L = torch.cholesky(sigma)
        LinvKx,_ = torch.triangular_solve(kx, L, upper = False)

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

        return mean, var_diag
//  [python]
def forward(self, x, x_var=0.):
        """
        Forward pass of the module.

        Args:
            x (torch.Tensor or list): Input tensor or list of tensors.
            x_var (float, optional): Variance of the input tensor. Defaults to 0.

        Returns:
            tuple: Tuple containing the mean and variance of the output.
        """
        x = self.check_single_tensor(x)

        if self.train_x is None:
            MFGP_LOG.e("gp model model hasn't been trained. predict failed")
            return None
        
        l1, h1, l2, h2 = self.l1, self.h1, self.l2, self.h2
        _noise = self.noise_box.get()

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

        kx = self.kernel(self.train_x, x, l1, h1, l2, h2)
        L = torch.linalg.cholesky(sigma)
        LinvKx,_ = torch.triangular_solve(kx, L, upper = False)

        mean = kx.t() @ torch.linalg.solve(L, self.train_y, upper=False)  # torch.linalg.cholesky()
        
        var_diag = self.kernel(x, x, l1, h1, l2, h2).diag().view(-1, 1) \
            - (LinvKx**2).sum(dim = 0).view(-1, 1)
        # add the noise uncertainty
        var_diag = var_diag + _noise.pow(-1)

        return mean, var_diag
//  [python]
def transformation(covs, device, engine='symeig'):
    covs = covs.to(linalg_device)
    if engine == 'cholesky':
        C = torch.cholesky(covs.to(linalg_device))
        W = torch.triangular_solve(torch.eye(C.size(-1)).expand_as(C).to(C), C, upper=False)[0].transpose(1,2).to(x.device)
    else:
        if engine == 'symeig':
            
            S, U = torch.symeig(covs.to(linalg_device), eigenvectors=True, upper=True)
        elif engine == 'svd':
            U, S, _ = torch.svd(covs.to(linalg_device))
        elif engine == 'svd_lowrank':
            U, S, _ = torch.svd_lowrank(covs.to(linalg_device))
        elif engine == 'pca_lowrank':
            U, S, _ = torch.pca_lowrank(covs.to(linalg_device), center=False)
        S, U = S.to(device), U.to(device)
        W = U.bmm(S.rsqrt().diag_embed()).bmm(U.transpose(1,2))
    return W
//  [python]
def transformation(covs, device, engine='symeig'):
    covs = covs.to(linalg_device)
    if engine == 'cholesky':
        C = torch.linalg.cholesky(covs.to(linalg_device))
        W = torch.triangular_solve(torch.eye(C.size(-1)).expand_as(C).to(C), C, upper=False)[0].transpose(1,2).to(x.device)
    else:
        if engine == 'symeig':
            S, U = torch.symeig(covs.to(linalg_device), eigenvectors=True, upper=True)
        elif engine == 'svd':
            U, S, _ = torch.svd(covs.to(linalg_device))
        elif engine == 'svd_lowrank':
            U, S, _ = torch.svd_lowrank(covs.to(linalg_device))
        elif engine == 'pca_lowrank':
            U, S, _ = torch.pca_lowrank(covs.to(linalg_device), center=False)
        S, U = S.to(device), U.to(device)
        W = U.bmm(S.rsqrt().diag_embed()).bmm(U.transpose(1,2))
    return W
//  [python]
def test_batch_cg_with_tridiag(self):
        batch = 5
        size = 10
        matrix = torch.randn(batch, size, size, dtype=torch.float64)
        matrix = matrix.matmul(matrix.transpose(-1, -2))
        matrix.div_(matrix.norm())
        matrix.add_(torch.eye(matrix.size(-1), dtype=torch.float64).mul_(1e-1))

        rhs = torch.randn(batch, size, 10, dtype=torch.float64)
        solves, t_mats = linear_cg(
            matrix.matmul, rhs=rhs, n_tridiag=8, max_iter=size, max_tridiag_iter=10, tolerance=0, eps=1e-30
        )

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

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

        rhs = torch.randn(batch, size, 10, dtype=torch.float64)
        solves, t_mats = linear_cg(
            matrix.matmul, rhs=rhs, n_tridiag=8, max_iter=size, max_tridiag_iter=10, tolerance=0, eps=1e-30
        )

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

        # Check tridiag
        for i in range(5):
            eigs = matrix[i].symeig()[0]
            for j in range(8):
                approx_eigs = t_mats[j, i].symeig()[0]
                self.assertTrue(torch.allclose(eigs, approx_eigs, atol=1e-3, rtol=1e-4))
//  [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 quasi_sample(self, n, fix_input = None, factor = 16): 
        assert fix_input is None
        if self.clip:
            return self.eff_space.sample(n)

        B    = torch.FloatTensor(self.proj_matrix)
        L    = torch.cholesky(B.mm(B.t()))
        samp = pd.DataFrame(columns = self.eff_space.numeric_names)
        while samp.shape[0] < n:
            samp_hd = self.space.sample(100)
            alpha   = B.mm(torch.FloatTensor(samp_hd.values.T))
            samp_ld = pd.DataFrame(factor * torch.cholesky_solve(alpha, L).t().numpy(), columns = samp.columns)
            samp_pj = self.project(samp_ld)
            samp_ld = samp_ld[samp_pj.max(axis = 1) <= 1.0]
            samp_ld = samp_ld[samp_pj.min(axis = 1) >= -1.0]
            if samp_ld.shape[0] == samp_hd.shape[0]:
                factor /= 0.8
                continue
            elif samp_ld.shape[0] == 0:
                factor *= 0.8
            samp = samp.append(samp_ld, ignore_index = True)
        return samp.head(n)
//  [python]
def quasi_sample(self, n, fix_input = None, factor = 16): 
        assert fix_input is None
        if self.clip:
            return self.eff_space.sample(n)

        B    = torch.FloatTensor(self.proj_matrix)
        L    = torch.linalg.cholesky(B.mm(B.t()))
        samp = pd.DataFrame(columns = self.eff_space.numeric_names)
        while samp.shape[0] < n:
            samp_hd = self.space.sample(100)
            alpha   = B.mm(torch.FloatTensor(samp_hd.values.T))
            samp_ld = pd.DataFrame(factor * torch.cholesky_solve(alpha, L).t().numpy(), columns = samp.columns)
            samp_pj = self.project(samp_ld)
            samp_ld = samp_ld[samp_pj.max(axis = 1) <= 1.0]
            samp_ld = samp_ld[samp_pj.min(axis = 1) >= -1.0]
            if samp_ld.shape[0] == samp_hd.shape[0]:
                factor /= 0.8
                continue
            elif samp_ld.shape[0] == 0:
                factor *= 0.8
            samp = samp.append(samp_ld, ignore_index = True)
        return samp.head(n)
//  [python]
def log_likelihood(self, pars):
        self.k.set_lengthscale(pars[0])
        #K = (pars[1] ** 2) * self.k(self.X, self.X)
        K = (self.signal_stdev ** 2) * self.k(self.X, self.X)
        K = (K + K.t()).mul(0.5)
        K = K + (pars[2]**2) * torch.eye(self.X.shape[0]).type(self.dtype).to(self.device)

        # Compute K's inverse with Cholesky factorisation.
        try:
            L = torch.cholesky(K)
            L_inv = L.inverse()
            K_inv = L_inv.t().mm(L_inv)
        except:
            K = K + 1e-1*torch.eye(self.X.shape[0]).type(self.dtype).to(self.device)
            L = torch.cholesky(K)
            L_inv = L.inverse()
            K_inv = L_inv.t().mm(L_inv)
            #K_inv = torch.inverse(K)

        # Compute the log likelihood.
        y_mean = self.y - pars[3]
        log_likelihood_dims = -0.5 * y_mean.t().mm(K_inv.mm(y_mean)).sum(dim=0)
        log_likelihood_dims -= L.diag().log().sum()
        log_likelihood_dims -= L.shape[0] / 2.0 * np.log(2 * np.pi)
        log_likelihood = log_likelihood_dims.sum(dim=-1)
        print("Negative log likelihood = ", -log_likelihood)
        return -log_likelihood
//  [python]
def log_likelihood(self, pars):
        self.k.set_lengthscale(pars[0])
        #K = (pars[1] ** 2) * self.k(self.X, self.X)
        K = (self.signal_stdev ** 2) * self.k(self.X, self.X)
        K = (K + K.t()).mul(0.5)
        K = K + (pars[2]**2) * torch.eye(self.X.shape[0]).type(self.dtype).to(self.device)

        # Compute K's inverse with Cholesky factorisation.
        try:
            L = torch.linalg.cholesky(K)
            L_inv = L.inverse()
            K_inv = L_inv.t().mm(L_inv)
        except:
            K = K + 1e-1*torch.eye(self.X.shape[0]).type(self.dtype).to(self.device)
            L = torch.linalg.cholesky(K)
            L_inv = L.inverse()
            K_inv = L_inv.t().mm(L_inv)
            #K_inv = torch.inverse(K)

        # Compute the log likelihood.
        y_mean = self.y - pars[3]
        log_likelihood_dims = -0.5 * y_mean.t().mm(K_inv.mm(y_mean)).sum(dim=0)
        log_likelihood_dims -= L.diag().log().sum()
        log_likelihood_dims -= L.shape[0] / 2.0 * np.log(2 * np.pi)
        log_likelihood = log_likelihood_dims.sum(dim=-1)
        print("Negative log likelihood = ", -log_likelihood)
        return -log_likelihood
//  [python]
def update(self, X):
        self.X = X
        dim = self.X.shape[1]

        if self.objective is not None:
            #sampler = MultivariateNormal(self.X, 1e-2*torch.eye(dim).type(self.dtype).to(self.device))
            #self.Xplus = torch.cat((self.X, sampler.sample((self.N,)).reshape([-1, dim])), 0)
            self.Xplus = torch.cat((self.X, self.X - 1e-2*self.X.reshape([-1, dim])), 0)
            self.yplus = self.objective(self.Xplus)
        else:
            self.Xplus = self.X
            self.yplus = self.y

        #Fix for potential nan
        nan_idx = torch.isnan(self.yplus)
        self.Xplus = self.Xplus[~nan_idx[:, 0], :]
        self.yplus = self.yplus[~nan_idx[:, 0]]

        K = (self.signal_stdev**2)*self.k(self.Xplus, self.Xplus)
        K = (K + K.t()).mul(0.5)
        K = K + (self.noise**2) * torch.eye(self.Xplus.shape[0]).type(self.dtype).to(self.device)

        # Compute K's inverse with Cholesky factorisation.
        try:
            self.L = torch.cholesky(K)
            self.L_inv = self.L.inverse()
            self.K_inv = self.L_inv.t().mm(self.L_inv)
        except:
            self.K_inv = torch.inverse(K)
            print("Warning: using pseudoinverse")
//  [python]
def update(self, X):
        self.X = X
        dim = self.X.shape[1]

        if self.objective is not None:
            #sampler = MultivariateNormal(self.X, 1e-2*torch.eye(dim).type(self.dtype).to(self.device))
            #self.Xplus = torch.cat((self.X, sampler.sample((self.N,)).reshape([-1, dim])), 0)
            self.Xplus = torch.cat((self.X, self.X - 1e-2*self.X.reshape([-1, dim])), 0)
            self.yplus = self.objective(self.Xplus)
        else:
            self.Xplus = self.X
            self.yplus = self.y

        #Fix for potential nan
        nan_idx = torch.isnan(self.yplus)
        self.Xplus = self.Xplus[~nan_idx[:, 0], :]
        self.yplus = self.yplus[~nan_idx[:, 0]]

        K = (self.signal_stdev**2)*self.k(self.Xplus, self.Xplus)
        K = (K + K.t()).mul(0.5)
        K = K + (self.noise**2) * torch.eye(self.Xplus.shape[0]).type(self.dtype).to(self.device)

        # Compute K's inverse with Cholesky factorisation.
        try:
            self.L = torch.linalg.cholesky(K)
            self.L_inv = torch.linalg.solve(self.L, torch.eye(self.L.shape[0]).type(self.dtype).to(self.device))
            self.K_inv = self.L_inv.t().mm(self.L_inv)
        except:
            self.K_inv = torch.inverse(K)
            print("Warning: using pseudoinverse")
//  [python]
def update(self, X):
        if self.objective is not None:
            self.X = torch.cat((self.X, X), 0)
            y = self.objective(X)
            self.y = torch.cat((self.y, y - self.mean), 0)

        size_X = self.X.shape[0]
        if size_X > self.max_K_size:
            self.X = self.X[size_X - self.max_K_size: -1, :]
            self.y = self.y[size_X - self.max_K_size: -1, :]

        K = (self.signal_stdev ** 2) * self.k(self.X, self.X)
        K = (K + K.t()).mul(0.5)
        K = K + (self.noise ** 2) * torch.eye(self.X.shape[0]).type(self.dtype).to(self.device)

        # Compute K's inverse with Cholesky factorisation.
        try:
            self.L = torch.cholesky(K)
            self.L_inv = self.L.inverse()
            self.K_inv = self.L_inv.t().mm(self.L_inv)
        except:
            self.K_inv = torch.inverse(K)
            print("Warning: using pseudoinverse")
        #print("Loglikelihood = ",
        #      self.log_likelihood([self.lenghscales, self.signal_stdev, self.noise, self.mean]))
        self.n_updates += 1
        if self.n_updates%self.n_optimize == 0:
            params = self.optimize(pars=[self.lenghscales,
                                         self.signal_stdev,
                                         self.noise,
                                         self.mean], n_steps=10)
            print(params)
            self.lenghscales = params[0]
            self.k.set_lengthscale(params[0])
            #self.signal_stdev = params[1]
            #self.noise = params[2]
            self.mean = params[3]
//  [python]
def update(self, X):
        if self.objective is not None:
            self.X = torch.cat((self.X, X), 0)
            y = self.objective(X)
            self.y = torch.cat((self.y, y - self.mean), 0)

        size_X = self.X.shape[0]
        if size_X > self.max_K_size:
            self.X = self.X[size_X - self.max_K_size: -1, :]
            self.y = self.y[size_X - self.max_K_size: -1, :]

        K = (self.signal_stdev ** 2) * self.k(self.X, self.X)
        K = (K + K.t()).mul(0.5)
        K = K + (self.noise ** 2) * torch.eye(self.X.shape[0]).type(self.dtype).to(self.device)

        # Compute K's inverse with Cholesky factorisation.
        try:
            self.L = torch.linalg.cholesky(K)
            self.L_inv = torch.linalg.solve(self.L, torch.eye(self.L.shape[0]).to(self.device))
            self.K_inv = self.L_inv.t().mm(self.L_inv)
        except:
            self.K_inv = torch.inverse(K)
            print("Warning: using pseudoinverse")
        #print("Loglikelihood = ",
        #      self.log_likelihood([self.lenghscales, self.signal_stdev, self.noise, self.mean]))
        self.n_updates += 1
        if self.n_updates%self.n_optimize == 0:
            params = self.optimize(pars=[self.lenghscales,
                                         self.signal_stdev,
                                         self.noise,
                                         self.mean], n_steps=10)
            print(params)
            self.lenghscales = params[0]
            self.k.set_lengthscale(params[0])
            #self.signal_stdev = params[1]
            #self.noise = params[2]
            self.mean = params[3]
//  [python]
def psd_safe_cholesky(A, upper=False, out=None, jitter=None):
	"""Compute the Cholesky decomposition of A. If A is only p.s.d, add a small jitter to the diagonal.
	Args:
		:attr:`A` (Tensor):
			The tensor to compute the Cholesky decomposition of
		:attr:`upper` (bool, optional):
			See torch.cholesky
		:attr:`out` (Tensor, optional):
			See torch.cholesky
		:attr:`jitter` (float, optional):
			The jitter to add to the diagonal of A in case A is only p.s.d. If omitted, chosen
			as 1e-6 (float) or 1e-8 (double)
	"""
	try:
		L = torch.cholesky(A, upper=upper, out=out)
		return L
	except RuntimeError as e:
		isnan = torch.isnan(A)
		if isnan.any():
			raise NanError(
				f"cholesky_cpu: {isnan.sum().item()} of {A.numel()} elements of the {A.shape} tensor are NaN."
			)

		if jitter is None:
			jitter = 1e-6 if A.dtype == torch.float32 else 1e-8
		Aprime = A.clone()
		jitter_prev = 0
		for i in range(5):
			jitter_new = jitter * (10 ** i)
			Aprime.diagonal(dim1=-2, dim2=-1).add_(jitter_new - jitter_prev)
			jitter_prev = jitter_new
			try:
				L = torch.cholesky(Aprime, upper=upper, out=out)
				warnings.warn(
					f"A not p.d., added jitter of {jitter_new} to the diagonal",
					RuntimeWarning,
				)
				return L
			except RuntimeError:
				continue
		raise e
//  [python]
def psd_safe_cholesky(A, upper=False, out=None, jitter=None):
	"""Compute the Cholesky decomposition of A. If A is only p.s.d, add a small jitter to the diagonal.
	Args:
		:attr:`A` (Tensor):
			The tensor to compute the Cholesky decomposition of
		:attr:`upper` (bool, optional):
			See torch.linalg.cholesky
		:attr:`out` (Tensor, optional):
			See torch.linalg.cholesky
		:attr:`jitter` (float, optional):
			The jitter to add to the diagonal of A in case A is only p.s.d. If omitted, chosen
			as 1e-6 (float) or 1e-8 (double)
	"""
	try:
		L = torch.linalg.cholesky(A, upper=upper, out=out)
		return L
	except RuntimeError as e:
		isnan = torch.isnan(A)
		if isnan.any():
			raise NanError(
				f"cholesky_cpu: {isnan.sum().item()} of {A.numel()} elements of the {A.shape} tensor are NaN."
			)

		if jitter is None:
			jitter = 1e-6 if A.dtype == torch.float32 else 1e-8
		Aprime = A.clone()
		jitter_prev = 0
		for i in range(5):
			jitter_new = jitter * (10 ** i)
			Aprime.diagonal(dim1=-2, dim2=-1).add_(jitter_new - jitter_prev)
			jitter_prev = jitter_new
			try:
				L = torch.linalg.cholesky(Aprime, upper=upper, out=out)
				warnings.warn(
					f"A not p.d., added jitter of {jitter_new} to the diagonal",
					RuntimeWarning,
				)
				return L
			except RuntimeError:
				continue
		raise e
//  [python]
def log_norm(self):
        mean, cov = self.params.mean, self.params.cov
        size = mean.shape
        dim = self.dim
        if len(size) <= 1:
            mean = mean.view(1, -1)
            cov = cov.view(1, dim, dim)

        L = torch.cholesky(cov, upper=False)
        logdet = 2 * torch.log(L[:, range(dim), range(dim)]).sum(dim=-1)
        prec = cov.inverse()
        Sm = torch.matmul(prec, mean[:, :, None]).view(-1, dim)
        mSm = (Sm * mean).sum(dim=-1)
        log_base_measure = .5 * dim * math.log(2 * math.pi)
        return .5 * (logdet + mSm) + log_base_measure
//  [python]
def log_norm(self):
        mean, cov = self.params.mean, self.params.cov
        size = mean.shape
        dim = self.dim
        if len(size) <= 1:
            mean = mean.view(1, -1)
            cov = cov.view(1, dim, dim)

        L = torch.linalg.cholesky(cov, upper=False)
        logdet = 2 * torch.log(L[:, range(dim), range(dim)]).sum(dim=-1)
        prec = cov.inverse()
        Sm = torch.matmul(prec, mean[:, :, None]).view(-1, dim)
        mSm = (Sm * mean).sum(dim=-1)
        log_base_measure = .5 * dim * math.log(2 * math.pi)
        return .5 * (logdet + mSm) + log_base_measure
//  [python]
def forward(self, stats, pdfwise=False):
        nparams = self.natural_parameters()
        mean = self.params.mean
        cov = self.params.cov
        size = mean.shape
        dim = self.dim
        if len(size) <= 1:
            mean = mean.view(1, -1)
            nparams = nparams.view(1, -1)

        # Get the precision matrix from the natural parameters not
        # to inverse the covariance matrix once more time.
        prec = nparams[:, dim: dim * (dim + 1)].reshape(-1, dim, dim)

        # Log determinant of all the covariance matrix.
        L = torch.cholesky(cov, upper=False)
        logdet = 2 * torch.log(L[:, range(dim), range(dim)]).sum(dim=-1)

        # Quadratic term of the log-normalizer: -.5 * mu^T S mu
        Sm = nparams[:, :dim]
        mSm = (Sm * mean).sum(dim=-1)
        lnorm = .5 * (logdet + mSm)
        log_basemeasure = -.5 * (dim * math.log(2 * math.pi))

        if pdfwise:
            return torch.sum(nparams * stats, dim=-1) - lnorm \
                   + log_basemeasure
        retval = nparams @ stats.t() - lnorm[:, None] + log_basemeasure
        if len(size) <= 1:
            return retval.reshape(-1)
        return retval
//  [python]
def forward(self, stats, pdfwise=False):
        nparams = self.natural_parameters()
        mean = self.params.mean
        cov = self.params.cov
        size = mean.shape
        dim = self.dim
        if len(size) <= 1:
            mean = mean.view(1, -1)
            nparams = nparams.view(1, -1)

        # Get the precision matrix from the natural parameters not
        # to inverse the covariance matrix once more time.
        prec = nparams[:, dim: dim * (dim + 1)].reshape(-1, dim, dim)

        # Log determinant of all the covariance matrix.
        L = torch.linalg.cholesky(cov, upper=False)
        logdet = 2 * torch.log(L[:, range(dim), range(dim)]).sum(dim=-1)

        # Quadratic term of the log-normalizer: -.5 * mu^T S mu
        Sm = nparams[:, :dim]
        mSm = (Sm * mean).sum(dim=-1)
        lnorm = .5 * (logdet + mSm)
        log_basemeasure = -.5 * (dim * math.log(2 * math.pi))

        if pdfwise:
            return torch.sum(nparams * stats, dim=-1) - lnorm \
                   + log_basemeasure
        retval = nparams @ stats.t() - lnorm[:, None] + log_basemeasure
        if len(size) <= 1:
            return retval.reshape(-1)
        return retval
//  [python]
def estimate_precisions(self):
        identities = (
            torch.eye(self.n_features)
            .unsqueeze(0)
            .repeat(self.n_mixtures, 1, 1)
            .to(self.device)
        )
        L = torch.cholesky(self.sigma, upper=False)
        L_inv, _ = torch.triangular_solve(L, identities, upper=False)

        return L_inv.contiguous()
//  [python]
def estimate_precisions(self):
        identities = (
            torch.eye(self.n_features)
            .unsqueeze(0)
            .repeat(self.n_mixtures, 1, 1)
            .to(self.device)
        )
        L = torch.linalg.cholesky(self.sigma, upper=False)
        L_inv, _ = torch.triangular_solve(L, identities, upper=False)

        return L_inv.contiguous()
//  [python]
def compute_log_likelihood(self, X = None, Y = None):
        """
        Construct a tensorflow function to compute the likelihood.

            \log p(Y | theta).

        """
        assert X is None and Y is None, "{} does not support minibatch mode".format(str(type(self)))
        K = self.kern.K(self.X) + torch.eye(self.X.size(0), dtype=self.X.dtype, device=self.X.device) * self.likelihood.variance.get()
        L = torch.cholesky(K, upper=False)
        m = self.mean_function(self.X)
        return densities.multivariate_normal(self.Y, m, L)
//  [python]
def compute_log_likelihood(self, X = None, Y = None):
        """
        Construct a tensorflow function to compute the likelihood.

            \log p(Y | theta).

        """
        assert X is None and Y is None, "{} does not support minibatch mode".format(str(type(self)))
        K = self.kern.K(self.X) + torch.eye(self.X.size(0), dtype=self.X.dtype, device=self.X.device) * self.likelihood.variance.get()
        L = torch.linalg.cholesky(K, upper=False)
        m = self.mean_function(self.X)
        return densities.multivariate_normal(self.Y, m, L)
//  [python]
def sample(self, nsamples):
        mean, cov = self.params.mean, self.params.cov
        size = mean.shape
        dim = self.dim
        if len(size) <= 1:
            mean = mean.view(1, -1)
            cov = cov.view(1, dim, dim)
        L = torch.cholesky(cov, upper=False)
        noise = torch.randn(mean.shape[0], nsamples, mean.shape[1],
                            dtype=mean.dtype, device=mean.device)
        retval = mean[:, None, :] + torch.matmul(noise, L.permute(0, 2, 1))
        if len(size) <= 1:
            return retval.reshape(nsamples, -1)
        return retval
//  [python]
def sample(self, nsamples):
        mean, cov = self.params.mean, self.params.cov
        size = mean.shape
        dim = self.dim
        if len(size) <= 1:
            mean = mean.view(1, -1)
            cov = cov.view(1, dim, dim)
        L = torch.linalg.cholesky(cov, upper=False)
        noise = torch.randn(mean.shape[0], nsamples, mean.shape[1],
                            dtype=mean.dtype, device=mean.device)
        retval = mean[:, None, :] + torch.matmul(noise, L.permute(0, 2, 1))
        if len(size) <= 1:
            return retval.reshape(nsamples, -1)
        return retval
//  [python]
def cholesky_solve_test_helper(self, A_dims, b_dims, upper, device, dtype):
        from common_utils import random_symmetric_pd_matrix

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

        b = torch.randn(*b_dims, dtype=dtype, device=device)
        A = random_symmetric_pd_matrix(*A_dims, dtype=dtype, device=device)
        L = torch.linalg.cholesky(A, upper=upper)
        return b, A, L
//  [python]
def predict_f(self, Xnew, full_cov=False):
        """
        Xnew is a data matrix, point at which we want to predict

        This method computes

            p(F* | Y )

        where F* are points on the GP at Xnew, Y are noisy observations at X.

        """
        Kx = self.kern.K(self.X, Xnew)
        K = self.kern.K(self.X) + torch.eye(self.X.size(0), dtype=self.X.dtype, device=self.X.device) * self.likelihood.variance.get()
        L = torch.cholesky(K, upper=False)
        A,_ = torch.solve(Kx, L)  # could use triangular solve, note gesv has B first, then A in AX=B
        V,_ = torch.solve(self.Y - self.mean_function(self.X),L) # could use triangular solve
        fmean = torch.mm(A.t(), V) + self.mean_function(Xnew)
        if full_cov:
            fvar = self.kern.K(Xnew) - torch.mm(A.t(), A)
            fvar = fvar.unsqueeze(2).expand(fvar.size(0), fvar.size(1), self.Y.size(1))
        else:
            fvar = self.kern.Kdiag(Xnew) - (A**2).sum(0)
            fvar = fvar.view(-1,1)
            fvar = fvar.expand(fvar.size(0),self.Y.size(1))
        return fmean, fvar
//  [python]
def predict_f(self, Xnew, full_cov=False):
        """
        Xnew is a data matrix, point at which we want to predict

        This method computes

            p(F* | Y )

        where F* are points on the GP at Xnew, Y are noisy observations at X.

        """
        Kx = self.kern.K(self.X, Xnew)
        K = self.kern.K(self.X) + torch.eye(self.X.size(0), dtype=self.X.dtype, device=self.X.device) * self.likelihood.variance.get()
        L = torch.linalg.cholesky(K, upper=False)
        A,_ = torch.solve(Kx, L)  # could use triangular solve, note gesv has B first, then A in AX=B
        V,_ = torch.solve(self.Y - self.mean_function(self.X),L) # could use triangular solve
        fmean = torch.mm(A.t(), V) + self.mean_function(Xnew)
        if full_cov:
            fvar = self.kern.K(Xnew) - torch.mm(A.t(), A)
            fvar = fvar.unsqueeze(2).expand(fvar.size(0), fvar.size(1), self.Y.size(1))
        else:
            fvar = self.kern.Kdiag(Xnew) - (A**2).sum(0)
            fvar = fvar.view(-1,1)
            fvar = fvar.expand(fvar.size(0),self.Y.size(1))
        return fmean, fvar
//  [python]
def test_cholesky_solve_batched_non_contiguous(self, device, dtype):
        from numpy.linalg import solve
        from common_utils import random_symmetric_pd_matrix

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

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

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

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

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

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

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

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

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

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

        def cholesky_test_helper(n, batch_dims, upper):
            A = random_symmetric_pd_matrix(n, *batch_dims, dtype=dtype, device=device)
            cholesky_exp = torch.stack([m.cholesky(upper=upper) for m in A.reshape(-1, n, n)])
            cholesky_exp = cholesky_exp.reshape_as(A)
            self.assertEqual(cholesky_exp, torch.cholesky(A, upper=upper))

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

        def cholesky_test_helper(n, batch_dims, upper):
            A = random_symmetric_pd_matrix(n, *batch_dims, dtype=dtype, device=device)
            cholesky_exp = torch.stack([torch.linalg.cholesky(m, upper=upper) for m in A.reshape(-1, n, n)])
            cholesky_exp = cholesky_exp.reshape_as(A)
            self.assertEqual(cholesky_exp, torch.linalg.cholesky(A, upper=upper))

        for upper, batchsize in product([True, False], [(3,), (3, 4), (2, 3, 4)]):
            cholesky_test_helper(3, batchsize, upper)
//  [python]
def test_cholesky_batched_many_batches(self, device, dtype):
        from common_utils import random_symmetric_pd_matrix

        def cholesky_test_helper(n, batchsize, device, upper):
            A = random_symmetric_pd_matrix(n, batchsize, dtype=dtype, device=device)
            chol_fact = torch.cholesky(A, upper=upper)
            if upper:
                # Correctness check
                self.assertEqual(A, chol_fact.transpose(-2, -1).matmul(chol_fact))
                # Upper triangular check
                self.assertEqual(chol_fact, chol_fact.triu())
            else:
                # Correctness check
                self.assertEqual(A, chol_fact.matmul(chol_fact.transpose(-2, -1)))
                # Lower triangular check
                self.assertEqual(chol_fact, chol_fact.tril())

        for upper, batchsize in product([True, False], [262144, 524288]):
            cholesky_test_helper(2, batchsize, device, upper)
//  [python]
def test_cholesky_batched_many_batches(self, device, dtype):
        from common_utils import random_symmetric_pd_matrix

        def cholesky_test_helper(n, batchsize, device, upper):
            A = random_symmetric_pd_matrix(n, batchsize, dtype=dtype, device=device)
            chol_fact = torch.linalg.cholesky(A, upper=upper)
            if upper:
                # Correctness check
                self.assertEqual(A, chol_fact.transpose(-2, -1).matmul(chol_fact))
                # Upper triangular check
                self.assertEqual(chol_fact, chol_fact.triu())
            else:
                # Correctness check
                self.assertEqual(A, chol_fact.matmul(chol_fact.transpose(-2, -1)))
                # Lower triangular check
                self.assertEqual(chol_fact, chol_fact.tril())

        for upper, batchsize in product([True, False], [262144, 524288]):
            cholesky_test_helper(2, batchsize, device, upper)
//  [python]
def test_cholesky_solve_batched_broadcasting(self, device, dtype):
        from numpy.linalg import solve
        from common_utils import random_symmetric_pd_matrix

        def run_test(A_dims, b_dims, upper):
            A_matrix_size = A_dims[-1]
            A_batch_dims = A_dims[:-2]
            A = random_symmetric_pd_matrix(A_matrix_size, *A_batch_dims,
                                           dtype=dtype, device='cpu')
            b = torch.randn(*b_dims, dtype=dtype, device='cpu')
            x_exp = torch.tensor(solve(A.numpy(), b.numpy()), dtype=dtype, device=device)
            A, b = A.to(dtype=dtype, device=device), b.to(dtype=dtype, device=device)
            L = torch.cholesky(A, upper)
            x = torch.cholesky_solve(b, L, upper=upper)
            self.assertEqual(x, x_exp)

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

        def run_test(A_dims, b_dims, upper):
            A_matrix_size = A_dims[-1]
            A_batch_dims = A_dims[:-2]
            A = random_symmetric_pd_matrix(A_matrix_size, *A_batch_dims,
                                           dtype=dtype, device='cpu')
            b = torch.randn(*b_dims, dtype=dtype, device='cpu')
            x_exp = torch.tensor(solve(A.numpy(), b.numpy()), dtype=dtype, device=device)
            A, b = A.to(dtype=dtype, device=device), b.to(dtype=dtype, device=device)
            L = torch.linalg.cholesky(A, upper)
            x = torch.cholesky_solve(b, L, upper=upper)
            self.assertEqual(x, x_exp)

        # test against numpy.linalg.solve
        for upper in [True, False]:
            run_test((2, 1, 3, 4, 4), (2, 1, 3, 4, 6), upper)  # no broadcasting
            run_test((2, 1, 3, 4, 4), (4, 6), upper)  # broadcasting b
            run_test((4, 4), (2, 1, 3, 4, 2), upper)  # broadcasting A
            run_test((1, 3, 1, 4, 4), (2, 1, 3, 4, 5), upper)
//  [python]
def cholesky_invert(A):
	L = torch.cholesky(A)
	L_inv = torch.inverse(L)
	A_inv = torch.mm(L_inv.T, L_inv)

	return A_inv
//  [python]
def cholesky_invert(A):
	L = torch.linalg.cholesky(A)
	L_inv = torch.inverse(L)
	A_inv = torch.mm(L_inv.T, L_inv)

	return A_inv
//  [python]
def log_norm(self):
        mean, scale, = self.params.mean, self.params.scale
        scale_matrix, dof = self.params.scale_matrix, self.params.dof
        dim = mean.shape[-1]
        mean_size = mean.shape if len(mean.shape) > 0 else 1
        if len(mean_size) == 1:
            mean, scale = mean.view(1, -1), scale.view(1, 1)
            scale_matrix, dof = scale_matrix.view(1, dim, dim), dof.view(1, 1)

        idxs = torch.arange(1, dim + 1, dtype=mean.dtype, device=mean.device)
        L = torch.cholesky(scale_matrix, upper=False)
        logdet = 2 * torch.log(L[:, range(dim), range(dim)]).sum(dim=-1,
                                                                 keepdim=True)
        return (.5 * dof * logdet + .5 * dof * dim * math.log(2) \
               + .25 * dim * (dim - 1) * math.log(math.pi) \
               + torch.lgamma(.5 * (dof + 1 - idxs)).sum(dim=-1, keepdim=True) \
               - .5 * dim * torch.log(scale) \
               + .5 * dim * math.log(2 * math.pi)).sum(dim=-1)
//  [python]
def log_norm(self):
        mean, scale, = self.params.mean, self.params.scale
        scale_matrix, dof = self.params.scale_matrix, self.params.dof
        dim = mean.shape[-1]
        mean_size = mean.shape if len(mean.shape) > 0 else 1
        if len(mean_size) == 1:
            mean, scale = mean.view(1, -1), scale.view(1, 1)
            scale_matrix, dof = scale_matrix.view(1, dim, dim), dof.view(1, 1)

        idxs = torch.arange(1, dim + 1, dtype=mean.dtype, device=mean.device)
        L = torch.linalg.cholesky(scale_matrix, upper=False)
        logdet = 2 * torch.log(L[:, range(dim), range(dim)]).sum(dim=-1,
                                                                 keepdim=True)
        return (.5 * dof * logdet + .5 * dof * dim * math.log(2) \
               + .25 * dim * (dim - 1) * math.log(math.pi) \
               + torch.lgamma(.5 * (dof + 1 - idxs)).sum(dim=-1, keepdim=True) \
               - .5 * dim * torch.log(scale) \
               + .5 * dim * math.log(2 * math.pi)).sum(dim=-1)
//  [python]
def expected_sufficient_statistics(self):
        '''
        stats = (
            S * mu,
            S,
            tr(S * mu * mu^T),
            ln |S|
        )

        E[stats] = (
            v * W * m,
            v * W,
            (D/k) + tr(v * W * m * m^T),
            ( \sum_i psi(.5 * (v + 1 - i)) ) + D * ln 2 + ln |W|
        )
        '''
        mean, scale, = self.params.mean, self.params.scale
        scale_matrix, dof = self.params.scale_matrix, self.params.dof
        dim = mean.shape[-1]
        mean_size = mean.shape if len(mean.shape) > 0 else 1
        if len(mean_size) == 1:
            mean, scale = mean.view(1, -1), scale.view(1, 1)
            scale_matrix, dof = scale_matrix.view(1, dim, dim), dof.view(1, 1)

        idxs = torch.arange(1, dim + 1, dtype=mean.dtype, device=mean.device)
        L = torch.cholesky(scale_matrix, upper=False)
        logdet = 2 * torch.log(L[:, range(dim), range(dim)]).sum(dim=-1,
                                                                 keepdim=True)
        mean_quad = mean[:, :, None] * mean[:, None, :]
        exp_prec = dof[:, :, None] * scale_matrix
        retval = torch.cat([
            torch.matmul(exp_prec, mean[:, :, None]).reshape(-1, dim),
            exp_prec.reshape(-1, dim ** 2),
            (dim / scale) + _batch_trace(exp_prec, mean_quad, keepdim=True),
            (torch.digamma(.5 * (dof + 1 - idxs)).sum(dim=-1, keepdim=True) \
                + dim * math.log(2) + logdet)
        ], dim=-1)

        if len(mean_size) == 1:
            return retval.view(-1)
        return retval