jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def fit(self, dataset: torch.utils.data.DataLoader) -> None:
        errors = []

        # Compute mean and covariance over the entire validation dataset
        counter = 0
        for b_inputs, b_targets in dataset:
            b_inputs = tuple(b_inp.to(self.dummy.device) for b_inp in b_inputs)
            b_targets = tuple(b_tar.to(self.dummy.device) for b_tar in b_targets)
            with torch.no_grad():
                pred = self.model(b_inputs)

            target, = b_targets

            error = target[:, -self.offset:] - pred[:, -self.offset:]
            for j in range(error.shape[0]):
                for j2 in range(error.shape[1]):
                    index = counter + j + j2
                    if len(errors) < index + 1:
                        errors.append([])
                    errors[counter + j + j2].append(error[j, j2])
            counter += error.shape[0]

        # We will have up to offset predictions for each point. We use the mean as the final prediction
        errors = [sum(error) / len(error) for error in errors]
        errors = torch.stack(errors, dim=0)

        mean = torch.mean(errors, dim=0)
        errors -= mean

        cov = torch.matmul(errors.T, errors)
        cov /= errors.shape[0] - 1

        # Add a small epsilon to the diagonal of the matrix to make it non-singular
        cov.diagonal().add_(1e-5)

        # This construction ensures that the resulting precision matrix is pos. semi-definite, even if the condition
        # number of the cov matrix is large
        cholesky = torch.cholesky(cov)
        precision = cov
        torch.cholesky_inverse(cholesky, out=precision)

        self.register_buffer('mean', mean)
        self.register_buffer('precision', precision)
//  [python]
def fit(self, dataset: torch.utils.data.DataLoader) -> None:
        errors = []

        # Compute mean and covariance over the entire validation dataset
        counter = 0
        for b_inputs, b_targets in dataset:
            b_inputs = tuple(b_inp.to(self.dummy.device) for b_inp in b_inputs)
            b_targets = tuple(b_tar.to(self.dummy.device) for b_tar in b_targets)
            with torch.no_grad():
                pred = self.model(b_inputs)

            target, = b_targets

            error = target[:, -self.offset:] - pred[:, -self.offset:]
            for j in range(error.shape[0]):
                for j2 in range(error.shape[1]):
                    index = counter + j + j2
                    if len(errors) < index + 1:
                        errors.append([])
                    errors[counter + j + j2].append(error[j, j2])
            counter += error.shape[0]

        # We will have up to offset predictions for each point. We use the mean as the final prediction
        errors = [sum(error) / len(error) for error in errors]
        errors = torch.stack(errors, dim=0)

        mean = torch.mean(errors, dim=0)
        errors -= mean

        cov = torch.matmul(errors.T, errors)
        cov /= errors.shape[0] - 1

        # Add a small epsilon to the diagonal of the matrix to make it non-singular
        cov.diagonal().add_(1e-5)

        # This construction ensures that the resulting precision matrix is pos. semi-definite, even if the condition
        # number of the cov matrix is large
        cholesky = torch.linalg.cholesky(cov)
        precision = cov
        torch.cholesky_inverse(cholesky, out=precision)

        self.register_buffer('mean', mean)
        self.register_buffer('precision', precision)
//  [python]
def fit(
        self,
        alpha: float,
        gamma: float,
        train_embeddings: torch.Tensor,
        train_y: torch.Tensor,
        device: torch.device = torch.device("cpu"),
    ) -> None:
        # Save values that we need later
        self.alpha = nn.Parameter(torch.tensor(alpha), requires_grad=False)
        self.gamma = nn.Parameter(torch.tensor(gamma), requires_grad=False)
        self.train_embeddings = nn.Parameter(train_embeddings, requires_grad=False)

        # Precompute what can be precomputed (cf. to SI)
        self.intercept = nn.Parameter(torch.mean(train_y), requires_grad=False)
        self.scale = nn.Parameter(torch.std(train_y), requires_grad=False)
        train_y = (train_y - self.intercept) / self.scale
        train_kernel = _compute_kernel(
            self.train_embeddings.data,
            self.train_embeddings.data,
            self.gamma.data,
            device=device,
        ) + self.alpha.to(device) * torch.eye(self.train_embeddings.shape[0], device=device)
        self.train_kernel_cholesky = nn.Parameter(
            torch.cholesky(train_kernel).cpu(), requires_grad=False
        )
        self.dual_coefficients = nn.Parameter(
            torch.cholesky_solve(train_y.unsqueeze(1), self.train_kernel_cholesky)[:, 0].cpu(),
            requires_grad=False,
        )
//  [python]
def fit(
        self,
        alpha: float,
        gamma: float,
        train_embeddings: torch.Tensor,
        train_y: torch.Tensor,
        device: torch.device = torch.device("cpu"),
    ) -> None:
        # Save values that we need later
        self.alpha = nn.Parameter(torch.tensor(alpha), requires_grad=False)
        self.gamma = nn.Parameter(torch.tensor(gamma), requires_grad=False)
        self.train_embeddings = nn.Parameter(train_embeddings, requires_grad=False)

        # Precompute what can be precomputed (cf. to SI)
        self.intercept = nn.Parameter(torch.mean(train_y), requires_grad=False)
        self.scale = nn.Parameter(torch.std(train_y), requires_grad=False)
        train_y = (train_y - self.intercept) / self.scale
        train_kernel = _compute_kernel(
            self.train_embeddings.data,
            self.train_embeddings.data,
            self.gamma.data,
            device=device,
        ) + self.alpha.to(device) * torch.eye(self.train_embeddings.shape[0], device=device)
        self.train_kernel_cholesky = nn.Parameter(
            torch.linalg.cholesky(train_kernel).cpu(), requires_grad=False
        )
        self.dual_coefficients = nn.Parameter(
            torch.cholesky_solve(train_y.unsqueeze(1), self.train_kernel_cholesky)[:, 0].cpu(),
            requires_grad=False,
        )
