jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def merge_quadratic_term_v2(model):
    nodes = list(model.nodes())
    for node in nodes:
        if (
            isinstance(node, BoundReduceSum)
            and isinstance(node.inputs[0], BoundMul)
            and isinstance(node.inputs[0].inputs[1], BoundMatMul)
            and node.inputs[0].inputs[1].inputs[0] == node.inputs[0].inputs[0]
            and not node.inputs[0].inputs[1].inputs[1].perturbed
        ):
            node_x = node.inputs[0].inputs[0]

            q = node.inputs[0].inputs[1].inputs[1].forward_value
            q_cholesky = torch.nn.Parameter(torch.linalg.cholesky(q))
            node_q_cholesky = BoundParams(f"{node.name}/q_cholesky", q_cholesky)
            node_q_cholesky.name = f"{node.name}/q_cholesky"

            node_x_matmul_q_cholesky = BoundMatMul(inputs=[node_x, node_q_cholesky])
            node_x_matmul_q_cholesky.name = f"{node.name}/x_matmul_q_cholesky"
            node_x_matmul_q_cholesky_sqr = BoundSqr(inputs=[node_x_matmul_q_cholesky])
            node_x_matmul_q_cholesky_sqr.name = f"{node.name}/x_matmul_q_cholesky_sqr"
            node_reduce_sum = BoundReduceSum(
                inputs=[node_x_matmul_q_cholesky_sqr], attr={"axes": node.axis}
            )
            node_reduce_sum.name = f"{node.name}/reduce_sum"

            node_quadratic = BoundQuadratic(
                inputs=[node_x, node_q_cholesky, node_reduce_sum]
            )
            node_quadratic.name = f"{node.name}/quadratic"

            model.add_nodes(
                [
                    node_q_cholesky,
                    node_x_matmul_q_cholesky,
                    node_x_matmul_q_cholesky_sqr,
                    node_reduce_sum,
                    node_quadratic,
                ]
            )
            model.replace_node(node, node_quadratic)
//  [python]
def initialize(self, data, lens):
        with torch.no_grad():
            lens = lens.lengths
            N = lens.sum()
            data_unrolled = torch.zeros(data.shape[1], N).to(data.device)
            start_idx = 0
            for b in range(data.shape[0]):
                data_unrolled[:, start_idx:start_idx+lens[b]] = data[b, :, :lens[b]]
                start_idx = start_idx + lens[b]

            input_mean = data_unrolled.mean(1, keepdims=True)
            centered_mels = data_unrolled - input_mean
            self.covar = (centered_mels @ centered_mels.transpose(0,1))/N
            inv_covar = torch.inverse(self.covar)
            whitening_matrix = torch.cholesky(inv_covar, upper=True).contiguous()
            if dist.is_initialized():
                dist.broadcast(whitening_matrix, 0)
                dist.broadcast(input_mean, 0)
            self.input_mean = input_mean
            # make sure to use in-place operations here so that it's tracked by pytorch
            self.upper_diag.copy_(torch.diag(whitening_matrix))
            self.upper.copy_(torch.triu(whitening_matrix, 1))
            self.initialized = ~self.initialized
//  [python]
def initialize(self, data, lens):
        with torch.no_grad():
            lens = lens.lengths
            N = lens.sum()
            data_unrolled = torch.zeros(data.shape[1], N).to(data.device)
            start_idx = 0
            for b in range(data.shape[0]):
                data_unrolled[:, start_idx:start_idx+lens[b]] = data[b, :, :lens[b]]
                start_idx = start_idx + lens[b]

            input_mean = data_unrolled.mean(1, keepdims=True)
            centered_mels = data_unrolled - input_mean
            self.covar = (centered_mels @ centered_mels.transpose(0,1))/N
            inv_covar = torch.inverse(self.covar)
            whitening_matrix = torch.linalg.cholesky(inv_covar, upper=True).contiguous()
            if dist.is_initialized():
                dist.broadcast(whitening_matrix, 0)
                dist.broadcast(input_mean, 0)
            self.input_mean = input_mean
            # make sure to use in-place operations here so that it's tracked by pytorch
            self.upper_diag.copy_(torch.diag(whitening_matrix))
            self.upper.copy_(torch.triu(whitening_matrix, 1))
            self.initialized = ~self.initialized
//  [python]
def update_dict_ridge(x, z, lambd=1e-4):
    """Update an (unconstrained) dictionary with ridge regression

    This is equivalent to a Newton step with the (L2-regularized) squared
    error objective:
        f(V) = (1/2N) * ||Vz - x||_2^2 + (lambd/2) * ||V||_2^2

    x : a batch of observations with shape (n_samples, n_features)
    z : a batch of code vectors with shape (n_samples, n_components)
    lambd : weight decay parameter
    """
    rhs = torch.mm(z.T, x)
    M = torch.mm(z.T, z)
    M.diagonal().add_(lambd * x.size(0))
    L = torch.cholesky(M)
    V = torch.cholesky_solve(rhs, L).T

    return V
//  [python]
def update_dict_ridge(x, z, lambd=1e-4):
    """Update an (unconstrained) dictionary with ridge regression

    This is equivalent to a Newton step with the (L2-regularized) squared
    error objective:
        f(V) = (1/2N) * ||Vz - x||_2^2 + (lambd/2) * ||V||_2^2

    x : a batch of observations with shape (n_samples, n_features)
    z : a batch of code vectors with shape (n_samples, n_components)
    lambd : weight decay parameter
    """
    rhs = torch.mm(z.T, x)
    M = torch.mm(z.T, z)
    M.diagonal().add_(lambd * x.size(0))
    L = torch.linalg.cholesky(M)
    V = torch.cholesky_solve(rhs, L).T

    return V
