jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _get_score(self, model, x, i, L=None, **kwargs):
        """
        Returns score: ∇_xs log p(xs)
        """
        if isinstance(i, float):
            i = torch.Tensor([i]).to(x)
        if i.shape[:-1] != x.shape[:-1]:
            i = i.view(*(1,) * len(x.shape)).expand_as(x[...,:1])

        beta_int = self.beta_fn.integral(i)
        std = (1 - torch.exp(-beta_int)).clamp(1e-5).sqrt()

        noise = model(x, i=i, **kwargs)

        if L is not None:
            # We have to compute the score using -Sigma.inv() @ noise / std
            # assuming noise~N(0, Sigma).
            # If `predict_gaussian_noise=False`, compute (LL^T).inv()
            # Else, we can simplify (LL^T).inv() @ L @ noise
            # to (L^T).inv() @ noise, where noise~N(0, I).
            # So we anyways have to do (L^T).inv(), and sometimes L.inv()
            if not self.predict_gaussian_noise:
                noise = torch.linalg.solve_triangular(L, noise, upper=False)
            noise = torch.linalg.solve_triangular(L.transpose(-1, -2), noise, upper=True)

        score = -noise / std
        return score
//  [python]
def update_precond_newton_math_(Q, v, h, step, tiny):
    # type: (Tensor, Tensor, Tensor, float, float) -> None
    """
    Update the classic Newton–Raphson type preconditioner P = Q'*Q with (v, h).
    """
    a = Q.mm(h)
    b = torch.triangular_solve(v, Q.t(), upper=False)[0]
    grad = torch.triu(a.mm(a.t()) - b.mm(b.t()))
    mu = step/(grad.abs().max() + tiny)
    Q.sub_(mu*grad.mm(Q))
//  [python]
def update_precond_newton_math_(Q, v, h, step, tiny):
    # type: (Tensor, Tensor, Tensor, float, float) -> None
    """
    Update the classic Newton–Raphson type preconditioner P = Q'*Q with (v, h).
    """
    a = Q.mm(h)
    b = torch.linalg.solve_triangular(Q.t(), v, upper=False)
    grad = torch.triu(a.mm(a.t()) - b.mm(b.t()))
    mu = step/(grad.abs().max() + tiny)
    Q.sub_(mu*grad.mm(Q))
//  [python]
def update_precond_dense(Q, dxs, dgs, step=0.01, _tiny=1.2e-38):
    # type: (Tensor, List[Tensor], List[Tensor], float, float) -> Tensor
    """
    update dense preconditioner P = Q^T*Q
    Q: Cholesky factor of preconditioner with positive diagonal entries
    dxs: list of perturbations of parameters
    dgs: list of perturbations of gradients
    step: update step size normalized to range [0, 1]
    _tiny: an offset to avoid division by zero
    """
    dx = torch.cat([torch.reshape(x, [-1, 1]) for x in dxs])
    dg = torch.cat([torch.reshape(g, [-1, 1]) for g in dgs])

    a = Q.mm(dg)
    b = torch.triangular_solve(dx, Q, upper=True, transpose=True)[0]

    grad = torch.triu(a.mm(a.t()) - b.mm(b.t()))
    step0 = step/(grad.abs().max() + _tiny)

    return Q - step0*grad.mm(Q)
//  [python]
def update_precond_dense(Q, dxs, dgs, step=0.01, _tiny=1.2e-38):
    # type: (Tensor, List[Tensor], List[Tensor], float, float) -> Tensor
    """
    update dense preconditioner P = Q^T*Q
    Q: Cholesky factor of preconditioner with positive diagonal entries
    dxs: list of perturbations of parameters
    dgs: list of perturbations of gradients
    step: update step size normalized to range [0, 1]
    _tiny: an offset to avoid division by zero
    """
    dx = torch.cat([torch.reshape(x, [-1, 1]) for x in dxs])
    dg = torch.cat([torch.reshape(g, [-1, 1]) for g in dgs])

    a = Q.mm(dg)
    #b = torch.triangular_solve(dx, Q, upper=True, transpose=True)[0]
    b = torch.linalg.solve_triangular(Q.t(), dx, upper=False)

    grad = torch.triu(a.mm(a.t()) - b.mm(b.t()))
    step0 = step/(grad.abs().max() + _tiny)

    return Q - step0*grad.mm(Q)
//  [python]
def solve_optimal_mean(self, y, L_A, K_mn, K_mn_y):
    n = y.shape[1]
    
    L_inv_Kmn_y = torch.triangular_solve(K_mn_y, L_A, upper=False)[0]
    Kmn_sum = torch.sum(K_mn, dim=(2), keepdim=True)
    L_inv_Kmn_sum = torch.triangular_solve(Kmn_sum, L_A, upper=False)[0]

    A1T_A1 = n
    A1T_b1 = torch.sum(y, dim=(1), keepdim=True)
    tmp = torch.transpose(L_inv_Kmn_sum, dim0=1, dim1=2)
    A2T_A2 = torch.matmul(tmp, L_inv_Kmn_sum)
    A2T_b2 = torch.matmul(tmp, L_inv_Kmn_y)

    mean = (A1T_b1-A2T_b2)/(A1T_A1-A2T_A2)
    return mean
//  [python]
def solve_optimal_mean(self, y, L_A, K_mn, K_mn_y):
    n = y.shape[1]
    
    L_inv_Kmn_y = torch.linalg.solve_triangular(L_A, K_mn_y, upper=False)
    Kmn_sum = torch.sum(K_mn, dim=(2), keepdim=True)
    L_inv_Kmn_sum = torch.linalg.solve_triangular(L_A, Kmn_sum, upper=False)

    A1T_A1 = n
    A1T_b1 = torch.sum(y, dim=(1), keepdim=True)
    tmp = torch.transpose(L_inv_Kmn_sum, dim0=1, dim1=2)
    A2T_A2 = torch.matmul(tmp, L_inv_Kmn_sum)
    A2T_b2 = torch.matmul(tmp, L_inv_Kmn_y)

    mean = (A1T_b1-A2T_b2)/(A1T_A1-A2T_A2)
    return mean
//  [python]
def _test_linalg_solve_triangular(self, A, B, upper, left, uni):
        X = torch.solve(B, A, upper=upper, left=left)
        if left:
            self.assertEqual(A @ X, B)
        else:
            self.assertEqual(X @ A, B)
        out = B
        # B may be expanded
        if not B.is_contiguous() and not B.transpose(-2, -1).is_contiguous():
            out = B.clone()
        torch.solve(B, A, upper=upper, left=left, out=out)
        self.assertEqual(X, out)
//  [python]
def _test_linalg_solve_triangular(self, A, B, upper, left, uni):
        X = torch.linalg.solve_triangular(A, B, upper=upper, left=left, unitriangular=uni)
        if left:
            self.assertEqual(A @ X, B)
        else:
            self.assertEqual(X @ A, B)
        out = B
        # B may be expanded
        if not B.is_contiguous() and not B.transpose(-2, -1).is_contiguous():
            out = B.clone()
        torch.linalg.solve_triangular(A, B, upper=upper, left=left, unitriangular=uni, out=out)
        self.assertEqual(X, out)