//  [python]
def is_pos_def(x):
    if torch.equal(x, x.t()):
        try:
            torch.cholesky(x)
            return True
        except RuntimeError:
            return False
    else:
        return False
//  [python]
def is_pos_def(x):
    if torch.equal(x, x.t()):
        try:
            torch.linalg.cholesky(x)
            return True
        except RuntimeError:
            return False
    else:
        return False
//  [python]
def fasterquant(
        self, blocksize=128, percdamp=.01, groupsize=-1
    ):
        W = self.layer.weight.data.clone()
        
        W = W.float()

        if not self.quantizer.ready():
            if self.fp_features > 0:
                # We calculate the quantization parameters only on the integer part of the weights
                # Then, we permute the columns so the integer part is on the left
                self.quantizer.find_params(W[:, self.int_indices])
                W = W[:, self.col_perm]
                self.H = self.H[self.col_perm, :][:, self.col_perm]
            else:
                self.quantizer.find_params(W)

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

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

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

        for i1 in range(0, self.columns, blocksize):
            i2 = min(i1 + blocksize, self.columns)
            if i1 >= self.columns - self.fp_features and self.fp_features > 0:
                Q[:, i1:i2] = W[:, i1:i2].clone()
                continue
            count = i2 - i1

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

            for i in range(count):
                if i + i1 >= self.columns - self.fp_features and self.fp_features > 0:
                    Q1[:, i] = W1[:, i]
                    continue
                w = W1[:, i]
                d = Hinv1[i, i]

                if groupsize != -1:
                    if (i1 + i) % groupsize == 0:
                        self.quantizer.find_params(W[:, (i1 + i):(i1 + i + groupsize)])
                
                q = self.quantizer.quantize(w.unsqueeze(1)).flatten()
                
                Q1[:, i] = q
                Losses1[:, i] = (w - q) ** 2 / d ** 2

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

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

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

        torch.cuda.synchronize()
        self.layer.weight.data = Q.reshape(self.layer.weight.shape).to(self.layer.weight.data.dtype)

        # Permut Back the weights
        if self.fp_features > 0:
            self.layer.weight.data = self.layer.weight.data[:, self.inv_col_perm]
//  [python]
def fasterquant(
        self, blocksize=128, percdamp=.01, groupsize=-1
    ):
        W = self.layer.weight.data.clone()
        
        W = W.float()

        if not self.quantizer.ready():
            if self.fp_features > 0:
                # We calculate the quantization parameters only on the integer part of the weights
                # Then, we permute the columns so the integer part is on the left
                self.quantizer.find_params(W[:, self.int_indices])
                W = W[:, self.col_perm]
                self.H = self.H[self.col_perm, :][:, self.col_perm]
            else:
                self.quantizer.find_params(W)

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

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

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

        for i1 in range(0, self.columns, blocksize):
            i2 = min(i1 + blocksize, self.columns)
            if i1 >= self.columns - self.fp_features and self.fp_features > 0:
                Q[:, i1:i2] = W[:, i1:i2].clone()
                continue
            count = i2 - i1

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

            for i in range(count):
                if i + i1 >= self.columns - self.fp_features and self.fp_features > 0:
                    Q1[:, i] = W1[:, i]
                    continue
                w = W1[:, i]
                d = Hinv1[i, i]

                if groupsize != -1:
                    if (i1 + i) % groupsize == 0:
                        self.quantizer.find_params(W[:, (i1 + i):(i1 + i + groupsize)])
                
                q = self.quantizer.quantize(w.unsqueeze(1)).flatten()
                
                Q1[:, i] = q
                Losses1[:, i] = (w - q) ** 2 / d ** 2

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

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

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

        torch.cuda.synchronize()
        self.layer.weight.data = Q.reshape(self.layer.weight.shape).to(self.layer.weight.data.dtype)

        # Permut Back the weights
        if self.fp_features > 0:
            self.layer.weight.data = self.layer.weight.data[:, self.inv_col_perm]