//  [python]
def reverse_propogator(pastdata=None, futuredata=None, epsilon=1e-6, method='cutoff'):

    """
    Compute the reverse propogator for the time-lagged sequence data with detailed balance contraint.
    This function is used to calculate the loss function for SRVNet.
    cholesky decomposition of reversed instantaneous correlation matrix: LL^{T} = (c_{00} + c_{tt}}) / 2
    propogator = L ((c_{0t} + c_{t0}}) / 2) L^{T}

    Parameters
    ----------
    pastdata: torch.Tensor, default=None
        The past part of the time-lagged sequence data. Can be generated from TimeLaggedDataset.
    futuredata: torch.Tensor, default=None
        The future part of the time-lagged sequence data. Can be generated from TimeLaggedDataset.
    epsilon: float, default=1e-6
        The epsilon modification used to calculate the inverse of the instantaneous matrix.
    method: string, default='cutoff'
        The method used to calculate the inverse of the instantaneous matix: : 'regularize', 'cutoff', 'clamp'
    
    Returns
    ----------
    _reverse_propogator: torch.Tensor, shape:(num_features, num_features)
        Reverse propogator matrix for the input time-lagged sequence data.
    """

    c00, c0t, ctt = covariance(pastdata=pastdata, futuredata=futuredata, remove_mean=True)
    _, ct0, _ = covariance(pastdata=futuredata, futuredata=pastdata, remove_mean=True)
    
    c0 = (c00 + ctt) / 2
    ct = (c0t + ct0) / 2

    _lower_triangle = torch.cholesky(c0)
    _lower_triangle_inv = torch.inverse(_lower_triangle)
    _reverse_propogator = torch.matmul(_lower_triangle_inv, torch.matmul(ct, _lower_triangle_inv.t()))

    return _reverse_propogator
//  [python]
def reverse_propogator(pastdata=None, futuredata=None, epsilon=1e-6, method='cutoff'):

    """
    Compute the reverse propogator for the time-lagged sequence data with detailed balance contraint.
    This function is used to calculate the loss function for SRVNet.
    cholesky decomposition of reversed instantaneous correlation matrix: LL^{T} = (c_{00} + c_{tt}}) / 2
    propogator = L ((c_{0t} + c_{t0}}) / 2) L^{T}

    Parameters
    ----------
    pastdata: torch.Tensor, default=None
        The past part of the time-lagged sequence data. Can be generated from TimeLaggedDataset.
    futuredata: torch.Tensor, default=None
        The future part of the time-lagged sequence data. Can be generated from TimeLaggedDataset.
    epsilon: float, default=1e-6
        The epsilon modification used to calculate the inverse of the instantaneous matrix.
    method: string, default='cutoff'
        The method used to calculate the inverse of the instantaneous matix: : 'regularize', 'cutoff', 'clamp'
    
    Returns
    ----------
    _reverse_propogator: torch.Tensor, shape:(num_features, num_features)
        Reverse propogator matrix for the input time-lagged sequence data.
    """

    c00, c0t, ctt = covariance(pastdata=pastdata, futuredata=futuredata, remove_mean=True)
    _, ct0, _ = covariance(pastdata=futuredata, futuredata=pastdata, remove_mean=True)
    
    c0 = (c00 + ctt) / 2
    ct = (c0t + ct0) / 2

    _lower_triangle = torch.linalg.cholesky(c0)
    _lower_triangle_inv = torch.inverse(_lower_triangle)
    _reverse_propogator = torch.matmul(_lower_triangle_inv, torch.matmul(ct, _lower_triangle_inv.t()))

    return _reverse_propogator
//  [python]
def calculate_vel_impulse(rela_v, J_c, mu, elasticity):

    # set up differentiable optimization problem
    f = cp.Variable((2, 1)) # [f_t, f_n]
    A_decom_p = cp.Parameter((2, 2))
    v_p = cp.Parameter((2, 1))
    mu_p = cp.Parameter((1, 1))
    objective = cp.Minimize(0.5 * cp.sum_squares(A_decom_p @ f) + cp.sum(cp.multiply(f, v_p)))
    constraints = [
        f[0] >= 0, # normal impulse should be creater than zero
        f[1] <= cp.multiply(mu_p[0], f[0]),
        f[1] >= -cp.multiply(mu_p[0], f[0]),
    ]
    problem = cp.Problem(objective, constraints)
    cvxpylayer = CvxpyLayer(problem, parameters=[A_decom_p, v_p, mu_p], variables=[f])

    # feed values to the layer
    A = J_c.transpose(0, 1) @ J_c
    A_decom = torch.cholesky(A)
    impulse, = cvxpylayer(
        A_decom,
        J_c @ rela_v.reshape(-1, 1),
        mu.reshape(-1, 1),
    )
    return (J_c.transpose(0, 1) @ impulse)[:, 0] * (1+elasticity)
//  [python]
def calculate_vel_impulse(rela_v, J_c, mu, elasticity):

    # set up differentiable optimization problem
    f = cp.Variable((2, 1)) # [f_t, f_n]
    A_decom_p = cp.Parameter((2, 2))
    v_p = cp.Parameter((2, 1))
    mu_p = cp.Parameter((1, 1))
    objective = cp.Minimize(0.5 * cp.sum_squares(A_decom_p @ f) + cp.sum(cp.multiply(f, v_p)))
    constraints = [
        f[0] >= 0, # normal impulse should be creater than zero
        f[1] <= cp.multiply(mu_p[0], f[0]),
        f[1] >= -cp.multiply(mu_p[0], f[0]),
    ]
    problem = cp.Problem(objective, constraints)
    cvxpylayer = CvxpyLayer(problem, parameters=[A_decom_p, v_p, mu_p], variables=[f])

    # feed values to the layer
    A = J_c.transpose(0, 1) @ J_c
    A_decom = torch.linalg.cholesky(A)
    impulse, = cvxpylayer(
        A_decom,
        J_c @ rela_v.reshape(-1, 1),
        mu.reshape(-1, 1),
    )
    return (J_c.transpose(0, 1) @ impulse)[:, 0] * (1+elasticity)
//  [python]
def inv(M, perc_damp, sub=True):
    dead = M.diag() == 0.0

    if dead.any() and sub:
        print("[WARNING] DIAGONAL VALUE 0 FOUND IN INV")

        M_out = full_to_sub(M_out, dead, dead)
        aug(M_out, perc_damp, inplace=True)

        M_out = torch.cholesky(M_out.double())
        M_out = torch.cholesky_inverse(M_out)
    else:
        M_out = aug(M, perc_damp, inplace=False)
        try:
            M_out = torch.cholesky(M_out.double())
        except:
            M_out = torch.cholesky(M_out.double())
        M_out = torch.cholesky_inverse(M_out)

    return M_out.float()