//  [python]
def _update_precond_dense_scale(Ql, qr, dX, dG, step=0.01, _tiny=1.2e-38):
    # type: (Tensor, Tensor, Tensor, Tensor, float, float) -> Tuple[Tensor, Tensor]
    """
    update (dense, scaling) preconditioner P = kron_prod(Qr^T*Qr, Ql^T*Ql), where
    dX and dG have shape (M, N)
    Ql has shape (M, M)
    qr has shape (1, N)
    qr is the diagonal part of Qr
    dX is perturbation of (matrix) parameter
    dG is perturbation of (matrix) gradient
    step: update step size
    _tiny: an offset to avoid division by zero
    """
    max_l = torch.max(torch.diag(Ql))
    max_r = torch.max(qr)

    rho = torch.sqrt(max_l/max_r)
    Ql /= rho
    qr *= rho

    A = Ql.mm( dG*qr )
    Bt = torch.triangular_solve(dX/qr, Ql, upper=False, transpose=True)[0]

    grad1 = torch.triu(A.mm(A.t()) - Bt.mm(Bt.t()))
    grad2 = torch.sum(A*A, dim=0, keepdim=True) - torch.sum(Bt*Bt, dim=0, keepdim=True)

    step1 = step/(torch.max(torch.abs(grad1)) + _tiny)
    step2 = step/(torch.max(torch.abs(grad2)) + _tiny)

    return Ql - step1*grad1.mm(Ql), qr - step2*grad2*qr
//  [python]
def _update_precond_dense_scale(Ql, qr, dX, dG, step=0.01, _tiny=1.2e-38):
    # type: (Tensor, Tensor, Tensor, Tensor, float, float) -> Tuple[Tensor, Tensor]
    """
    update (dense, scaling) preconditioner P = kron_prod(Qr^T*Qr, Ql^T*Ql), where
    dX and dG have shape (M, N)
    Ql has shape (M, M)
    qr has shape (1, N)
    qr is the diagonal part of Qr
    dX is perturbation of (matrix) parameter
    dG is perturbation of (matrix) gradient
    step: update step size
    _tiny: an offset to avoid division by zero
    """
    max_l = torch.max(torch.diag(Ql))
    max_r = torch.max(qr)

    rho = torch.sqrt(max_l/max_r)
    Ql /= rho
    qr *= rho

    A = Ql.mm( dG*qr )
    #Bt = torch.triangular_solve(dX/qr, Ql, upper=True, transpose=True)[0]
    Bt = torch.linalg.solve_triangular(Ql.t(), dX/qr, upper=False)

    grad1 = torch.triu(A.mm(A.t()) - Bt.mm(Bt.t()))
    grad2 = torch.sum(A*A, dim=0, keepdim=True) - torch.sum(Bt*Bt, dim=0, keepdim=True)

    step1 = step/(torch.max(torch.abs(grad1)) + _tiny)
    step2 = step/(torch.max(torch.abs(grad2)) + _tiny)

    return Ql - step1*grad1.mm(Ql), qr - step2*grad2*qr
//  [python]
def test_linalg_solve_triangular_broadcasting(self, device, dtype):
        make_arg = partial(make_tensor, dtype=dtype, device=device)

        sizes = (((2, 1, 3, 4, 4), (2, 1, 3, 4, 6)),
                 ((2, 1, 3, 4, 4), (4, 6)),
                 ((4, 4), (2, 1, 3, 4, 2)),
                 ((1, 3, 1, 4, 4), (2, 1, 3, 4, 5)))
        for size_A, size_B in sizes:
            for left, upper, uni in itertools.product([True, False], repeat=3):
                A = make_arg(size_A)
                if upper:
                    A.triu_()
                else:
                    A.tril_()
                diag = A.diagonal(0, -2, -1)
                if uni:
                    diag.fill_(1.)
                else:
                    diag[diag.abs() < 1e-6] = 1.
                B = make_arg(size_B)
                if not left:
                    B.transpose_(-2, -1)

                X = torch.triangular_solve(B, A, upper=upper, left=left, unitriangular=uni)[0]
                if left:
                    B_other = A @ X
                else:
                    B_other = X @ A

                self.assertEqual(*torch.broadcast_tensors(B, B_other))
//  [python]
def test_linalg_solve_triangular_broadcasting(self, device, dtype):
        make_arg = partial(make_tensor, dtype=dtype, device=device)

        sizes = (((2, 1, 3, 4, 4), (2, 1, 3, 4, 6)),
                 ((2, 1, 3, 4, 4), (4, 6)),
                 ((4, 4), (2, 1, 3, 4, 2)),
                 ((1, 3, 1, 4, 4), (2, 1, 3, 4, 5)))
        for size_A, size_B in sizes:
            for left, upper, uni in itertools.product([True, False], repeat=3):
                A = make_arg(size_A)
                if upper:
                    A.triu_()
                else:
                    A.tril_()
                diag = A.diagonal(0, -2, -1)
                if uni:
                    diag.fill_(1.)
                else:
                    diag[diag.abs() < 1e-6] = 1.
                B = make_arg(size_B)
                if not left:
                    B.transpose_(-2, -1)

                X = torch.linalg.solve_triangular(A, B, upper=upper, left=left, unitriangular=uni)
                if left:
                    B_other = A @ X
                else:
                    B_other = X @ A

                self.assertEqual(*torch.broadcast_tensors(B, B_other))
//  [python]
def predict_f(self, X, full=False):
        with torch.inference_mode():
            X = self._check_input(X)  # MxD
            if self.mean is not None:
                y = self.y - self.mean(self.X).reshape(-1,1)  # Nx1
            else:
                y = self.y  # Nx1

            Kff = self.kernel.K(self.X)
            Kff += self._index_channel(self.likelihood.scale().square(), self.X) * self.eye  # NxN
            if self.data_variance is not None:
                Kff += self.data_variance
            Kfs = self.kernel.K(self.X,X)  # NxM

            Lff = self._cholesky(Kff, add_jitter=True)  # NxN
            v = torch.triangular_solve(Kfs,Lff,upper=False).solution  # NxM

            mu = Kfs.T.mm(torch.cholesky_solve(y,Lff))  # Mx1
            if self.mean is not None:
                mu += self.mean(X).reshape(-1,1)  # Mx1

            if full:
                Kss = self.kernel.K(X)  # MxM
                var = Kss - v.T.mm(v)  # MxM
            else:
                Kss_diag = self.kernel.K_diag(X)  # M
                var = Kss_diag - v.T.square().sum(dim=1)  # M
                var = var.reshape(-1,1)
            return mu, var