//  [python]
def test_utilities(head_size):
    size = 5
    cov = torch.randn(size, size)
    cov = torch.mm(cov, cov.t())

    mask = torch.ones(size, size)
    mask[head_size:, head_size:] = 0.0
    mask.view(-1)[:: size + 1][head_size:] = 1.0
    arrowhead_full = mask * cov
    expected = torch.flip(
        torch.cholesky(torch.flip(arrowhead_full, (-2, -1))), (-2, -1)
    )
    # test if those flip ops give expected upper triangular values
    assert_close(expected.triu(), expected)
    assert_close(expected.matmul(expected.t()), arrowhead_full)

    # test sqrt
    arrowhead = SymmArrowhead(cov[:head_size], cov.diag()[head_size:])
    actual = sqrt(arrowhead)
    assert_close(actual.top, expected[:head_size])
    assert_close(actual.bottom_diag, expected.diag()[head_size:])

    # test triu_inverse
    expected = expected.inverse()
    actual = triu_inverse(actual)
    assert_close(actual.top, expected[:head_size])
    assert_close(actual.bottom_diag, expected.diag()[head_size:])

    # test triu_matvecmul
    v = torch.randn(size)
    assert_close(triu_matvecmul(actual, v), expected.matmul(v))
    assert_close(triu_matvecmul(actual, v, transpose=True), expected.t().matmul(v))

    # test triu_gram
    actual = triu_gram(actual)
    expected = (
        arrowhead_full.inverse()
        if head_size > 0
        else arrowhead_full.diag().reciprocal()
    )
    assert_close(actual, expected)
//  [python]
def test_utilities(head_size):
    size = 5
    cov = torch.randn(size, size)
    cov = torch.mm(cov, cov.t())

    mask = torch.ones(size, size)
    mask[head_size:, head_size:] = 0.0
    mask.view(-1)[:: size + 1][head_size:] = 1.0
    arrowhead_full = mask * cov
    expected = torch.flip(
        torch.linalg.cholesky(torch.flip(arrowhead_full, (-2, -1))), (-2, -1)
    )
    # test if those flip ops give expected upper triangular values
    assert_close(expected.triu(), expected)
    assert_close(expected.matmul(expected.t()), arrowhead_full)

    # test sqrt
    arrowhead = SymmArrowhead(cov[:head_size], cov.diag()[head_size:])
    actual = sqrt(arrowhead)
    assert_close(actual.top, expected[:head_size])
    assert_close(actual.bottom_diag, expected.diag()[head_size:])

    # test triu_inverse
    expected = expected.inverse()
    actual = triu_inverse(actual)
    assert_close(actual.top, expected[:head_size])
    assert_close(actual.bottom_diag, expected.diag()[head_size:])

    # test triu_matvecmul
    v = torch.randn(size)
    assert_close(triu_matvecmul(actual, v), expected.matmul(v))
    assert_close(triu_matvecmul(actual, v, transpose=True), expected.t().matmul(v))

    # test triu_gram
    actual = triu_gram(actual)
    expected = (
        arrowhead_full.inverse()
        if head_size > 0
        else arrowhead_full.diag().reciprocal()
    )
    assert_close(actual, expected)
//  [python]
def forward(self, x, L, mass, evals, evecs):

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

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

        if self.method == "spectral":

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

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

            # Transform back to per-vertex
            x_diffuse = from_basis(x_diffuse_spec, evecs)

        elif self.method == "implicit_dense":
            V = x.shape[-2]

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

            # Factor the system
            cholesky_factors = torch.cholesky(mat_dense)

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

        else:
            raise ValueError("unrecognized method")

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

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

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

        if self.method == "spectral":

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

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

            # Transform back to per-vertex
            x_diffuse = from_basis(x_diffuse_spec, evecs)

        elif self.method == "implicit_dense":
            V = x.shape[-2]

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

            # Factor the system
            cholesky_factors = torch.linalg.cholesky(mat_dense)

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

        else:
            raise ValueError("unrecognized method")

        return x_diffuse
//  [python]
def fit(self, C, v):
        ## iterates are in c, the number of concepts
        c = C.shape[0]

        ## size: c x c
        Q = 2*C @ C.T + (torch.eye(c)*self.rho).to(self.device)

        # factor Q for quicker solve -- this is critical.
        Q_cho = torch.cholesky(Q)

        # iterates, size: c x batch
        x = torch.randn((c, v.shape[0])).to(self.device)
        z = torch.randn((c, v.shape[0])).to(self.device)
        u = torch.randn((c, v.shape[0])).to(self.device)

        for ix in range(self.max_iter):
            z_old = z
            x, z, u = self.step(C@v.T, Q_cho, z, u)

            res_prim = torch.linalg.norm(x-z, dim=0)
            res_dual = torch.linalg.norm(self.rho*(z-z_old), dim=0)

            if (res_prim.max() < self.tol) and (res_dual.max() < self.tol):
                break
        if self.verbose:
            print("Stopping at iteration {}".format(ix))
            print("Prime Residual, r_k: {}".format(res_prim.mean()))
            print("Dual Residual, s_k: {}".format(res_dual.mean()))
        return z.T
//  [python]
def fit(self, C, v):
        ## iterates are in c, the number of concepts
        c = C.shape[0]

        ## size: c x c
        Q = 2*C @ C.T + (torch.eye(c)*self.rho).to(self.device)

        # factor Q for quicker solve -- this is critical.
        Q_cho = torch.linalg.cholesky(Q)

        # iterates, size: c x batch
        x = torch.randn((c, v.shape[0])).to(self.device)
        z = torch.randn((c, v.shape[0])).to(self.device)
        u = torch.randn((c, v.shape[0])).to(self.device)

        for ix in range(self.max_iter):
            z_old = z
            x, z, u = self.step(C@v.T, Q_cho, z, u)

            res_prim = torch.linalg.norm(x-z, dim=0)
            res_dual = torch.linalg.norm(self.rho*(z-z_old), dim=0)

            if (res_prim.max() < self.tol) and (res_dual.max() < self.tol):
                break
        if self.verbose:
            print("Stopping at iteration {}".format(ix))
            print("Prime Residual, r_k: {}".format(res_prim.mean()))
            print("Dual Residual, s_k: {}".format(res_dual.mean()))
        return z.T