//  [python]
def inv(M, perc_damp, sub=True):
    dead = M.diag() == 0.0

    if dead.any() and sub:
        print("[WARNING] DIAGONAL VALUE 0 FOUND IN INV")

        M_out = full_to_sub(M_out, dead, dead)
        aug(M_out, perc_damp, inplace=True)

        M_out = torch.linalg.cholesky(M_out.double())
        M_out = torch.cholesky_inverse(M_out)
    else:
        M_out = aug(M, perc_damp, inplace=False)
        try:
            M_out = torch.linalg.cholesky(M_out.double())
        except:
            M_out = torch.linalg.cholesky(M_out.double())
        M_out = torch.cholesky_inverse(M_out)

    return M_out.float()
//  [python]
def mvn_hellinger_distance(
    p_mean: Tensor, q_mean: Tensor, p_covar: Tensor, q_covar: Tensor
) -> Tensor:
    """Computes the (2)-Hellinger distance between two multivariate normal
        distributions.
    Args:
        p_mean: A `batch_shape x dist_shape x 1`-dim Tensor of means of
            the first distribution.
        q_mean: A `batch_shape x dist_shape x dist_shape`-dim Tensor of
            covariances of the first distribution,  where the covariances
            are (optionally) in the q-batch dim.
        p_mean: A `batch_shape x dist_shape x 1`-dim Tensor of means of
            the second distribution.
        q_mean: A `batch_shape x dist_shape x dist_shape`-dim Tensor of
            covariances of the second distribution where the covariances
            are (optionally) in the q-batch dim.
    Returns:
        A tensor of shape `batch_shape x dist_shape` denoting the KL-divergence
        between the multivariate Gaussian distributions p and q.
    """
    p_logdet = torch.logdet(p_covar).unsqueeze(-1)
    q_logdet = torch.logdet(q_covar).unsqueeze(-1)
    avg_covar = (p_covar + q_covar) / 2

    # we need to re-use the cholesky decomp, so we compute it once here
    L_avg = torch.cholesky(avg_covar)
    L_inv = torch.inverse(L_avg)

    # removes one dimension, which needs to be recouped
    pq_logdet = (
        torch.pow(torch.diagonal(L_avg, dim1=-2, dim2=-1), 2)
        .prod(dim=-1, keepdim=True)
        .log()
    )
    base_logterm = 0.25 * (p_logdet + q_logdet) - 0.5 * pq_logdet

    mean_diff = p_mean - q_mean
    L_mean_diff = torch.matmul(L_inv, mean_diff)
    exp_logterm = -0.125 * torch.matmul(L_mean_diff.transpose(-2, -1), L_mean_diff)
    sq_hdist = 1 - (base_logterm + exp_logterm.squeeze(-1)).exp()
    return sq_hdist.sqrt()
//  [python]
def mvn_hellinger_distance(
    p_mean: Tensor, q_mean: Tensor, p_covar: Tensor, q_covar: Tensor
) -> Tensor:
    """Computes the (2)-Hellinger distance between two multivariate normal
        distributions.
    Args:
        p_mean: A `batch_shape x dist_shape x 1`-dim Tensor of means of
            the first distribution.
        q_mean: A `batch_shape x dist_shape x dist_shape`-dim Tensor of
            covariances of the first distribution,  where the covariances
            are (optionally) in the q-batch dim.
        p_mean: A `batch_shape x dist_shape x 1`-dim Tensor of means of
            the second distribution.
        q_mean: A `batch_shape x dist_shape x dist_shape`-dim Tensor of
            covariances of the second distribution where the covariances
            are (optionally) in the q-batch dim.
    Returns:
        A tensor of shape `batch_shape x dist_shape` denoting the KL-divergence
        between the multivariate Gaussian distributions p and q.
    """
    p_logdet = torch.logdet(p_covar).unsqueeze(-1)
    q_logdet = torch.logdet(q_covar).unsqueeze(-1)
    avg_covar = (p_covar + q_covar) / 2

    # we need to re-use the cholesky decomp, so we compute it once here
    L_avg = torch.linalg.cholesky(avg_covar)
    L_inv = torch.inverse(L_avg)

    # removes one dimension, which needs to be recouped
    pq_logdet = (
        torch.pow(torch.diagonal(L_avg, dim1=-2, dim2=-1), 2)
        .prod(dim=-1, keepdim=True)
        .log()
    )
    base_logterm = 0.25 * (p_logdet + q_logdet) - 0.5 * pq_logdet

    mean_diff = p_mean - q_mean
    L_mean_diff = torch.matmul(L_inv, mean_diff)
    exp_logterm = -0.125 * torch.matmul(L_mean_diff.transpose(-2, -1), L_mean_diff)
    sq_hdist = 1 - (base_logterm + exp_logterm.squeeze(-1)).exp()
    return sq_hdist.sqrt()
//  [python]
def __init__(self, mu, sigma):
        chol_sigma = torch.cholesky(sigma)

        self._mu = torch.nn.Parameter(mu)
        self._chol_sigma = torch.nn.Parameter(chol_sigma)

        super().__init__()

        self._add_save_attr(
            _mu='torch',
            _chol_sigma='torch'
        )
//  [python]
def __init__(self, mu, sigma):
        chol_sigma = torch.linalg.cholesky(sigma)

        self._mu = torch.nn.Parameter(mu)
        self._chol_sigma = torch.nn.Parameter(chol_sigma)

        super().__init__()

        self._add_save_attr(
            _mu='torch',
            _chol_sigma='torch'
        )
//  [python]
def cholesky_precision(covariances: torch.Tensor, covariance_type: CovarianceType) -> torch.Tensor:
    """
    Computes the Cholesky decompositions of the precision matrices induced by the provided
    covariance matrices.

    Args:
        covariances: A tensor of shape ``[num_components, dim, dim]``, ``[dim, dim]``,
            ``[num_components, dim]``, ``[dim]`` or ``[num_components]`` depending on the
            ``covariance_type``. These are the covariance matrices of multivariate Normal
            distributions.
        covariance_type: The type of covariance for the covariance matrices given.

    Returns:
        A tensor of the same shape as ``covariances``, providing the lower-triangular Cholesky
        decompositions of the precision matrices.
    """
    if covariance_type in ("tied", "full"):
        # Compute Cholesky decomposition
        cholesky = torch.cholesky(covariances)
        # Invert
        num_features = covariances.size(-1)
        target = torch.eye(num_features, dtype=covariances.dtype, device=covariances.device)
        if covariance_type == "full":
            num_components = covariances.size(0)
            target = target.unsqueeze(0).expand(num_components, -1, -1)
        return torch.triangular_solve(target, cholesky, upper=False)[0].transpose(-2, -1)

    # "Simple" kind of covariance
    return covariances.sqrt().reciprocal()