//  [python]
def predict_f(self, X, full=False):
        with torch.inference_mode():
            X = self._check_input(X)  # MxD
            if self.mean is not None:
                y = self.y - self.mean(self.X).reshape(-1,1)  # Nx1
            else:
                y = self.y  # Nx1

            Kff = self.kernel.K(self.X)
            Kff += self._index_channel(self.likelihood.scale().square(), self.X) * self.eye  # NxN
            if self.data_variance is not None:
                Kff += self.data_variance
            Kfs = self.kernel.K(self.X,X)  # NxM

            Lff = self._cholesky(Kff, add_jitter=True)  # NxN
            v = torch.linalg.solve_triangular(Lff,Kfs,upper=False)  # NxM

            mu = Kfs.T.mm(torch.cholesky_solve(y,Lff))  # Mx1
            if self.mean is not None:
                mu += self.mean(X).reshape(-1,1)  # Mx1

            if full:
                Kss = self.kernel.K(X)  # MxM
                var = Kss - v.T.mm(v)  # MxM
            else:
                Kss_diag = self.kernel.K_diag(X)  # M
                var = Kss_diag - v.T.square().sum(dim=1)  # M
                var = var.reshape(-1,1)
            return mu, var
//  [python]
def _update_precond_norm_dense(ql, Qr, dX, dG, step=0.01, _tiny=1.2e-38):
    # type: (Tensor, Tensor, Tensor, Tensor, float, float) -> Tuple[Tensor, Tensor]
    """
    update (normalization, dense) Kronecker product preconditioner P = kron_prod(Qr^T*Qr, Ql^T*Ql), where
    dX and dG have shape (M, N)
    ql has shape (2, M)
    Qr has shape (N, N)
    ql[0] is the diagonal part of Ql
    ql[1,0:-1] is the last column of Ql, excluding the last entry
    dX is perturbation of (matrix) parameter
    dG is perturbation of (matrix) gradient
    step: update step size normalized to range [0, 1]
    _tiny: an offset to avoid division by zero
    """
    # make sure that Ql and Qr have similar dynamic range
    max_l = torch.max(ql[0])
    max_r = torch.max(torch.diag(Qr))
    rho = torch.sqrt(max_l/max_r)
    ql /= rho
    Qr *= rho

    # refer to https://arxiv.org/abs/1512.04202 for details
    A = ql[0:1].t()*dG + ql[1:].t().mm( dG[-1:] ) # Ql*dG
    A = A.mm(Qr.t())

    Bt = dX/ql[0:1].t()
    Bt[-1:] -= (ql[1:]/(ql[0:1]*ql[0,-1])).mm(dX)
    Bt = torch.triangular_solve(Bt.t(), Qr, upper=True, transpose=True)[0].t()
    #Bt = torch.linalg.solve_triangular(Qr, Bt, upper=True, left=False)

    grad1_diag = torch.sum(A*A, dim=1) - torch.sum(Bt*Bt, dim=1)
    grad1_bias = A[:-1].mm(A[-1:].t()) - Bt[:-1].mm(Bt[-1:].t())
    grad1_bias = torch.cat([torch.squeeze(grad1_bias), grad1_bias.new_zeros(1)])

    step1 = step/(torch.max(torch.max(torch.abs(grad1_diag)),
                            torch.max(torch.abs(grad1_bias))) + _tiny)
    new_ql0 = ql[0] - step1*grad1_diag*ql[0]
    new_ql1 = ql[1] - step1*(grad1_diag*ql[1] + ql[0,-1]*grad1_bias)

    grad2 = torch.triu(A.t().mm(A) - Bt.t().mm(Bt))
    step2 = step/(torch.max(torch.abs(grad2)) + _tiny)

    return torch.stack((new_ql0, new_ql1)), Qr - step2*grad2.mm(Qr)
//  [python]
def _update_precond_norm_dense(ql, Qr, dX, dG, step=0.01, _tiny=1.2e-38):
    # type: (Tensor, Tensor, Tensor, Tensor, float, float) -> Tuple[Tensor, Tensor]
    """
    update (normalization, dense) Kronecker product preconditioner P = kron_prod(Qr^T*Qr, Ql^T*Ql), where
    dX and dG have shape (M, N)
    ql has shape (2, M)
    Qr has shape (N, N)
    ql[0] is the diagonal part of Ql
    ql[1,0:-1] is the last column of Ql, excluding the last entry
    dX is perturbation of (matrix) parameter
    dG is perturbation of (matrix) gradient
    step: update step size normalized to range [0, 1]
    _tiny: an offset to avoid division by zero
    """
    # make sure that Ql and Qr have similar dynamic range
    max_l = torch.max(ql[0])
    max_r = torch.max(torch.diag(Qr))
    rho = torch.sqrt(max_l/max_r)
    ql /= rho
    Qr *= rho

    # refer to https://arxiv.org/abs/1512.04202 for details
    A = ql[0:1].t()*dG + ql[1:].t().mm( dG[-1:] ) # Ql*dG
    A = A.mm(Qr.t())

    Bt = dX/ql[0:1].t()
    Bt[-1:] -= (ql[1:]/(ql[0:1]*ql[0,-1])).mm(dX)
    #Bt = torch.triangular_solve(Bt.t(), Qr, upper=True, transpose=True)[0].t()
    Bt = torch.linalg.solve_triangular(Qr, Bt, upper=True, left=False)

    grad1_diag = torch.sum(A*A, dim=1) - torch.sum(Bt*Bt, dim=1)
    grad1_bias = A[:-1].mm(A[-1:].t()) - Bt[:-1].mm(Bt[-1:].t())
    grad1_bias = torch.cat([torch.squeeze(grad1_bias), grad1_bias.new_zeros(1)])

    step1 = step/(torch.max(torch.max(torch.abs(grad1_diag)),
                            torch.max(torch.abs(grad1_bias))) + _tiny)
    new_ql0 = ql[0] - step1*grad1_diag*ql[0]
    new_ql1 = ql[1] - step1*(grad1_diag*ql[1] + ql[0,-1]*grad1_bias)

    grad2 = torch.triu(A.t().mm(A) - Bt.t().mm(Bt))
    step2 = step/(torch.max(torch.abs(grad2)) + _tiny)

    return torch.stack((new_ql0, new_ql1)), Qr - step2*grad2.mm(Qr)
//  [python]
def log_marginal_likelihood(self):
        if self.mean is not None:
            y = self.y - self.mean(self.X).reshape(-1,1)  # Nx1
        else:
            y = self.y  # Nx1

        Kff_diag = self.kernel.K_diag(self.X)  # N
        Kuf = self.kernel.K(self.Z(),self.X)  # MxN
        Kuu = self.kernel.K(self.Z())  # MxM

        Luu = self._cholesky(Kuu, add_jitter=True)  # MxM;  Luu = Kuu^(1/2)
        v = torch.triangular_solve(Kuf,Luu,upper=False)[0]  # MxN;  Kuu^(-1/2).Kuf
        g = Kff_diag - v.T.square().sum(dim=1) + self._index_channel(self.likelihood.scale().square(), self.X)  # N;  diag(Kff-Qff) + sigma^2.I
        G = torch.diagflat(1.0/g)  # N
        L = self._cholesky(v.mm(G).mm(v.T) + self.eye)  # MxM;  (Kuu^(-1/2).Kuf.G.Kfu.Kuu^(-1/2) + I)^(1/2)

        c = torch.triangular_solve(v.mm(G).mm(y),L,upper=False)[0]  # Mx1;  L^(-1).Kuu^(-1/2).Kuf.G.y

        p = -self.log_marginal_likelihood_constant
        p -= L.diagonal().log().sum() # 0.5 is taken as the square root of L
        p -= 0.5*g.log().sum()
        p -= 0.5*y.T.mm(G).mm(y).squeeze()
        p += 0.5*c.T.mm(c).squeeze()
        return p