//  [python]
def mat_inv(x):
    u = torch.cholesky(x)
    return torch.cholesky_inverse(u)
//  [python]
def mat_inv(x):
    u = torch.linalg.cholesky(x)
    return torch.cholesky_inverse(u)
//  [python]
def to_cholesky(diag_vector=None, off_diag_vector=None,
                L=None, cov_matrix=None):
    """
    Compute Cholesky matrix
    Args:
        diag_vector: diag elements in a vector
        off_diag_vector: off-diagonal elements in a vector
        L: Cholesky matrix
        cov_matrix: Covariance matrix

    Returns:
        Cholesky matrix L
    """
    if L is not None:
        pass
    elif cov_matrix is None:
        assert diag_vector is not None and off_diag_vector is not None
        L = build_lower_matrix(diag_vector, off_diag_vector)
    elif diag_vector is None and off_diag_vector is None:
        L = torch.cholesky(cov_matrix)
    else:
        raise RuntimeError("Unexpected behaviours")
    return L
//  [python]
def to_cholesky(diag_vector=None, off_diag_vector=None,
                L=None, cov_matrix=None):
    """
    Compute Cholesky matrix
    Args:
        diag_vector: diag elements in a vector
        off_diag_vector: off-diagonal elements in a vector
        L: Cholesky matrix
        cov_matrix: Covariance matrix

    Returns:
        Cholesky matrix L
    """
    if L is not None:
        pass
    elif cov_matrix is None:
        assert diag_vector is not None and off_diag_vector is not None
        L = build_lower_matrix(diag_vector, off_diag_vector)
    elif diag_vector is None and off_diag_vector is None:
        L = torch.linalg.cholesky(cov_matrix)
    else:
        raise RuntimeError("Unexpected behaviours")
    return L
//  [python]
def decompose_kernels(gpr,pred_matrix,kernels):
    '''
    A function to calculate different kernels contribution to final prediction

    ------------------Inputs----------------------------
    gpr: an optimized pyro GP regression model
    pred_matrix: a torch tensor of prediction matrix containing the data points for prediction
    kernels: a list of pyro kernels for decomposition

    ------------------Outputs---------------------------
    output: a list of tuples, each tuple contains the mean and covariance of the prediction for each kernel 
    '''
    N = len(gpr.X)
    M = pred_matrix.size(0)
    f_loc = gpr.y - gpr.mean_function(gpr.X)
    latent_shape = f_loc.shape[:-1]
    loc_shape = latent_shape + (M,)
    f_loc = f_loc.permute(-1, *range(len(latent_shape)))
    f_loc_2D = f_loc.reshape(N, -1)
    loc_shape = latent_shape + (M,)
    v_2D = f_loc_2D
    Kff = gpr.kernel(gpr.X).contiguous()
    if gpr.noise.dim() <=1:
        Kff.view(-1)[:: N + 1] += gpr.jitter + gpr.noise   # add noise to the diagonal
    elif gpr.noise.dim() ==2:
        Kff = Kff + gpr.noise
        Kff.view(-1)[:: N + 1] += gpr.jitter 
    Lff = torch.cholesky(Kff)
    
    output = []
    for kernel in kernels:
        Kfs = kernel(gpr.X, pred_matrix)
        pack = torch.cat((f_loc_2D, Kfs), dim=1)
        Lffinv_pack = pack.triangular_solve(Lff, upper=False)[0]
        W = Lffinv_pack[:, f_loc_2D.size(1):f_loc_2D.size(1) + M].t()
        Qss = W.matmul(W.t())
        v_2D = Lffinv_pack[:, :f_loc_2D.size(1)]
        loc = W.matmul(v_2D).t().reshape(loc_shape)
        Kss = kernel(pred_matrix)
        cov = Kss - Qss
        output.append((loc, cov))
        
    return output
//  [python]
def decompose_kernels(gpr,pred_matrix,kernels):
    '''
    A function to calculate different kernels contribution to final prediction

    ------------------Inputs----------------------------
    gpr: an optimized pyro GP regression model
    pred_matrix: a torch tensor of prediction matrix containing the data points for prediction
    kernels: a list of pyro kernels for decomposition

    ------------------Outputs---------------------------
    output: a list of tuples, each tuple contains the mean and covariance of the prediction for each kernel 
    '''
    N = len(gpr.X)
    M = pred_matrix.size(0)
    f_loc = gpr.y - gpr.mean_function(gpr.X)
    latent_shape = f_loc.shape[:-1]
    loc_shape = latent_shape + (M,)
    f_loc = f_loc.permute(-1, *range(len(latent_shape)))
    f_loc_2D = f_loc.reshape(N, -1)
    loc_shape = latent_shape + (M,)
    v_2D = f_loc_2D
    Kff = gpr.kernel(gpr.X).contiguous()
    if gpr.noise.dim() <=1:
        Kff.view(-1)[:: N + 1] += gpr.jitter + gpr.noise   # add noise to the diagonal
    elif gpr.noise.dim() ==2:
        Kff = Kff + gpr.noise
        Kff.view(-1)[:: N + 1] += gpr.jitter 
    Lff = torch.linalg.cholesky(Kff)
    
    output = []
    for kernel in kernels:
        Kfs = kernel(gpr.X, pred_matrix)
        pack = torch.cat((f_loc_2D, Kfs), dim=1)
        Lffinv_pack = pack.triangular_solve(Lff, upper=False)[0]
        W = Lffinv_pack[:, f_loc_2D.size(1):f_loc_2D.size(1) + M].t()
        Qss = W.matmul(W.t())
        v_2D = Lffinv_pack[:, :f_loc_2D.size(1)]
        loc = W.matmul(v_2D).t().reshape(loc_shape)
        Kss = kernel(pred_matrix)
        cov = Kss - Qss
        output.append((loc, cov))
        
    return output