//  [python]
def cholesky_precision(covariances: torch.Tensor, covariance_type: CovarianceType) -> torch.Tensor:
    """
    Computes the Cholesky decompositions of the precision matrices induced by the provided
    covariance matrices.

    Args:
        covariances: A tensor of shape ``[num_components, dim, dim]``, ``[dim, dim]``,
            ``[num_components, dim]``, ``[dim]`` or ``[num_components]`` depending on the
            ``covariance_type``. These are the covariance matrices of multivariate Normal
            distributions.
        covariance_type: The type of covariance for the covariance matrices given.

    Returns:
        A tensor of the same shape as ``covariances``, providing the lower-triangular Cholesky
        decompositions of the precision matrices.
    """
    if covariance_type in ("tied", "full"):
        # Compute Cholesky decomposition
        cholesky = torch.linalg.cholesky(covariances)
        # Invert
        num_features = covariances.size(-1)
        target = torch.eye(num_features, dtype=covariances.dtype, device=covariances.device)
        if covariance_type == "full":
            num_components = covariances.size(0)
            target = target.unsqueeze(0).expand(num_components, -1, -1)
        return torch.linalg.solve_triangular(cholesky, target, upper=False).transpose(-2, -1)

    # "Simple" kind of covariance
    return covariances.sqrt().reciprocal()
//  [python]
def test_cholesky(self):
        linear_op = self.create_linear_op()
        evaluated = self.evaluate_linear_op(linear_op)
        for upper in (False, True):
            res = torch.cholesky(linear_op, upper=upper).to_dense()
            actual = torch.cholesky(evaluated, upper=upper)
            self.assertAllClose(res, actual, **self.tolerances["cholesky"])
//  [python]
def test_cholesky(self):
        linear_op = self.create_linear_op()
        evaluated = self.evaluate_linear_op(linear_op)
        for upper in (False, True):
            res = torch.linalg.cholesky(linear_op, upper=upper).to_dense()
            actual = torch.linalg.cholesky(evaluated, upper=upper)
            self.assertAllClose(res, actual, **self.tolerances["cholesky"])
//  [python]
def compute_L(self):
        if self.L is None:
            try:
                self.L = torch.cholesky(torch.inverse(precision))
            except:
                print("Cholesky failed. Using standard normal.")
                self.L = torch.cholesky(torch.inverse(precision.cuda()*0.+torch.eye(mu.shape[-1]).cuda()).cuda()).cuda()
//  [python]
def compute_L(self):
        if self.L is None:
            try:
                self.L = torch.linalg.cholesky(torch.inverse(precision))
            except:
                print("Cholesky failed. Using standard normal.")
                self.L = torch.linalg.cholesky(torch.inverse(precision.cuda()*0.+torch.eye(mu.shape[-1]).cuda()).cuda()).cuda()
//  [python]
def Hsummary(H, percdamp=0.01):
    assert H.shape[0] == H.shape[1]
    n = H.shape[0]
    damp = percdamp * torch.mean(torch.diag(H))
    diag = torch.arange(n)
    H[diag, diag] += damp
    L = torch.cholesky(H)
    D = torch.diag(L).square()
    a = D.sum() / H.trace()
    k00 = torch.linalg.matrix_rank(H) / n
    k01 = torch.linalg.matrix_rank(H, rtol=0.01) / n
    _, Q = torch.linalg.eigh(H)
    mu = torch.linalg.matrix_norm(Q) * np.sqrt(n)
    return a, k00, k01, mu
//  [python]
def Hsummary(H, percdamp=0.01):
    assert H.shape[0] == H.shape[1]
    n = H.shape[0]
    damp = percdamp * torch.mean(torch.diag(H))
    diag = torch.arange(n)
    H[diag, diag] += damp
    L = torch.linalg.cholesky(H)
    D = torch.diag(L).square()
    a = D.sum() / H.trace()
    k00 = torch.linalg.matrix_rank(H) / n
    k01 = torch.linalg.matrix_rank(H, rtol=0.01) / n
    _, Q = torch.linalg.eigh(H)
    mu = torch.linalg.matrix_norm(Q) * np.sqrt(n)
    return a, k00, k01, mu
//  [python]
def forward(self, X, return_prior_samples=False):
        
        
        Ku = self.kernel(self.inducing_points) + 2e-6 * torch.eye(self.num_inducing)

        Kf = self.kernel(X) + 2e-6 * torch.eye(X.shape[0])
        Kfu = self.kernel(X, self.inducing_points)
        
        self.Lu = torch.cholesky(Ku +  1e-4 * torch.eye(self.num_inducing))

        A = torch.triangular_solve(Kfu, self.Lu.permute(0,2,1), upper=True)[0]
        A = torch.triangular_solve(A, self.Lu, upper=False)[0]

        mean = torch.einsum("dnm, md -> nd", A, self.q_mu)

        # Shape (S, S, D)
        self.q_sqrt = (
            torch.zeros((self.num_inducing, self.num_inducing, self.output_dim))
            .to(self.dtype)
            .to(self.device)
        )
        li, lj = torch.tril_indices(self.num_inducing, self.num_inducing)
        self.q_sqrt[li, lj] = self.q_sqrt_tri
        
        SK = torch.einsum("nmd, bmd->dnb", self.q_sqrt, self.q_sqrt) - Ku
        
        B = torch.einsum("dmb, dnb -> dmn", SK, A)
        
        delta_cov = torch.sum(A * B.permute(0, 2, 1), -1)
        K = torch.diagonal(Kf, dim1=1, dim2=2) + delta_cov
        
        # Add layer noise to variance
        if self.log_layer_noise is not None:
            K = K + torch.exp(self.log_layer_noise)

        # Add mean function
        if self.mean_function is not None:
            mean = mean + self.mean_function(X)
            
        if return_prior_samples:
            return mean, K, torch.zeros(20, X.shape[0], 1)
        return mean, K