//  [python]
def log_marginal_likelihood(self):
        if self.mean is not None:
            y = self.y - self.mean(self.X).reshape(-1,1)  # Nx1
        else:
            y = self.y  # Nx1

        Kff_diag = self.kernel.K_diag(self.X)  # N
        Kuf = self.kernel.K(self.Z(),self.X)  # MxN
        Kuu = self.kernel.K(self.Z())  # MxM

        Luu = self._cholesky(Kuu, add_jitter=True)  # MxM;  Luu = Kuu^(1/2)
        v = torch.linalg.solve_triangular(Luu,Kuf,upper=False)  # MxN;  Kuu^(-1/2).Kuf
        g = Kff_diag - v.T.square().sum(dim=1) + self._index_channel(self.likelihood.scale().square(), self.X)  # N;  diag(Kff-Qff) + sigma^2.I
        G = torch.diagflat(1.0/g)  # N
        L = self._cholesky(v.mm(G).mm(v.T) + self.eye)  # MxM;  (Kuu^(-1/2).Kuf.G.Kfu.Kuu^(-1/2) + I)^(1/2)

        c = torch.linalg.solve_triangular(L,v.mm(G).mm(y),upper=False)  # Mx1;  L^(-1).Kuu^(-1/2).Kuf.G.y

        p = -self.log_marginal_likelihood_constant
        p -= L.diagonal().log().sum() # 0.5 is taken as the square root of L
        p -= 0.5*g.log().sum()
        p -= 0.5*y.T.mm(G).mm(y).squeeze()
        p += 0.5*c.T.mm(c).squeeze()
        return p
//  [python]
def predict_f(self, X, full=False):
        with torch.inference_mode():
            X = self._check_input(X)  # MxD

            Kff = self.kernel(self.X)
            Kfs = self.kernel(self.X,X)  # NxS

            L = self._cholesky(Kff + (1.0/self.q_lambda().square()).diagflat())  # NxN
            a = torch.triangular_solve(Kfs,L,upper=False).solution  # NxS;  Kuu^(-1/2).Kus

            mu = Kfs.T.mm(self.q_nu())  # Sx1
            if self.mean is not None:
                mu += self.mean(X).reshape(-1,1)  # Sx1

            if full:
                Kss = self.kernel(X)  # SxS
                var = Kss - a.T.mm(a)  # SxS
            else:
                Kss_diag = self.kernel.K_diag(X)  # M
                var = Kss_diag - a.T.square().sum(dim=1)  # M
                var = var.reshape(-1,1)
            return mu, var
//  [python]
def predict_f(self, X, full=False):
        with torch.inference_mode():
            X = self._check_input(X)  # MxD

            Kff = self.kernel(self.X)
            Kfs = self.kernel(self.X,X)  # NxS

            L = self._cholesky(Kff + (1.0/self.q_lambda().square()).diagflat())  # NxN
            a = torch.linalg.solve_triangular(L,Kfs,upper=False)  # NxS;  Kuu^(-1/2).Kus

            mu = Kfs.T.mm(self.q_nu())  # Sx1
            if self.mean is not None:
                mu += self.mean(X).reshape(-1,1)  # Sx1

            if full:
                Kss = self.kernel(X)  # SxS
                var = Kss - a.T.mm(a)  # SxS
            else:
                Kss_diag = self.kernel.K_diag(X)  # M
                var = Kss_diag - a.T.square().sum(dim=1)  # M
                var = var.reshape(-1,1)
            return mu, var
//  [python]
def predict_f(self, X, full=False):
        with torch.inference_mode():
            X = self._check_input(X)  # MxD
            if self.mean is not None:
                y = self.y - self.mean(self.X).reshape(-1,1)  # Nx1
            else:
                y = self.y  # Nx1

            Kff_diag = self.kernel.K_diag(self.X)  # N
            Kuf = self.kernel.K(self.Z(),self.X)  # MxN
            Kuu = self.kernel.K(self.Z())  # MxM
            Kus = self.kernel.K(self.Z(),X)  # MxS

            Luu = self._cholesky(Kuu, add_jitter=True)  # MxM;  Kuu^(1/2)
            v = torch.triangular_solve(Kuf,Luu,upper=False)[0]  # MxN;  Kuu^(-1/2).Kuf
            g = Kff_diag - v.T.square().sum(dim=1) + self._index_channel(self.likelihood.scale().square(), self.X)
            G = torch.diagflat(1.0/g)  # N
            L = self._cholesky(v.mm(G).mm(v.T) + self.eye)  # MxM;  (Kuu^(-1/2).Kuf.G.Kfu.Kuu^(-1/2) + I)^(1/2)

            a = torch.triangular_solve(Kus,Luu,upper=False)[0]  # NxM
            b = torch.triangular_solve(a,L,upper=False)[0]
            c = torch.triangular_solve(v.mm(G).mm(y),L,upper=False)[0]  # Mx1;  L^(-1).Kuu^(-1/2).Kuf.G.y

            mu = b.T.mm(c)  # Mx1
            if self.mean is not None:
                mu += self.mean(X).reshape(-1,1)  # Mx1

            if full:
                Kss = self.kernel(X)  # MxM
                var = Kss - a.T.mm(w) + b.T.mm(u)  # MxM
            else:
                Kss_diag = self.kernel.K_diag(X)  # M
                var = Kss_diag - a.T.square().sum(dim=1) + b.T.square().sum(dim=1)  # M
                var = var.reshape(-1,1)
            return mu, var
//  [python]
def predict_f(self, X, full=False):
        with torch.inference_mode():
            X = self._check_input(X)  # MxD
            if self.mean is not None:
                y = self.y - self.mean(self.X).reshape(-1,1)  # Nx1
            else:
                y = self.y  # Nx1

            Kff_diag = self.kernel.K_diag(self.X)  # N
            Kuf = self.kernel.K(self.Z(),self.X)  # MxN
            Kuu = self.kernel.K(self.Z())  # MxM
            Kus = self.kernel.K(self.Z(),X)  # MxS

            Luu = self._cholesky(Kuu, add_jitter=True)  # MxM;  Kuu^(1/2)
            v = torch.linalg.solve_triangular(Luu,Kuf,upper=False)  # MxN;  Kuu^(-1/2).Kuf
            g = Kff_diag - v.T.square().sum(dim=1) + self._index_channel(self.likelihood.scale().square(), self.X)
            G = torch.diagflat(1.0/g)  # N
            L = self._cholesky(v.mm(G).mm(v.T) + self.eye)  # MxM;  (Kuu^(-1/2).Kuf.G.Kfu.Kuu^(-1/2) + I)^(1/2)

            a = torch.linalg.solve_triangular(Luu,Kus,upper=False)  # NxM
            b = torch.linalg.solve_triangular(L,a,upper=False)
            c = torch.linalg.solve_triangular(L,v.mm(G).mm(y),upper=False)  # Mx1;  L^(-1).Kuu^(-1/2).Kuf.G.y

            mu = b.T.mm(c)  # Mx1
            if self.mean is not None:
                mu += self.mean(X).reshape(-1,1)  # Mx1

            if full:
                Kss = self.kernel(X)  # MxM
                var = Kss - a.T.mm(w) + b.T.mm(u)  # MxM
            else:
                Kss_diag = self.kernel.K_diag(X)  # M
                var = Kss_diag - a.T.square().sum(dim=1) + b.T.square().sum(dim=1)  # M
                var = var.reshape(-1,1)
            return mu, var