//  [python]
def _solve_weighted_least_square_regression(kernelSHAP_weights: FloatTensor,
                                            coalitions: BoolTensor,
                                            regression_target: FloatTensor | Tensor,
                                            ) -> FloatTensor:
    weighted_regression_target = regression_target * kernelSHAP_weights
    ZtWvx = coalitions.t() @ weighted_regression_target
    L = torch.cholesky(coalitions.t() @ (coalitions * kernelSHAP_weights))

    return torch.cholesky_solve(ZtWvx, L).detach()
//  [python]
def _solve_weighted_least_square_regression(kernelSHAP_weights: FloatTensor,
                                            coalitions: BoolTensor,
                                            regression_target: FloatTensor | Tensor,
                                            ) -> FloatTensor:
    weighted_regression_target = regression_target * kernelSHAP_weights
    ZtWvx = coalitions.t() @ weighted_regression_target
    L = torch.linalg.cholesky(coalitions.t() @ (coalitions * kernelSHAP_weights))

    return torch.cholesky_solve(ZtWvx, L).detach()
//  [python]
def negative_log_likelihood(self, **kwargs):

        X = self.dkl_nn.get_intermediate_representation(self.X, self.nn_feature_level).view(10, -1)

        y_num, y_dimension = self.Y.shape
        Sigma = self.kernel(X, X) + self.log_beta.exp().pow(-1) * torch.eye(
            X.size(0)).to(self.device) + self.jitter * torch.eye(X.size(0)).to(self.device)

        L = torch.cholesky(Sigma)
        # option 1 (use this if torch supports)
        Gamma, _ = torch.triangular_solve(self.Y, L, upper=False)
        # option 2
        # gamma = L.inverse() @ Y       # we can use this as an alternative because L is a lower triangular matrix.

        nll = 0.5 * (Gamma ** 2).sum() + L.diag().log().sum() * y_dimension  \
            + 0.5 * y_num * torch.log(2 * torch.tensor(self.pi)) * y_dimension

        loss_dict = {'all_loss_all': nll.item()}
        return nll, loss_dict
//  [python]
def negative_log_likelihood(self, **kwargs):

        X = self.dkl_nn.get_intermediate_representation(self.X, self.nn_feature_level).view(10, -1)

        y_num, y_dimension = self.Y.shape
        Sigma = self.kernel(X, X) + self.log_beta.exp().pow(-1) * torch.eye(
            X.size(0)).to(self.device) + self.jitter * torch.eye(X.size(0)).to(self.device)

        L = torch.linalg.cholesky(Sigma)
        # option 1 (use this if torch supports)
        Gamma, _ = torch.triangular_solve(self.Y, L, upper=False)
        # option 2
        # gamma = L.inverse() @ Y       # we can use this as an alternative because L is a lower triangular matrix.

        nll = 0.5 * (Gamma ** 2).sum() + L.diag().log().sum() * y_dimension  \
            + 0.5 * y_num * torch.log(2 * torch.tensor(self.pi)) * y_dimension

        loss_dict = {'all_loss_all': nll.item()}
        return nll, loss_dict
//  [python]
def forward(self, data: AtomicDataDict.Type, nk: Optional[int]=None) -> AtomicDataDict.Type:
        num_k = data[AtomicDataDict.KPOINT_KEY].shape[0]
        kpoints = data[AtomicDataDict.KPOINT_KEY]
        eigvals = []
        if nk is None:
            nk = num_k
        for i in range(int(np.ceil(num_k / nk))):
            data[AtomicDataDict.KPOINT_KEY] = kpoints[i*nk:(i+1)*nk]
            data = self.h2k(data)
            if self.overlap:
                data = self.s2k(data)
                chklowt = torch.cholesky(data[self.s_out_field])
                chklowtinv = torch.inverse(chklowt)
                data[self.h_out_field] = (chklowtinv @ data[self.h_out_field] @ torch.transpose(chklowtinv,dim0=1,dim1=2).conj())
            else:
                data[self.h_out_field] = data[self.h_out_field]
            
            eigvals.append(torch.linalg.eigvalsh(data[self.h_out_field]))
        data[self.out_field] = torch.cat(eigvals, dim=0)

        return data