//  [python]
def forward(self, X, return_prior_samples=False):
        
        
        Ku = self.kernel(self.inducing_points) + 2e-6 * torch.eye(self.num_inducing)

        Kf = self.kernel(X) + 2e-6 * torch.eye(X.shape[0])
        Kfu = self.kernel(X, self.inducing_points)
        
        self.Lu = torch.linalg.cholesky(Ku +  1e-4 * torch.eye(self.num_inducing))

        A = torch.linalg.solve_triangular(self.Lu.permute(0,2,1), Kfu, upper=True, left= False)
        A = torch.linalg.solve_triangular(self.Lu, A, upper=False, left = False)

        mean = torch.einsum("dnm, md -> nd", A, self.q_mu)

        # Shape (S, S, D)
        self.q_sqrt = (
            torch.zeros((self.num_inducing, self.num_inducing, self.output_dim))
            .to(self.dtype)
            .to(self.device)
        )
        li, lj = torch.tril_indices(self.num_inducing, self.num_inducing)
        self.q_sqrt[li, lj] = self.q_sqrt_tri
        
        SK = torch.einsum("nmd, bmd->dnb", self.q_sqrt, self.q_sqrt) - Ku
        
        B = torch.einsum("dmb, dnb -> dmn", SK, A)
        
        delta_cov = torch.sum(A * B.permute(0, 2, 1), -1)
        K = torch.diagonal(Kf, dim1=1, dim2=2) + delta_cov
        
        # Add layer noise to variance
        if self.log_layer_noise is not None:
            K = K + torch.exp(self.log_layer_noise)

        # Add mean function
        if self.mean_function is not None:
            mean = mean + self.mean_function(X)
            
        if return_prior_samples:
            return mean, K, torch.zeros(20, X.shape[0], 1)
        return mean, K
//  [python]
def cholesky(
    x: torch.Tensor, /, *, upper: bool = False, out: Optional[torch.Tensor] = None
) -> torch.Tensor:
    if not upper:
        return torch.cholesky(x, out=out)
    else:
        ret = torch.transpose(
            torch.cholesky(
                torch.transpose(x, dim0=len(x.shape) - 1, dim1=len(x.shape) - 2)
            ),
            dim0=len(x.shape) - 1,
            dim1=len(x.shape) - 2,
        )
        if ivy.exists(out):
            return ivy.inplace_update(out, ret)
        return ret
//  [python]
def cholesky(
    x: torch.Tensor, /, *, upper: bool = False, out: Optional[torch.Tensor] = None
) -> torch.Tensor:
    if not upper:
        return torch.linalg.cholesky(x, out=out)
    else:
        ret = torch.transpose(
            torch.linalg.cholesky(
                torch.transpose(x, dim0=len(x.shape) - 1, dim1=len(x.shape) - 2)
            ),
            dim0=len(x.shape) - 1,
            dim1=len(x.shape) - 2,
        )
        if ivy.exists(out):
            return ivy.inplace_update(out, ret)
        return ret
//  [python]
def thompson_sampling_multivariate(
    f_mean: torch.Tensor, f_cov: torch.Tensor,
    gamma: float = 1., random_state: int = 1
) -> torch.Tensor:
    """
    Multi objective sampling.

    Parameters:
    -----------
    f_mean: torch.Tensor
        Shape (B, K)

    f_cov: torch.Tensor
        Shape (B, K, K)

    Returns:
    --------
    ts: torch.Tensor
        Shape (B, K)
    """
    assert len(f_mean.shape) == 2 and len(f_cov.shape) == 3
    device = f_mean.device
    generator = torch.Generator(device=device).manual_seed(random_state)
    return f_mean + gamma * torch.einsum(
        'bij,bj->bi',
        torch.cholesky(f_cov),  # (B, K, K)
        torch.randn(f_mean.shape, device=device, generator=generator)  # (B, K)
    )
//  [python]
def thompson_sampling_multivariate(
    f_mean: torch.Tensor, f_cov: torch.Tensor,
    gamma: float = 1., random_state: int = 1
) -> torch.Tensor:
    """
    Multi objective sampling.

    Parameters:
    -----------
    f_mean: torch.Tensor
        Shape (B, K)

    f_cov: torch.Tensor
        Shape (B, K, K)

    Returns:
    --------
    ts: torch.Tensor
        Shape (B, K)
    """
    assert len(f_mean.shape) == 2 and len(f_cov.shape) == 3
    device = f_mean.device
    generator = torch.Generator(device=device).manual_seed(random_state)
    return f_mean + gamma * torch.einsum(
        'bij,bj->bi',
        torch.linalg.cholesky(f_cov),  # (B, K, K)
        torch.randn(f_mean.shape, device=device, generator=generator)  # (B, K)
    )
//  [python]
def trKsquare(B, A):
    G = torch.cholesky(B)
    inv_G = torch.inverse(G)
    W = torch.einsum("...ij,...jk,...lk->...il", inv_G, A, inv_G)
    lamda , _ = se.apply(W.reshape((-1,3,3)))
    lamda = lamda.reshape((*W.shape[:-2],3))
    result = torch.sum(torch.log(lamda) ** 2, (-1))
    return result
//  [python]
def trKsquare(B, A):
    G = torch.linalg.cholesky(B)
    inv_G = torch.inverse(G)
    W = torch.einsum("...ij,...jk,...lk->...il", inv_G, A, inv_G)
    lamda , _ = se.apply(W.reshape((-1,3,3)))
    lamda = lamda.reshape((*W.shape[:-2],3))
    result = torch.sum(torch.log(lamda) ** 2, (-1))
    return result