//  [python]
def solve_triangular(A, b, **kwargs):
    return torch.solve(b.unsqueeze(1), A, **kwargs)[0].squeeze(1)
//  [python]
def solve_triangular(A, b, **kwargs):
    return torch.linalg.solve_triangular(A, b.unsqueeze(1), **kwargs)[0].squeeze(1)
//  [python]
def inverse_no_cache(self, inputs):
        """Cost:
            output = O(D^2N + KDN)
            logabsdet = O(D)
        where:
            K = num of householder transforms
            D = num of features
            N = num of inputs
        """
        upper = self._create_upper()
        outputs = inputs - self.bias
        outputs, _ = self.orthogonal.inverse(
            outputs
        )  # Ignore logabsdet since we know it's zero.
        outputs = torch.triangular_solve(outputs.t(), upper, upper=True)[0]
        outputs = outputs.t()
        logabsdet = -self.logabsdet()
        logabsdet = logabsdet * outputs.new_ones(outputs.shape[0])
        return outputs, logabsdet
//  [python]
def inverse_no_cache(self, inputs):
        """Cost:
            output = O(D^2N + KDN)
            logabsdet = O(D)
        where:
            K = num of householder transforms
            D = num of features
            N = num of inputs
        """
        upper = self._create_upper()
        outputs = inputs - self.bias
        outputs, _ = self.orthogonal.inverse(
            outputs
        )  # Ignore logabsdet since we know it's zero.
        outputs = torch.linalg.solve_triangular(upper, outputs.t(), upper=True)
        outputs = outputs.t()
        logabsdet = -self.logabsdet()
        logabsdet = logabsdet * outputs.new_ones(outputs.shape[0])
        return outputs, logabsdet
//  [python]
def weight_inverse(self):
        """Cost:
            inverse = O(D^3 + KD^2)
        where:
            K = num of householder transforms
            D = num of features
        """
        upper = self._create_upper()
        identity = torch.eye(self.features, self.features)
        upper_inv = torch.triangular_solve(identity, upper, upper=True).solution
        weight_inv, _ = self.orthogonal(upper_inv)
        return weight_inv
//  [python]
def weight_inverse(self):
        """Cost:
            inverse = O(D^3 + KD^2)
        where:
            K = num of householder transforms
            D = num of features
        """
        upper = self._create_upper()
        identity = torch.eye(self.features, self.features)
        upper_inv = torch.linalg.solve_triangular(upper, identity, upper=True)
        weight_inv, _ = self.orthogonal(upper_inv)
        return weight_inv
//  [python]
def elbo(self):
        if self.mean is not None:
            y = self.y - self.mean(self.X).reshape(-1,1)  # Nx1
        else:
            y = self.y  # Nx1

        Kff_diag = self.kernel.K_diag(self.X)  # N
        Kuf = self.kernel(self.Z(),self.X)  # MxN
        Kuu = self.kernel(self.Z())  # MxM

        Luu = self._cholesky(Kuu, add_jitter=True)  # MxM;  Kuu^(1/2)
        v = torch.triangular_solve(Kuf,Luu,upper=False)[0]  # MxN;  Kuu^(-1/2).Kuf
        Q = v.mm(v.T)  # MxM;  Kuu^(-1/2).Kuf.Kfu.Kuu^(-1/2)
        L = self._cholesky(Q/self.likelihood.scale().square() + self.eye)  # MxM;  (Q/sigma^2 + I)^(1/2)

        c = torch.triangular_solve(v.mm(y),L,upper=False)[0]/self.likelihood.scale().square()  # Mx1;  L^(-1).Kuu^(-1/2).Kuf.y

        # p = log N(0, Kfu.Kuu^(-1).Kuf + I/sigma^2) - 1/(2.sigma^2).Trace(Kff - Kfu.Kuu^(-1).Kuf)
        p = -self.log_marginal_likelihood_constant
        p -= L.diagonal().log().sum() # 0.5 is taken as the square root of L
        p -= self.X.shape[0]*self.likelihood.scale().log()
        p -= 0.5*y.T.mm(y).squeeze()/self.likelihood.scale().square()
        p += 0.5*c.T.mm(c).squeeze()
        p -= 0.5*(Kff_diag.sum() - Q.trace())/self.likelihood.scale().square() # trace
        return p
//  [python]
def elbo(self):
        if self.mean is not None:
            y = self.y - self.mean(self.X).reshape(-1,1)  # Nx1
        else:
            y = self.y  # Nx1

        Kff_diag = self.kernel.K_diag(self.X)  # N
        Kuf = self.kernel(self.Z(),self.X)  # MxN
        Kuu = self.kernel(self.Z())  # MxM

        Luu = self._cholesky(Kuu, add_jitter=True)  # MxM;  Kuu^(1/2)
        v = torch.linalg.solve_triangular(Luu,Kuf,upper=False)  # MxN;  Kuu^(-1/2).Kuf
        Q = v.mm(v.T)  # MxM;  Kuu^(-1/2).Kuf.Kfu.Kuu^(-1/2)
        L = self._cholesky(Q/self.likelihood.scale().square() + self.eye)  # MxM;  (Q/sigma^2 + I)^(1/2)

        c = torch.linalg.solve_triangular(L,v.mm(y),upper=False)/self.likelihood.scale().square()  # Mx1;  L^(-1).Kuu^(-1/2).Kuf.y

        # p = log N(0, Kfu.Kuu^(-1).Kuf + I/sigma^2) - 1/(2.sigma^2).Trace(Kff - Kfu.Kuu^(-1).Kuf)
        p = -self.log_marginal_likelihood_constant
        p -= L.diagonal().log().sum() # 0.5 is taken as the square root of L
        p -= self.X.shape[0]*self.likelihood.scale().log()
        p -= 0.5*y.T.mm(y).squeeze()/self.likelihood.scale().square()
        p += 0.5*c.T.mm(c).squeeze()
        p -= 0.5*(Kff_diag.sum() - Q.trace())/self.likelihood.scale().square() # trace
        return p