//  [python]
def forward(self, data: AtomicDataDict.Type, nk: Optional[int]=None) -> AtomicDataDict.Type:
        num_k = data[AtomicDataDict.KPOINT_KEY].shape[0]
        kpoints = data[AtomicDataDict.KPOINT_KEY]
        eigvals = []
        if nk is None:
            nk = num_k
        for i in range(int(np.ceil(num_k / nk))):
            data[AtomicDataDict.KPOINT_KEY] = kpoints[i*nk:(i+1)*nk]
            data = self.h2k(data)
            if self.overlap:
                data = self.s2k(data)
                chklowt = torch.linalg.cholesky(data[self.s_out_field])
                chklowtinv = torch.linalg.inv(chklowt)
                data[self.h_out_field] = (chklowtinv @ data[self.h_out_field] @ torch.transpose(chklowtinv,dim0=1,dim1=2).conj())
            else:
                data[self.h_out_field] = data[self.h_out_field]
            
            eigvals.append(torch.linalg.eigvalsh(data[self.h_out_field]))
        data[self.out_field] = torch.cat(eigvals, dim=0)

        return data
//  [python]
def affine_invariant_distance_torch(x1, x2, diagonal_distance=False):
    """
    Compute the affine invariant distance between points on the SPD manifold

    Parameters
    ----------
    :param x1: set of SPD matrices      (N1 x d x d or b1 x ... x bk x N1 x d x d)
    :param x2: set of SPD matrices      (N2 x d x d or b1 x ... x bk x N1 x d x d)

    Optional parameters
    -------------------
    :param diagonal_distance: Whole distance matrix, or just the diagonal? If True, we must have `x1 == x2`.

    Returns
    -------
    :return: matrix of manifold affine-invariante distance between points in x1 and x2  (N1 x N2 or b1 x ... x bk x N1 x N2)
    """
    # If diag, x1 must be equal to x2 and we can return zeros.
    if diagonal_distance is True:
        shape = list(x2.shape)[:-2]
        shape.append(1)
        return torch.zeros(shape, dtype=x1.dtype)

    dim = x1.shape[-1]

    # Expand dimensions to compute all matrix-matrix distances
    x1 = x1.unsqueeze(-3)
    x2 = x2.unsqueeze(-4)

    # Method 1: compute x1^(-1)
    # x1_inv = torch.inverse(x1)  # method 1: uses inv(x1)x2

    # Method 2: uses the Cholesky decomposition to compute x1^(-0.5)
    x1_chol = torch.cholesky(x1)
    x1_chol_inv = torch.inverse(x1_chol)

    # Repeat x and y data along 1- and 0- dimensions to have ndata_x x ndata_y x dim x dim arrays
    # x1_inv = torch.cat(x2.shape[-3] * [x1_inv], dim=-3) # method 1
    x1_chol_inv = torch.cat(x2.shape[-3] * [x1_chol_inv], dim=-3)
    x2 = torch.cat(x1.shape[-4] * [x2], dim=-4)

    # Compute the distance between each pair of matrices
    # Method 1: compute x1\x2
    # x1_inv_x2 = torch.bmm(x1_inv.view(-1, dim, dim), x2.view(-1, dim, dim))  # method 1
    # x1_inv_x2 = torch.solve(x1, x2).solution.view(-1, dim, dim)

    # Method 2: compute x1^(-0.5)*x2*x1^(-0.5)
    # The advantage of this method is that the resulting matrix is symmetric => we can use symeig for the eigenvalues.
    x1_inv_x2_x1_inv = torch.bmm(torch.bmm(x1_chol_inv.view(-1, dim, dim), x2.view(-1, dim, dim)),
                                 x1_chol_inv.view(-1, dim, dim).transpose(-2, -1))

    # x1_inv_x2_x1_inv += 1e-10 * torch.rand(x1_inv_x2_x1_inv.shape)

    # Compute norm(logm(x1\x2))
    eig_values = torch.zeros(x1_inv_x2_x1_inv.shape[0], dim).to(device)
    for i in range(x1_inv_x2_x1_inv.shape[0]):
        eig_values[i] = torch.linalg.eigvalsh(x1_inv_x2_x1_inv[i])
        # eig_values[i] = torch.symeig(x1_inv_x2_x1_inv[i], eigenvectors=True).eigenvalues  # Eigenvalue True necessary for derivation
        # eigv[i] = torch.svd(x1_inv_x2_x1_inv[i], some=False).S
        # eigv[i] = torch.eig(x1_inv_x2[i]).eigenvalues[:, 0]  # Method 1: first/second column of eigenvalues = real/imaginary parts. Problem: the gradient of eig is not implemented.

    # Reshape
    shape = list(x2.shape)[:-2]
    shape.append(eig_values.shape[-1])
    eigv = eig_values.view(shape)

    logeigv = torch.log(eigv)
    logeigv2 = logeigv * logeigv
    sumlogeigv2 = torch.sum(logeigv2, dim=-1)
    return torch.sqrt(sumlogeigv2 + 1e-15).double()