//  [python]
def __init__(self, n, alpha=1e-3, beta=2, kappa=0, sqrt_method=None):
        """Constructor for the MerweScaledUTF.

        :param n: Dimensionality of the state. 2n+1 weights will be generated.
        :type n: int
        :param alpha: Defines the spread of the sigma points around the mean.
            Usually a small positive value. Defaults to 1e-3.
        :type alpha: float
        :param beta: Incorporates prior knowledge of the distribution of the
            mean. For Gaussian distributions, beta=2 is optimal. Defaults to 2.
        :type beta: float
        :param kappa: Secondary scaling parameter usually set to 0 or to 3-n.
            Defaults to 0.
        :type kappa: float
        :param sqrt_method: Defines how we compute the square root of a matrix,
            which has no unique answer. Cholesky is the default choice due to
            its speed. As of [1] this was not a well researched area and
            different choices affect how the sigma points are arranged relative
            to the eigenvectors of the covariance matrix. If the default is not
            used, the alternative method must return an upper triangular matrix
            If None, defaults to `torch.linalg.cholesky`.
        :type sqrt_method: function
        """
        self.n = n
        self.pts = 2 * n + 1
        self.alpha = alpha
        self.beta = beta
        self.kappa = kappa
        if sqrt_method is None:
            self.sqrt = (
                lambda A: torch.cholesky(A.transpose(-2, -1).conj())
                .transpose(-2, -1)
                .conj()
            )
        else:
            self.sqrt = sqrt_method
        self._compute_weights()
//  [python]
def __init__(self, n, alpha=1e-3, beta=2, kappa=0, sqrt_method=None):
        """Constructor for the MerweScaledUTF.

        :param n: Dimensionality of the state. 2n+1 weights will be generated.
        :type n: int
        :param alpha: Defines the spread of the sigma points around the mean.
            Usually a small positive value. Defaults to 1e-3.
        :type alpha: float
        :param beta: Incorporates prior knowledge of the distribution of the
            mean. For Gaussian distributions, beta=2 is optimal. Defaults to 2.
        :type beta: float
        :param kappa: Secondary scaling parameter usually set to 0 or to 3-n.
            Defaults to 0.
        :type kappa: float
        :param sqrt_method: Defines how we compute the square root of a matrix,
            which has no unique answer. Cholesky is the default choice due to
            its speed. As of [1] this was not a well researched area and
            different choices affect how the sigma points are arranged relative
            to the eigenvectors of the covariance matrix. If the default is not
            used, the alternative method must return an upper triangular matrix
            If None, defaults to `torch.linalg.cholesky`.
        :type sqrt_method: function
        """
        self.n = n
        self.pts = 2 * n + 1
        self.alpha = alpha
        self.beta = beta
        self.kappa = kappa
        if sqrt_method is None:
            self.sqrt = (
                lambda A: torch.linalg.cholesky(A.transpose(-2, -1).conj())
                .transpose(-2, -1)
                .conj()
            )
        else:
            self.sqrt = sqrt_method
        self._compute_weights()
//  [python]
def forward(self, X, return_prior_samples=False):

        # Let S = num_coeffs, D = output_dim and N = num_samples
        # Shape (S, N, 1)
        X_and_Z = torch.cat([X, self.inducing_points], dim = 0)
        f = self.generative_function(X_and_Z)
        # Compute mean value, shape (1 , N, 1)
        f_mean = torch.mean(f, dim=0, keepdims=True)

        # Compute regresion function, shape (S , N, 1)
        phi = (f - f_mean) / torch.sqrt(torch.tensor(f.shape[0] - 1).type(self.dtype))
        f_K = torch.einsum("and, amd -> nmd", phi, phi)
        
        Ku = f_K[X.shape[0]:, X.shape[0]:].permute(2,0,1) + 2e-6 * torch.eye(self.num_inducing)
        Kf = f_K[:X.shape[0], :X.shape[0]].permute(2,0,1) + 2e-6 * torch.eye(X.shape[0])
        Kfu = f_K[:X.shape[0], X.shape[0]:].permute(2,0,1)
        
        self.Lu = torch.cholesky(Ku+ 1e-4 * torch.eye(self.num_inducing))

        A = torch.triangular_solve(Kfu, self.Lu.permute(0,2,1), upper=True)[0]
        A = torch.triangular_solve(A, self.Lu, upper=False)[0]

        mean = torch.einsum("dnm, md -> nd", A, self.q_mu)

        # Shape (S, S, D)
        self.q_sqrt = (
            torch.zeros((self.num_inducing, self.num_inducing, self.output_dim))
            .to(self.dtype)
            .to(self.device)
        )
        li, lj = torch.tril_indices(self.num_inducing, self.num_inducing)
        self.q_sqrt[li, lj] = self.q_sqrt_tri
        
        SK = torch.einsum("nmd, bmd->dnb", self.q_sqrt, self.q_sqrt) - Ku
        
        B = torch.einsum("dmb, dnb -> dmn", SK, A)
        
        delta_cov = torch.sum(A * B.permute(0, 2, 1), -1)
        K = torch.diagonal(Kf, dim1=1, dim2=2) + delta_cov
        
        # Add layer noise to variance
        if self.log_layer_noise is not None:
            K = K + torch.exp(self.log_layer_noise).unsqueeze(-1)

        # Add mean function
        if self.mean_function is not None:
            mean = mean + self.mean_function(X)
            
        if return_prior_samples:
            return mean, K, f[:, :X.shape[0], :]
        return mean, K