//  [python]
def elbo(self):
        if self.mean is not None:
            y = self.y - self.mean(self.X).reshape(-1,1)  # Nx1
        else:
            y = self.y  # Nx1

        q_nu = self.q_nu()
        q_lambda = self.q_lambda()

        Kff = self.kernel(self.X)  # NxN
        L = self._cholesky(q_lambda*q_lambda.T*Kff + self.eye)
        invL = torch.triangular_solve(self.eye,L,upper=False)[0]  # NxN

        qf_mu = Kff.mm(q_nu)
        qf_var_diag = 1.0/q_lambda.square() - (invL.T.mm(invL)/q_lambda/q_lambda.T).diagonal().reshape(-1,1)

        kl = q_nu.T.mm(qf_mu).squeeze()  # Mahalanobis
        kl += L.diagonal().square().log().sum()  # determinant TODO: is this correct?
        #kl += invL.diagonal().square().sum()  # trace
        kl += invL.square().sum()  # trace
        kl -= q_nu.shape[0]

        if self.mean is not None:
            qf_mu = qf_mu - self.mean(self.X).reshape(-1,1)  # Sx1
        var_exp = self.likelihood.variational_expectation(self.X, y, qf_mu, qf_var_diag)

        #eye = torch.eye(q_lambda.shape[0], device=config.device, dtype=config.dtype)
        #qf_var = (1.0/q_lambda.square())*eye - invL.T.mm(invL)/q_lambda/q_lambda.T
        #kl = -q_nu.shape[0]
        #kl += q_nu.T.mm(qf_mu).squeeze()  # Mahalanobis
        #kl -= qf_var.det().log()  # determinant
        #kl += invL.diagonal().square().sum()  # trace

        #kl = -q_nu.shape[0]
        #kl += q_nu.T.mm(q_nu).squeeze()  # Mahalanobis
        #kl -= qf_var.det().log()  # determinant
        #kl += qf_var_diag.sum()  # trace
        return var_exp - 0.5*kl
//  [python]
def elbo(self):
        if self.mean is not None:
            y = self.y - self.mean(self.X).reshape(-1,1)  # Nx1
        else:
            y = self.y  # Nx1

        q_nu = self.q_nu()
        q_lambda = self.q_lambda()

        Kff = self.kernel(self.X)  # NxN
        L = self._cholesky(q_lambda*q_lambda.T*Kff + self.eye)
        invL = torch.linalg.solve_triangular(L,self.eye,upper=False)  # NxN

        qf_mu = Kff.mm(q_nu)
        qf_var_diag = 1.0/q_lambda.square() - (invL.T.mm(invL)/q_lambda/q_lambda.T).diagonal().reshape(-1,1)

        kl = q_nu.T.mm(qf_mu).squeeze()  # Mahalanobis
        kl += L.diagonal().square().log().sum()  # determinant TODO: is this correct?
        #kl += invL.diagonal().square().sum()  # trace
        kl += invL.square().sum()  # trace
        kl -= q_nu.shape[0]

        if self.mean is not None:
            qf_mu = qf_mu - self.mean(self.X).reshape(-1,1)  # Sx1
        var_exp = self.likelihood.variational_expectation(self.X, y, qf_mu, qf_var_diag)

        #eye = torch.eye(q_lambda.shape[0], device=config.device, dtype=config.dtype)
        #qf_var = (1.0/q_lambda.square())*eye - invL.T.mm(invL)/q_lambda/q_lambda.T
        #kl = -q_nu.shape[0]
        #kl += q_nu.T.mm(qf_mu).squeeze()  # Mahalanobis
        #kl -= qf_var.det().log()  # determinant
        #kl += invL.diagonal().square().sum()  # trace

        #kl = -q_nu.shape[0]
        #kl += q_nu.T.mm(q_nu).squeeze()  # Mahalanobis
        #kl -= qf_var.det().log()  # determinant
        #kl += qf_var_diag.sum()  # trace
        return var_exp - 0.5*kl
//  [python]
def _predict_f(self, X, full=False):
        Kuu = self.kernel(self.Z())
        Kus = self.kernel(self.Z(),X)  # NxS

        Luu = self._cholesky(Kuu, add_jitter=True)  # NxN
        a = torch.triangular_solve(Kus,Luu,upper=False)[0]  # NxS;  Kuu^(-1/2).Kus
        b = self.q_sqrt().tril().T.mm(torch.triangular_solve(Kus,Luu,upper=False)[0])

        mu = Kus.T.mm(torch.triangular_solve(self.q_mu(),Luu.T,upper=True)[0])  # Sx1
        if full:
            Kss = self.kernel(X)  # SxS
            var = Kss - a.T.mm(a) + b.T.mm(b)  # SxS
        else:
            Kss_diag = self.kernel.K_diag(X)  # M
            var = Kss_diag - a.T.square().sum(dim=1) + b.T.square().sum(dim=1)  # M
            var = var.reshape(-1,1)
        return mu, var
//  [python]
def _predict_f(self, X, full=False):
        Kuu = self.kernel(self.Z())
        Kus = self.kernel(self.Z(),X)  # NxS

        Luu = self._cholesky(Kuu, add_jitter=True)  # NxN
        a = torch.linalg.solve_triangular(Luu,Kus,upper=False)  # NxS;  Kuu^(-1/2).Kus
        b = self.q_sqrt().tril().T.mm(torch.linalg.solve_triangular(Luu,Kus,upper=False))

        mu = Kus.T.mm(torch.linalg.solve_triangular(Luu.T,self.q_mu(),upper=True))  # Sx1
        if full:
            Kss = self.kernel(X)  # SxS
            var = Kss - a.T.mm(a) + b.T.mm(b)  # SxS
        else:
            Kss_diag = self.kernel.K_diag(X)  # M
            var = Kss_diag - a.T.square().sum(dim=1) + b.T.square().sum(dim=1)  # M
            var = var.reshape(-1,1)
        return mu, var
//  [python]
def test_namedtuple():
    a = ProvenanceTensor(torch.randn(3, 3), frozenset({"a"}))
    b = ProvenanceTensor(torch.randn(3, 3), frozenset({"b"}))
    result = torch.solve(a, b, upper=True)[0]
    actual = getattr(result, "_provenance", frozenset())
    assert actual == frozenset({"a", "b"})
//  [python]
def test_namedtuple():
    a = ProvenanceTensor(torch.randn(3, 3), frozenset({"a"}))
    b = ProvenanceTensor(torch.randn(3, 3), frozenset({"b"}))
    result = torch.linalg.solve_triangular(b, a, upper=True)
    actual = getattr(result, "_provenance", frozenset())
    assert actual == frozenset({"a", "b"})
//  [python]
def _inverse(self, y):
        """
        :param y: the output of the bijection
        :type y: torch.Tensor

        Inverts y => x.
        """
        return torch.triangular_solve(
            (y - self.loc).unsqueeze(-1), self.scale_tril, upper=False
        ).solution.squeeze(-1)
//  [python]
def _inverse(self, y):
        """
        :param y: the output of the bijection
        :type y: torch.Tensor

        Inverts y => x.
        """
        return torch.linalg.solve_triangular(
            self.scale_tril, (y - self.loc).unsqueeze(-1), upper=False
        ).squeeze(-1)
//  [python]
def _inducing_inv_root(self):
        if not self.training and hasattr(self, "_cached_kernel_inv_root"):
            return self._cached_kernel_inv_root
        else:
            chol = psd_safe_cholesky(self._inducing_mat, upper=True)
            eye = torch.eye(chol.size(-1), device=chol.device, dtype=chol.dtype)
            inv_root = torch.triangular_solve(eye, chol, upper=True).solution

            res = inv_root
            if not self.training:
                self._cached_kernel_inv_root = res
            return res