//  [python]
def affine_invariant_distance_torch(x1, x2, diagonal_distance=False):
    """
    Compute the affine invariant distance between points on the SPD manifold

    Parameters
    ----------
    :param x1: set of SPD matrices      (N1 x d x d or b1 x ... x bk x N1 x d x d)
    :param x2: set of SPD matrices      (N2 x d x d or b1 x ... x bk x N1 x d x d)

    Optional parameters
    -------------------
    :param diagonal_distance: Whole distance matrix, or just the diagonal? If True, we must have `x1 == x2`.

    Returns
    -------
    :return: matrix of manifold affine-invariante distance between points in x1 and x2  (N1 x N2 or b1 x ... x bk x N1 x N2)
    """
    # If diag, x1 must be equal to x2 and we can return zeros.
    if diagonal_distance is True:
        shape = list(x2.shape)[:-2]
        shape.append(1)
        return torch.zeros(shape, dtype=x1.dtype)

    dim = x1.shape[-1]

    # Expand dimensions to compute all matrix-matrix distances
    x1 = x1.unsqueeze(-3)
    x2 = x2.unsqueeze(-4)

    # Method 1: compute x1^(-1)
    # x1_inv = torch.inverse(x1)  # method 1: uses inv(x1)x2

    # Method 2: uses the Cholesky decomposition to compute x1^(-0.5)
    x1_chol = torch.linalg.cholesky(x1)
    x1_chol_inv = torch.inverse(x1_chol)

    # Repeat x and y data along 1- and 0- dimensions to have ndata_x x ndata_y x dim x dim arrays
    # x1_inv = torch.cat(x2.shape[-3] * [x1_inv], dim=-3) # method 1
    x1_chol_inv = torch.cat(x2.shape[-3] * [x1_chol_inv], dim=-3)
    x2 = torch.cat(x1.shape[-4] * [x2], dim=-4)

    # Compute the distance between each pair of matrices
    # Method 1: compute x1\x2
    # x1_inv_x2 = torch.bmm(x1_inv.view(-1, dim, dim), x2.view(-1, dim, dim))  # method 1
    # x1_inv_x2 = torch.solve(x1, x2).solution.view(-1, dim, dim)

    # Method 2: compute x1^(-0.5)*x2*x1^(-0.5)
    # The advantage of this method is that the resulting matrix is symmetric => we can use symeig for the eigenvalues.
    x1_inv_x2_x1_inv = torch.bmm(torch.bmm(x1_chol_inv.view(-1, dim, dim), x2.view(-1, dim, dim)),
                                 x1_chol_inv.view(-1, dim, dim).transpose(-2, -1))

    # x1_inv_x2_x1_inv += 1e-10 * torch.rand(x1_inv_x2_x1_inv.shape)

    # Compute norm(logm(x1\x2))
    eig_values = torch.zeros(x1_inv_x2_x1_inv.shape[0], dim).to(device)
    for i in range(x1_inv_x2_x1_inv.shape[0]):
        eig_values[i] = torch.linalg.eigvalsh(x1_inv_x2_x1_inv[i])
        # eig_values[i] = torch.symeig(x1_inv_x2_x1_inv[i], eigenvectors=True).eigenvalues  # Eigenvalue True necessary for derivation
        # eigv[i] = torch.svd(x1_inv_x2_x1_inv[i], some=False).S
        # eigv[i] = torch.eig(x1_inv_x2[i]).eigenvalues[:, 0]  # Method 1: first/second column of eigenvalues = real/imaginary parts. Problem: the gradient of eig is not implemented.

    # Reshape
    shape = list(x2.shape)[:-2]
    shape.append(eig_values.shape[-1])
    eigv = eig_values.view(shape)

    logeigv = torch.log(eigv)
    logeigv2 = logeigv * logeigv
    sumlogeigv2 = torch.sum(logeigv2, dim=-1)
    return torch.sqrt(sumlogeigv2 + 1e-15).double()
//  [python]
def gmm(all_fea, pi, mu, all_output):    
    Cov = []
    dist = []
    log_probs = []
    
    for i in range(len(mu)):
        temp = all_fea - mu[i]
        predi = all_output[:,i].unsqueeze(dim=-1)
        Covi = torch.matmul(temp.t(), temp * predi.expand_as(temp)) / (predi.sum()) + args.epsilon * torch.eye(temp.shape[1]).cuda()
        try:
            chol = torch.cholesky(Covi)
        except RuntimeError:
            Covi += args.epsilon * torch.eye(temp.shape[1]).cuda() * 100
            chol = torch.cholesky(Covi)
        chol_inv = torch.inverse(chol)
        Covi_inv = torch.matmul(chol_inv.t(), chol_inv)
        logdet = torch.logdet(Covi)
        mah_dist = (torch.matmul(temp, Covi_inv) * temp).sum(dim=1)
        log_prob = -0.5*(Covi.shape[0] * np.log(2*math.pi) + logdet + mah_dist) + torch.log(pi)[i]
        Cov.append(Covi)
        log_probs.append(log_prob)
        dist.append(mah_dist)
    Cov = torch.stack(Cov, dim=0)
    dist = torch.stack(dist, dim=0).t()
    log_probs = torch.stack(log_probs, dim=0).t()
    zz = log_probs - torch.logsumexp(log_probs, dim=1, keepdim=True).expand_as(log_probs)
    gamma = torch.exp(zz)
    
    return zz, gamma