//  [python]
def forward(self, X, return_prior_samples=False):

        # Let S = num_coeffs, D = output_dim and N = num_samples
        # Shape (S, N, 1)
        X_and_Z = torch.cat([X, self.inducing_points], dim = 0)
        f = self.generative_function(X_and_Z)
        # Compute mean value, shape (1 , N, 1)
        f_mean = torch.mean(f, dim=0, keepdims=True)

        # Compute regresion function, shape (S , N, 1)
        phi = (f - f_mean) / torch.sqrt(torch.tensor(f.shape[0] - 1).type(self.dtype))
        f_K = torch.einsum("and, amd -> nmd", phi, phi)
        
        Ku = f_K[X.shape[0]:, X.shape[0]:].permute(2,0,1) + 2e-6 * torch.eye(self.num_inducing)
        Kf = f_K[:X.shape[0], :X.shape[0]].permute(2,0,1) + 2e-6 * torch.eye(X.shape[0])
        Kfu = f_K[:X.shape[0], X.shape[0]:].permute(2,0,1)
        
        self.Lu = torch.linalg.cholesky(Ku+ 1e-4 * torch.eye(self.num_inducing))

        A = torch.linalg.solve_triangular(self.Lu.permute(0,2,1), Kfu, upper=True, left= False)
        A = torch.linalg.solve_triangular(self.Lu, A, upper=False, left = False)

        mean = torch.einsum("dnm, md -> nd", A, self.q_mu)

        # Shape (S, S, D)
        self.q_sqrt = (
            torch.zeros((self.num_inducing, self.num_inducing, self.output_dim))
            .to(self.dtype)
            .to(self.device)
        )
        li, lj = torch.tril_indices(self.num_inducing, self.num_inducing)
        self.q_sqrt[li, lj] = self.q_sqrt_tri
        
        SK = torch.einsum("nmd, bmd->dnb", self.q_sqrt, self.q_sqrt) - Ku
        
        B = torch.einsum("dmb, dnb -> dmn", SK, A)
        
        delta_cov = torch.sum(A * B.permute(0, 2, 1), -1)
        K = torch.diagonal(Kf, dim1=1, dim2=2) + delta_cov
        
        # Add layer noise to variance
        if self.log_layer_noise is not None:
            K = K + torch.exp(self.log_layer_noise).unsqueeze(-1)

        # Add mean function
        if self.mean_function is not None:
            mean = mean + self.mean_function(X)
            
        if return_prior_samples:
            return mean, K, f[:, :X.shape[0], :]
        return mean, K
//  [python]
def _cholesky(a: Numeric):
    return torch.cholesky(a)
//  [python]
def _cholesky(a: Numeric):
    return torch.linalg.cholesky(a)
//  [python]
def _invchol(x, n, ret_chol=False):
    l = torch.cholesky(x)
    eye = torch.eye(n, out=x.new(n, n))
    l_inv = torch.solve(eye, l)
    if not ret_chol:
        return l_inv, None
    return l_inv, l
//  [python]
def _invchol(x, n, ret_chol=False):
    l = torch.linalg.cholesky(x)
    eye = torch.eye(n, out=x.new(n, n))
    l_inv = torch.linalg.solve_triangular(l, eye, upper=False)
    if not ret_chol:
        return l_inv, None
    return l_inv, l
//  [python]
def logm_invB_A(B, A):
#     inputs: A/B.shape = [hxwxd, 3, 3]
#     output: shape = [hxwxd, 3, 3]
    G = torch.cholesky(B)
    inv_G = torch.inverse(G)
    W = torch.einsum("...ij,...jk,...lk->...il", inv_G, A, inv_G)
    lamda, Q = se.apply(W)
    log_lamda = torch.diag_embed(torch.log(lamda))
    V = torch.einsum('...ji,...jk->...ik', inv_G, Q)
    inv_V = torch.inverse(V)
    return torch.einsum('...ij,...jk,...kl->...il', V, log_lamda, inv_V)
//  [python]
def logm_invB_A(B, A):
#     inputs: A/B.shape = [hxwxd, 3, 3]
#     output: shape = [hxwxd, 3, 3]
    G = torch.linalg.cholesky(B)
    inv_G = torch.inverse(G)
    W = torch.einsum("...ij,...jk,...lk->...il", inv_G, A, inv_G)
    lamda, Q = se.apply(W)
    log_lamda = torch.diag_embed(torch.log(lamda))
    V = torch.einsum('...ji,...jk->...ik', inv_G, Q)
    inv_V = torch.inverse(V)
    return torch.einsum('...ij,...jk,...kl->...il', V, log_lamda, inv_V)
//  [python]
def inner(self, x, u, v, keepdim=False):
        l = torch.cholesky(x)
        x_inv_u = torch.cholesky_solve(u, l)
        x_inv_v = torch.cholesky_solve(v, l)
        return _trace(x_inv_u @ x_inv_v, keepdim=keepdim)
//  [python]
def inner(self, x, u, v, keepdim=False):
        l = torch.linalg.cholesky(x)
        x_inv_u = torch.cholesky_solve(u, l)
        x_inv_v = torch.cholesky_solve(v, l)
        return _trace(x_inv_u @ x_inv_v, keepdim=keepdim)
//  [python]
def estimate_c_tilde_matrix(data: torch.Tensor, data_lagged: torch.Tensor, reversible = True):
    """ This method can be applied to compute the C\tilde matrix from time-instant and time-lagged data.

    Parameters
    ----------
    data : torch.Tensor
        The time-instant data of shape [batch_size, num_basis].
    data_lagged : torch.Tensor
        The time-lagged data of shape [batch_size, num_basis].

    Returns
    -------
    C_tilde : torch.Tensor
        The C tilde matrix of shape [num_basis, num_basis].
    """
    
    cov_00, cov_01, cov_11 = compute_covariance_matrix(data, data_lagged)
    _, cov_10, _ = compute_covariance_matrix(data_lagged, data)
    if reversible:
        c_0 = 0.5 * (cov_00 + cov_11)
        c_1 = 0.5 * (cov_01 + cov_10)
    else:
        c_0 = cov_00
        c_1 = cov_01

    L = torch.cholesky(c_0)
    L_inv = torch.inverse(L)

    try:
    # C_tilde = torch.chain_matmul(L_inv, c_1, L_inv.t())
        C_tilde = torch.linalg.multi_dot((L_inv, c_1, L_inv.t()))
    except:
        C_tilde = torch.chain_matmul(L_inv, c_1, L_inv.t())

    return C_tilde
//  [python]
def estimate_c_tilde_matrix(data: torch.Tensor, data_lagged: torch.Tensor, reversible = True):
    """ This method can be applied to compute the C\tilde matrix from time-instant and time-lagged data.

    Parameters
    ----------
    data : torch.Tensor
        The time-instant data of shape [batch_size, num_basis].
    data_lagged : torch.Tensor
        The time-lagged data of shape [batch_size, num_basis].

    Returns
    -------
    C_tilde : torch.Tensor
        The C tilde matrix of shape [num_basis, num_basis].
    """
    
    cov_00, cov_01, cov_11 = compute_covariance_matrix(data, data_lagged)
    _, cov_10, _ = compute_covariance_matrix(data_lagged, data)
    if reversible:
        c_0 = 0.5 * (cov_00 + cov_11)
        c_1 = 0.5 * (cov_01 + cov_10)
    else:
        c_0 = cov_00
        c_1 = cov_01

    L = torch.linalg.cholesky(c_0)
    L_inv = torch.inverse(L)

    try:
    # C_tilde = torch.chain_matmul(L_inv, c_1, L_inv.t())
        C_tilde = torch.linalg.multi_dot((L_inv, c_1, L_inv.t()))
    except:
        C_tilde = torch.chain_matmul(L_inv, c_1, L_inv.t())

    return C_tilde