//  [python]
def _inducing_inv_root(self):
        if not self.training and hasattr(self, "_cached_kernel_inv_root"):
            return self._cached_kernel_inv_root
        else:
            chol = psd_safe_cholesky(self._inducing_mat, upper=True)
            eye = torch.eye(chol.size(-1), device=chol.device, dtype=chol.dtype)
            inv_root = torch.linalg.solve_triangular(chol, eye, upper=True)

            res = inv_root
            if not self.training:
                self._cached_kernel_inv_root = res
            return res
//  [python]
def _inverse(self, y):
        """
        :param y: the output of the bijection
        :type y: torch.Tensor

        Inverts y => x.
        """

        """
        NOTE: This method is equivalent to the following two lines. Using
        Tensor.inverse() would be numerically unstable, however.

        filters = (self.permutation @ self.L @ self.U).inverse()[..., None, None]
        x = F.conv2d(y.view(-1, *y.shape[-3:]), filters)
        return x.view_as(y)

        """

        # Do a matrix vector product over the channel dimension
        # in order to apply inverse permutation matrix
        y_flat = y.flatten(start_dim=-2)
        LUx = (y_flat.unsqueeze(-3) * self.permutation.T.unsqueeze(-1)).sum(-2)

        # Solve L(Ux) = P^1y
        Ux = torch.triangular_solve(LUx, self.L, upper=False)[0]

        # Solve Ux = (PL)^-1y
        x = torch.triangular_solve(Ux, self.U, upper=True)[0]

        # Unflatten x (works when context variable has batch dim)
        return x.reshape(x.shape[:-1] + y.shape[-2:])
//  [python]
def _inverse(self, y):
        """
        :param y: the output of the bijection
        :type y: torch.Tensor

        Inverts y => x.
        """

        """
        NOTE: This method is equivalent to the following two lines. Using
        Tensor.inverse() would be numerically unstable, however.

        filters = (self.permutation @ self.L @ self.U).inverse()[..., None, None]
        x = F.conv2d(y.view(-1, *y.shape[-3:]), filters)
        return x.view_as(y)

        """

        # Do a matrix vector product over the channel dimension
        # in order to apply inverse permutation matrix
        y_flat = y.flatten(start_dim=-2)
        LUx = (y_flat.unsqueeze(-3) * self.permutation.T.unsqueeze(-1)).sum(-2)

        # Solve L(Ux) = P^1y
        Ux = torch.linalg.solve_triangular(self.L, LUx, upper=False)

        # Solve Ux = (PL)^-1y
        x = torch.linalg.solve_triangular(self.U, Ux, upper=True)

        # Unflatten x (works when context variable has batch dim)
        return x.reshape(x.shape[:-1] + y.shape[-2:])