//  [python]
def gmm(all_fea, pi, mu, all_output):    
    Cov = []
    dist = []
    log_probs = []
    
    for i in range(len(mu)):
        temp = all_fea - mu[i]
        predi = all_output[:,i].unsqueeze(dim=-1)
        Covi = torch.matmul(temp.t(), temp * predi.expand_as(temp)) / (predi.sum()) + args.epsilon * torch.eye(temp.shape[1]).cuda()
        try:
            chol = torch.linalg.cholesky(Covi)
        except RuntimeError:
            Covi += args.epsilon * torch.eye(temp.shape[1]).cuda() * 100
            chol = torch.linalg.cholesky(Covi)
        chol_inv = torch.inverse(chol)
        Covi_inv = torch.matmul(chol_inv.t(), chol_inv)
        logdet = torch.logdet(Covi)
        mah_dist = (torch.matmul(temp, Covi_inv) * temp).sum(dim=1)
        log_prob = -0.5*(Covi.shape[0] * np.log(2*math.pi) + logdet + mah_dist) + torch.log(pi)[i]
        Cov.append(Covi)
        log_probs.append(log_prob)
        dist.append(mah_dist)
    Cov = torch.stack(Cov, dim=0)
    dist = torch.stack(dist, dim=0).t()
    log_probs = torch.stack(log_probs, dim=0).t()
    zz = log_probs - torch.logsumexp(log_probs, dim=1, keepdim=True).expand_as(log_probs)
    gamma = torch.exp(zz)
    
    return zz, gamma
//  [python]
def _solve_sytem(self, Atb: torch.Tensor, AtA: torch.Tensor) -> torch.Tensor:
        lower = torch.cholesky(AtA)
        return torch.cholesky_solve(Atb, lower).squeeze(2)
//  [python]
def _solve_sytem(self, Atb: torch.Tensor, AtA: torch.Tensor) -> torch.Tensor:
        lower = torch.linalg.cholesky(AtA)
        return torch.cholesky_solve(Atb, lower).squeeze(2)
//  [python]
def jitcholesky(A, jit=1e-6, jitbase=2):
    ridge = 0
    try:
        L = torch.cholesky(A)
    except RuntimeError:
        scale = A.diag().mean()
        if scale == torch.zeros(1):
            scale = torch.finfo().eps
        ridge = jit * scale
        done = False
        while not done:
            try:
                L = torch.cholesky(A + ridge * torch.eye(*A.size(), dtype=A.dtype))
                done = True
            except RuntimeError:
                ridge *= jitbase
            if ridge > scale:
                raise RuntimeError("cholesky was not successful!")
    return L, ridge
//  [python]
def jitcholesky(A, jit=1e-6, jitbase=2):
    ridge = 0
    try:
        L = torch.linalg.cholesky(A)
    except RuntimeError:
        scale = A.diag().mean()
        if scale == torch.zeros(1):
            scale = torch.finfo().eps
        ridge = jit * scale
        done = False
        while not done:
            try:
                L = torch.linalg.cholesky(A + ridge * torch.eye(*A.size(), dtype=A.dtype))
                done = True
            except RuntimeError:
                ridge *= jitbase
            if ridge > scale:
                raise RuntimeError("cholesky was not successful!")
    return L, ridge
//  [python]
def cholesky(a: ArrayLike):
    a = _atleast_float_1(a)
    return t.cholesky(a)
//  [python]
def cholesky(a: ArrayLike):
    a = _atleast_float_1(a)
    return torch.linalg.cholesky(a)
//  [python]
def test_helper_functions(self) -> None:
        for batch_shape in (torch.Size(), torch.Size([2])):
            tkwargs = {"device": self.device, "dtype": self.dtype}
            # M is borderline PSD
            M = torch.ones((*batch_shape, 2, 2), **tkwargs)
            with self.assertRaises(torch._C._LinAlgError):
                torch.cholesky(M)
            # This should work fine
            _ensure_psd_with_jitter(M)

            bad_M = torch.tensor([[1.0, 2.0], [2.0, 1.0]], **tkwargs).expand(
                (*batch_shape, 2, 2)
            )
            with self.assertRaises(NotPSDError):
                _ensure_psd_with_jitter(bad_M)
//  [python]
def test_helper_functions(self) -> None:
        for batch_shape in (torch.Size(), torch.Size([2])):
            tkwargs = {"device": self.device, "dtype": self.dtype}
            # M is borderline PSD
            M = torch.ones((*batch_shape, 2, 2), **tkwargs)
            with self.assertRaises(torch._C._LinAlgError):
                torch.linalg.cholesky(M)
            # This should work fine
            _ensure_psd_with_jitter(M)

            bad_M = torch.tensor([[1.0, 2.0], [2.0, 1.0]], **tkwargs).expand(
                (*batch_shape, 2, 2)
            )
            with self.assertRaises(NotPSDError):
                _ensure_psd_with_jitter(bad_M)
//  [python]
def merge_quadratic_term(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
        ):
            q = node.inputs[0].inputs[1].inputs[1].forward_value
            q_cholesky = torch.nn.Parameter(torch.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.inputs[0].inputs[0], 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"

            model.add_nodes(
                [
                    node_q_cholesky,
                    node_x_matmul_q_cholesky,
                    node_x_matmul_q_cholesky_sqr,
                    node_reduce_sum,
                ]
            )
            model.replace_node(node, node_reduce_sum)
//  [python]
def merge_quadratic_term(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
        ):
            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.inputs[0].inputs[0], 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"

            model.add_nodes(
                [
                    node_q_cholesky,
                    node_x_matmul_q_cholesky,
                    node_x_matmul_q_cholesky_sqr,
                    node_reduce_sum,
                ]
            )
            model.replace_node(node, node_reduce_sum)
//  [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.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)