//  [python]
def safe_cholesky(covariance_matrix, jitter=1e-6):
    """Perform a safe cholesky decomposition of the covariance matrix.

    If cholesky decomposition raises Runtime error, it adds jitter to the covariance
    matrix.

    Parameters
    ----------
    covariance_matrix: torch.Tensor.
        Tensor with dimensions batch x dim x dim.
    jitter: float, optional.
        Jitter to add to the covariance matrix.
    """
    try:
        return torch.cholesky(covariance_matrix)
    except RuntimeError:
        dim = covariance_matrix.shape[-1]
        if jitter > 1:
            # When jitter is too big, then there is some numerical issue and this avoids
            # stack overflow.
            warnings.warn("Jitter too big. Maybe some numerical issue somewhere.")
            return torch.eye(dim)
        return safe_cholesky(
            covariance_matrix + jitter * torch.eye(dim), jitter=10 * jitter
        )
//  [python]
def safe_cholesky(covariance_matrix, jitter=1e-6):
    """Perform a safe cholesky decomposition of the covariance matrix.

    If cholesky decomposition raises Runtime error, it adds jitter to the covariance
    matrix.

    Parameters
    ----------
    covariance_matrix: torch.Tensor.
        Tensor with dimensions batch x dim x dim.
    jitter: float, optional.
        Jitter to add to the covariance matrix.
    """
    try:
        return torch.linalg.cholesky(covariance_matrix)
    except RuntimeError:
        dim = covariance_matrix.shape[-1]
        if jitter > 1:
            # When jitter is too big, then there is some numerical issue and this avoids
            # stack overflow.
            warnings.warn("Jitter too big. Maybe some numerical issue somewhere.")
            return torch.eye(dim)
        return safe_cholesky(
            covariance_matrix + jitter * torch.eye(dim), jitter=10 * jitter
        )
//  [python]
def test_batch_cg(self):
        batch = 5
        size = 100
        matrix = torch.randn(batch, size, size, dtype=torch.float64)
        matrix = matrix.matmul(matrix.mT)
        matrix.div_(matrix.norm())
        matrix.add_(torch.eye(matrix.size(-1), dtype=torch.float64).mul_(1e-1))

        rhs = torch.randn(batch, size, 50, dtype=torch.float64)
        solves = linear_cg(matrix.matmul, rhs=rhs, max_iter=size)

        # 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))
//  [python]
def test_batch_cg(self):
        batch = 5
        size = 100
        matrix = torch.randn(batch, size, size, dtype=torch.float64)
        matrix = matrix.matmul(matrix.mT)
        matrix.div_(matrix.norm())
        matrix.add_(torch.eye(matrix.size(-1), dtype=torch.float64).mul_(1e-1))

        rhs = torch.randn(batch, size, 50, dtype=torch.float64)
        solves = linear_cg(matrix.matmul, rhs=rhs, max_iter=size)

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

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

        # 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
        eigs = torch.linalg.eigvalsh(matrix)
        for i in range(5):
            approx_eigs = torch.linalg.eigvalsh(t_mats[i])
            self.assertTrue(torch.allclose(eigs, approx_eigs, atol=1e-3, rtol=1e-4))
//  [python]
def test_cg_with_tridiag(self):
        size = 10
        matrix = torch.randn(size, size, dtype=torch.float64)
        matrix = matrix.matmul(matrix.mT)
        matrix.div_(matrix.norm())
        matrix.add_(torch.eye(matrix.size(-1), dtype=torch.float64).mul_(1e-1))

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

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

        # Check tridiag
        eigs = torch.linalg.eigvalsh(matrix)
        for i in range(5):
            approx_eigs = torch.linalg.eigvalsh(t_mats[i])
            self.assertTrue(torch.allclose(eigs, approx_eigs, atol=1e-3, rtol=1e-4))
//  [python]
def belief_covariance(self, covariance: types.CovarianceTorch):
        self._belief_scale_tril = torch.cholesky(covariance)
        self._belief_covariance = covariance
//  [python]
def belief_covariance(self, covariance: types.CovarianceTorch):
        self._belief_scale_tril = torch.linalg.cholesky(covariance)
        self._belief_covariance = covariance
//  [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.mT)
        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 = torch.linalg.eigvalsh(matrix[i])
            for j in range(8):
                approx_eigs = torch.linalg.eigvalsh(t_mats[j, i])
                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.mT)
        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.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 = torch.linalg.eigvalsh(matrix[i])
            for j in range(8):
                approx_eigs = torch.linalg.eigvalsh(t_mats[j, i])
                self.assertTrue(torch.allclose(eigs, approx_eigs, atol=1e-3, rtol=1e-4))
//  [python]
def test_batch_cg_init(self):
        batch = 5
        size = 100
        matrix = torch.randn(batch, size, size, dtype=torch.float64)
        matrix = matrix.matmul(matrix.mT)
        matrix.div_(matrix.norm())
        matrix.add_(torch.eye(matrix.size(-1), dtype=torch.float64).mul_(1e-1))

        # Initial solve
        rhs = torch.randn(batch, size, 50, dtype=torch.float64)
        solves = linear_cg(matrix.matmul, rhs=rhs, max_iter=size, max_tridiag_iter=0)

        # Initialize with solve
        solves_with_init = linear_cg(matrix.matmul, rhs=rhs, max_iter=1, initial_guess=solves, max_tridiag_iter=0)

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

        # Initial solve
        rhs = torch.randn(batch, size, 50, dtype=torch.float64)
        solves = linear_cg(matrix.matmul, rhs=rhs, max_iter=size, max_tridiag_iter=0)

        # Initialize with solve
        solves_with_init = linear_cg(matrix.matmul, rhs=rhs, max_iter=1, initial_guess=solves, max_tridiag_iter=0)

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