//  [python]
def forward(
        self, x: torch.Tensor, embedded: bool = False, return_std: bool = False
    ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
        """
        Apply inference on the given input sequences.

        :param x: Input sequences.
        :param embedded: If True, x is expected to contain the embeddings of shape
            batch x embedding_dim. Else, x should contain the tokens of the input sequence
            of shape batch x sequence.
        :param return_std: Whether to compute an estimate of the standard deviation.
        :return: The predicted mean and optionally the estimated standard deviation.
        """
        # Compute embeddings if necessary
        if embedded:
            embeddings = x
        else:
            embeddings = self._embed(x)

        # Compute kernel matrix of query and training samples
        kernel = _compute_kernel(
            embeddings, self.train_embeddings.data, self.gamma.data, device=embeddings.device
        )

        # Compute mean (cf. to SI)
        mean = torch.matmul(kernel, self.dual_coefficients) * self.scale + self.intercept

        # If requested, compute standard deviation estimate (cf. to SI)
        if return_std:
            tmp = torch.triangular_solve(kernel.T, self.train_kernel_cholesky, upper=False)[0]
            diagonal = torch.sum(tmp**2, dim=0)
            std = torch.sqrt(1 - diagonal) * self.scale
            return mean, std
        else:
            return mean
//  [python]
def forward(
        self, x: torch.Tensor, embedded: bool = False, return_std: bool = False
    ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
        """
        Apply inference on the given input sequences.

        :param x: Input sequences.
        :param embedded: If True, x is expected to contain the embeddings of shape
            batch x embedding_dim. Else, x should contain the tokens of the input sequence
            of shape batch x sequence.
        :param return_std: Whether to compute an estimate of the standard deviation.
        :return: The predicted mean and optionally the estimated standard deviation.
        """
        # Compute embeddings if necessary
        if embedded:
            embeddings = x
        else:
            embeddings = self._embed(x)

        # Compute kernel matrix of query and training samples
        kernel = _compute_kernel(
            embeddings, self.train_embeddings.data, self.gamma.data, device=embeddings.device
        )

        # Compute mean (cf. to SI)
        mean = torch.matmul(kernel, self.dual_coefficients) * self.scale + self.intercept

        # If requested, compute standard deviation estimate (cf. to SI)
        if return_std:
            tmp = torch.linalg.solve_triangular(self.train_kernel_cholesky, kernel.T, upper=False)
            diagonal = torch.sum(tmp**2, dim=0)
            std = torch.sqrt(1 - diagonal) * self.scale
            return mean, std
        else:
            return mean
//  [python]
def predict_f(self, X, full=False):
        with torch.inference_mode():
            X = self._check_input(X)  # MxD
            if self.mean is not None:
                y = self.y - self.mean(self.X).reshape(-1,1)  # Nx1
            else:
                y = self.y  # Nx1

            Kus = self.kernel(self.Z(),X)  # MxS
            Kuf = self.kernel(self.Z(),self.X)  # MxN
            Kuu = self.kernel(self.Z())  # MxM

            Luu = self._cholesky(Kuu, add_jitter=True)  # MxM;  Kuu^(1/2)
            v = torch.triangular_solve(Kuf,Luu,upper=False).solution  # MxN;  Kuu^(-1/2).Kuf
            L = self._cholesky(v.mm(v.T)/self.likelihood.scale().square() + self.eye)  # MxM;  (Kuu^(-1/2).Kuf.Kfu.Kuu^(-1/2)/sigma^2 + I)^(1/2)

            a = torch.triangular_solve(Kus,Luu,upper=False).solution  # MxS;  Kuu^(-1/2).Kus
            b = torch.triangular_solve(a,L,upper=False).solution  # MxS;  L^(-1).Kuu^(-1/2).Kus
            c = torch.triangular_solve(v.mm(y),L,upper=False).solution/self.likelihood.scale().square()  # Mx1;  L^(-1).Kuu^(-1/2).Kuf.y

            # mu = sigma^(-2).Ksu.Kuu^(-1/2).(sigma^(-2).Kuu^(-1/2).Kuf.Kfu.Kuu^(-1/2) + I)^(-1).Kuu^(-1/2).Kuf.y
            mu = b.T.mm(c)  # Mx1
            if self.mean is not None:
                mu += self.mean(X).reshape(-1,1)  # Mx1

            # var = Kss - Qsf.(Qff + sigma^2 I)^(-1).Qfs
            # below is the equivalent but more stable version by using the matrix inversion lemma
            # var = Kss - Ksu.Kuu^(-1).Kus + Ksu.Kuu^(-1/2).(sigma^(-2).Kuu^(-1/2).Kuf.Kfu.Kuu^(-1/2) + I)^(-1).Kuu^(-1/2).Kus
            if full:
                Kss = self.kernel(X)  # MxM
                var = Kss - a.T.mm(a) + b.T.mm(b)  # MxM
            else:
                Kss_diag = self.kernel.K_diag(X)  # M
                var = Kss_diag - a.T.square().sum(dim=1) + b.T.square().sum(dim=1)  # M
                var = var.reshape(-1,1)
            return mu, var
//  [python]
def predict_f(self, X, full=False):
        with torch.inference_mode():
            X = self._check_input(X)  # MxD
            if self.mean is not None:
                y = self.y - self.mean(self.X).reshape(-1,1)  # Nx1
            else:
                y = self.y  # Nx1

            Kus = self.kernel(self.Z(),X)  # MxS
            Kuf = self.kernel(self.Z(),self.X)  # MxN
            Kuu = self.kernel(self.Z())  # MxM

            Luu = self._cholesky(Kuu, add_jitter=True)  # MxM;  Kuu^(1/2)
            v = torch.linalg.solve_triangular(Luu,Kuf,upper=False)  # MxN;  Kuu^(-1/2).Kuf
            L = self._cholesky(v.mm(v.T)/self.likelihood.scale().square() + self.eye)  # MxM;  (Kuu^(-1/2).Kuf.Kfu.Kuu^(-1/2)/sigma^2 + I)^(1/2)

            a = torch.linalg.solve_triangular(Luu,Kus,upper=False)  # MxS;  Kuu^(-1/2).Kus
            b = torch.linalg.solve_triangular(L,a,upper=False)  # MxS;  L^(-1).Kuu^(-1/2).Kus
            c = torch.linalg.solve_triangular(L,v.mm(y),upper=False)/self.likelihood.scale().square()  # Mx1;  L^(-1).Kuu^(-1/2).Kuf.y

            # mu = sigma^(-2).Ksu.Kuu^(-1/2).(sigma^(-2).Kuu^(-1/2).Kuf.Kfu.Kuu^(-1/2) + I)^(-1).Kuu^(-1/2).Kuf.y
            mu = b.T.mm(c)  # Mx1
            if self.mean is not None:
                mu += self.mean(X).reshape(-1,1)  # Mx1

            # var = Kss - Qsf.(Qff + sigma^2 I)^(-1).Qfs
            # below is the equivalent but more stable version by using the matrix inversion lemma
            # var = Kss - Ksu.Kuu^(-1).Kus + Ksu.Kuu^(-1/2).(sigma^(-2).Kuu^(-1/2).Kuf.Kfu.Kuu^(-1/2) + I)^(-1).Kuu^(-1/2).Kus
            if full:
                Kss = self.kernel(X)  # MxM
                var = Kss - a.T.mm(a) + b.T.mm(b)  # MxM
            else:
                Kss_diag = self.kernel.K_diag(X)  # M
                var = Kss_diag - a.T.square().sum(dim=1) + b.T.square().sum(dim=1)  # M
                var = var.reshape(-1,1)
            return mu, var
//  [python]
def test_solve_triangular(self):
        linear_op = self.create_linear_op()
        rhs = torch.randn(linear_op.size(-1))
        res = torch.triangular_solve(rhs, linear_op, upper=False).solution
        res_actual = rhs / linear_op.diagonal()
        self.assertAllClose(res, res_actual)
        res = torch.triangular_solve(rhs, linear_op, upper=True).solution
        res_actual = rhs / linear_op.diagonal()
        self.assertAllClose(res, res_actual)
        # unittriangular case
        with self.assertRaisesRegex(RuntimeError, "Received `unitriangular=True`"):
            torch.triangular_solve(rhs, linear_op, upper=False, unitriangular=True).solution
        linear_op = DiagLinearOperator(torch.ones(4))  # TODO: Test gradients
        res = torch.triangular_solve(rhs, linear_op, upper=False, unitriangular=True).solution
        self.assertAllClose(res, rhs)
//  [python]
def test_solve_triangular(self):
        linear_op = self.create_linear_op()
        rhs = torch.randn(linear_op.size(-1))
        res = torch.linalg.solve_triangular(linear_op, rhs, upper=False)
        res_actual = rhs / linear_op.diagonal()
        self.assertAllClose(res, res_actual)
        res = torch.linalg.solve_triangular(linear_op, rhs, upper=True)
        res_actual = rhs / linear_op.diagonal()
        self.assertAllClose(res, res_actual)
        # unittriangular case
        with self.assertRaisesRegex(RuntimeError, "Received `unitriangular=True`"):
            torch.linalg.solve_triangular(linear_op, rhs, upper=False, unitriangular=True)
        linear_op = DiagLinearOperator(torch.ones(4))  # TODO: Test gradients
        res = torch.linalg.solve_triangular(linear_op, rhs, upper=False, unitriangular=True)
        self.assertAllClose(res, rhs)
//  [python]
def test_solve_triangular(self):
        linear_op = self.create_linear_op()
        rhs = torch.randn(linear_op.size(-1))
        with self.assertRaisesRegex(NotImplementedError, "torch.linalg.solve_triangular"):
            torch.triangular_solve(rhs, linear_op, upper=True)[0]
//  [python]
def test_solve_triangular(self):
        linear_op = self.create_linear_op()
        rhs = torch.randn(linear_op.size(-1))
        with self.assertRaisesRegex(NotImplementedError, "torch.linalg.solve_triangular"):
            torch.linalg.solve_triangular(linear_op, rhs, upper=True)
//  [python]
def forward(self, x):
        x = x.unsqueeze(2).unsqueeze(3)
        m = x.mean(0).view(self.num_features, -1).mean(-1).view(1, -1, 1, 1)
        if not self.training and self.track_running_stats:  # for inference
            m = self.running_mean
        xn = x - m

        T = xn.permute(1, 0, 2, 3).contiguous().view(self.num_features, -1)
        f_cov = torch.mm(T, T.permute(1, 0)) / (T.shape[-1] - 1)

        eye = torch.eye(self.num_features).type(f_cov.type())

        if not self.training and self.track_running_stats:  # for inference
            f_cov = self.running_variance

        f_cov_shrinked = (1 - self.eps) * f_cov + self.eps * eye

        inv_sqrt = torch.solve(eye, torch.cholesky(f_cov_shrinked))[0]
        
        inv_sqrt = inv_sqrt.contiguous().view(
            self.num_features, self.num_features, 1, 1
        )

        decorrelated = conv2d(xn, inv_sqrt)

        if self.training and self.track_running_stats:
            self.running_mean = torch.add(
                self.momentum * m.detach(),
                (1 - self.momentum) * self.running_mean,
                out=self.running_mean,
            )
            self.running_variance = torch.add(
                self.momentum * f_cov.detach(),
                (1 - self.momentum) * self.running_variance,
                out=self.running_variance,
            )

        return decorrelated.squeeze(2).squeeze(2)