jiahuimbzuai/CoQuIR huggingface dataset:
//  [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.linalg.solve_triangular(
            torch.linalg.cholesky(f_cov_shrinked),
            eye, 
            upper=False
            )
        
        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)
//  [python]
def triangular_solve(a: Numeric, b: Numeric, lower_a: bool = True):
    return torch.triangular_solve(b, a, upper=not lower_a)[0]
//  [python]
def triangular_solve(a: Numeric, b: Numeric, lower_a: bool = True):
    return torch.linalg.solve_triangular(a, b, upper=not lower_a)
//  [python]
def KL(self):

        # self.q_sqrt_tri stores the triangular matrix using indexes
        #  (0,0), (1,0), (1,1), (2,0), (2,1), (2,2), (3,0)....
        #  knowing this, the diagonal is stored at positions 0, 2, 5, 9, 13...
        #  which can be created using np.cumsum
        diag_indexes = np.cumsum(np.arange(1, self.num_inducing + 1)) - 1
        diag = self.q_sqrt_tri[diag_indexes]
        # Constant dimensionality term
        KL = -0.5 * self.output_dim * self.num_inducing

        # Log of determinant of covariance matrix.
        # Det(Sigma) = Det(q_sqrt q_sqrt^T) = Det(q_sqrt) Det(q_sqrt^T)
        #            = prod(diag_s_sqrt)^2
        KL -= torch.sum(torch.log(torch.abs(diag)))
        
        KL += torch.sum(torch.log(torch.abs(torch.diagonal(self.Lu, dim1=1, dim2=2))))
        \
            
        KL += 0.5 * torch.sum(torch.square(torch.triangular_solve(self.q_sqrt, self.Lu, upper=False)[0]))
        
        Kinv_m = torch.triangular_solve(self.q_mu, self.Lu, upper = False)[0]
        KL += 0.5 * torch.sum(self.q_mu * Kinv_m)
        if self.add_prior_regularization:
            KL += self.generative_function.KL()
        return KL
//  [python]
def KL(self):

        # self.q_sqrt_tri stores the triangular matrix using indexes
        #  (0,0), (1,0), (1,1), (2,0), (2,1), (2,2), (3,0)....
        #  knowing this, the diagonal is stored at positions 0, 2, 5, 9, 13...
        #  which can be created using np.cumsum
        diag_indexes = np.cumsum(np.arange(1, self.num_inducing + 1)) - 1
        diag = self.q_sqrt_tri[diag_indexes]
        # Constant dimensionality term
        KL = -0.5 * self.output_dim * self.num_inducing

        # Log of determinant of covariance matrix.
        # Det(Sigma) = Det(q_sqrt q_sqrt^T) = Det(q_sqrt) Det(q_sqrt^T)
        #            = prod(diag_s_sqrt)^2
        KL -= torch.sum(torch.log(torch.abs(diag)))
        
        KL += torch.sum(torch.log(torch.abs(torch.diagonal(self.Lu, dim1=1, dim2=2))))
        \
            
        KL += 0.5 * torch.sum(torch.square(torch.linalg.solve_triangular(self.Lu, self.q_sqrt, upper=False)))
        
        Kinv_m = torch.linalg.solve_triangular(self.Lu, self.q_mu, upper = False)
        KL += 0.5 * torch.sum(self.q_mu * Kinv_m)
        if self.add_prior_regularization:
            KL += self.generative_function.KL()
        return KL
//  [python]
def solve_triangular_dense(self, matrix, rhs, lower: bool, unit_diagonal: bool):
        matrix, rhs = self.auto_cast(matrix, rhs, int_to_float=True, bool_to_int=True)
        rhs = self.expand_dims(rhs, -1)
        x = torch.triangular_solve(rhs, matrix, upper=not lower, unitriangular=unit_diagonal)[0]
        return x[..., 0]
//  [python]
def solve_triangular_dense(self, matrix, rhs, lower: bool, unit_diagonal: bool):
        matrix, rhs = self.auto_cast(matrix, rhs, int_to_float=True, bool_to_int=True)
        rhs = self.expand_dims(rhs, -1)
        x = torch.linalg.solve_triangular(matrix, rhs, upper=not lower, unitriangular=unit_diagonal)
        return x[..., 0]
//  [python]
def log_prob(self, value):
        if self._validate_args:
            self._validate_sample(value)
        n = self.loc.size(-1)
        y = torch.solve(
            (value - self.loc).unsqueeze(-1), self.scale_tril
        ).solution.squeeze(-1)
        Z = (
            self.scale_tril.diagonal(dim1=-2, dim2=-1).log().sum(-1)
            + 0.5 * n * self.df.log()
            + 0.5 * n * math.log(math.pi)
            + torch.lgamma(0.5 * self.df)
            - torch.lgamma(0.5 * (self.df + n))
        )
        return -0.5 * (self.df + n) * torch.log1p(y.pow(2).sum(-1) / self.df) - Z
//  [python]
def log_prob(self, value):
        if self._validate_args:
            self._validate_sample(value)
        n = self.loc.size(-1)
        y = torch.linalg.solve_triangular(
            self.scale_tril, (value - self.loc).unsqueeze(-1), upper=False
        ).squeeze(-1)
        Z = (
            self.scale_tril.diagonal(dim1=-2, dim2=-1).log().sum(-1)
            + 0.5 * n * self.df.log()
            + 0.5 * n * math.log(math.pi)
            + torch.lgamma(0.5 * self.df)
            - torch.lgamma(0.5 * (self.df + n))
        )
        return -0.5 * (self.df + n) * torch.log1p(y.pow(2).sum(-1) / self.df) - Z
//  [python]
def _triu_inverse(x):
    if x.dim() == 1:
        return x.reciprocal()
    else:
        identity = torch.eye(x.size(-1), dtype=x.dtype, device=x.device)
        return torch.triangular_solve(identity, x, upper=True).solution
//  [python]
def _triu_inverse(x):
    if x.dim() == 1:
        return x.reciprocal()
    else:
        identity = torch.eye(x.size(-1), dtype=x.dtype, device=x.device)
        return torch.linalg.solve_triangular(x, identity, upper=True)
//  [python]
def solve_triangular(a, b, lower=False):
    if b.ndim == a.ndim - 1:
        b = torch.unsqueeze(b, axis=-1)
        return torch.solve(b, a, upper=not lower).solution.squeeze(
            axis=-1
        )
    return torch.solve(b, a, upper=not lower).solution
//  [python]
def solve_triangular(a, b, lower=False):
    if b.ndim == a.ndim - 1:
        b = torch.unsqueeze(b, axis=-1)
        return torch.linalg.solve_triangular(a, b, upper=not lower).squeeze(
            axis=-1
        )
    return torch.linalg.solve_triangular(a, b, upper=not lower)
//  [python]
def precision_matrix(self):
        # We use "Woodbury matrix identity" to take advantage of low rank form::
        #     inv(W @ W.T + D) = inv(D) - inv(D) @ W @ inv(C) @ W.T @ inv(D)
        # where :math:`C` is the capacitance matrix.
        Wt_Dinv = (
            self._unbroadcasted_cov_factor.mT
            / self._unbroadcasted_cov_diag.unsqueeze(-2)
        )
        A = torch.triangular_solve(Wt_Dinv, self._capacitance_tril, upper=False)[0]
        precision_matrix = (
            torch.diag_embed(self._unbroadcasted_cov_diag.reciprocal()) - A.mT @ A
        )
        return precision_matrix.expand(
            self._batch_shape + self._event_shape + self._event_shape
        )
//  [python]
def precision_matrix(self):
        # We use "Woodbury matrix identity" to take advantage of low rank form::
        #     inv(W @ W.T + D) = inv(D) - inv(D) @ W @ inv(C) @ W.T @ inv(D)
        # where :math:`C` is the capacitance matrix.
        Wt_Dinv = (
            self._unbroadcasted_cov_factor.mT
            / self._unbroadcasted_cov_diag.unsqueeze(-2)
        )
        A = torch.linalg.solve_triangular(self._capacitance_tril, Wt_Dinv, upper=False)
        precision_matrix = (
            torch.diag_embed(self._unbroadcasted_cov_diag.reciprocal()) - A.mT @ A
        )
        return precision_matrix.expand(
            self._batch_shape + self._event_shape + self._event_shape
        )
//  [python]
def solvetri(matrix, rhs, lower=True):
    upper = not lower
    out = torch.triangular_solve(rhs, matrix, upper=upper)[0]
    return out
//  [python]
def solvetri(matrix, rhs, lower=True):
    upper = not lower
    out = torch.linalg.solve_triangular(matrix, rhs, upper=upper)
    return out
//  [python]
def inverse_no_cache(self, inputs):
        """
        Cost:

        
            output = O(D^2N)
            logabsdet = O(D)
        

        where:

        
            D = num of features
            N = num of inputs
        
        """
        lower, upper = self._create_lower_upper()
        outputs = inputs - self.bias
        try:
            outputs = torch.triangular_solve(
                outputs.t(), lower, upper=False, unitriangular=True
            )[0]
            outputs = torch.triangular_solve(
                outputs, upper, upper=True, unitriangular=False
            )[0]
        except:
            outputs, _ = torch.triangular_solve(
                outputs.t(), lower, upper=False, unitriangular=True
            )
            outputs, _ = torch.triangular_solve(
                outputs, upper, upper=True, unitriangular=False
            )
        outputs = outputs.t()

        logabsdet = -self.logabsdet()
        logabsdet = logabsdet * inputs.new_ones(outputs.shape[0])

        return outputs, logabsdet
//  [python]
def inverse_no_cache(self, inputs):
        """
        Cost:

        ```
            output = O(D^2N)
            logabsdet = O(D)
        ```

        where:

        ```
            D = num of features
            N = num of inputs
        ```
        """
        lower, upper = self._create_lower_upper()
        outputs = inputs - self.bias
        try:
            outputs = torch.linalg.solve_triangular(
                lower, outputs.t(), upper=False, unitriangular=True
            )
            outputs = torch.linalg.solve_triangular(
                upper, outputs, upper=True, unitriangular=False
            )
        except:
            outputs, _ = torch.triangular_solve(
                outputs.t(), lower, upper=False, unitriangular=True
            )
            outputs, _ = torch.triangular_solve(
                outputs, upper, upper=True, unitriangular=False
            )
        outputs = outputs.t()

        logabsdet = -self.logabsdet()
        logabsdet = logabsdet * inputs.new_ones(outputs.shape[0])

        return outputs, logabsdet
//  [python]
def objective(self,L):
        inv=torch.triangular_solve(self.Y,L)[0]
        return -torch.diag(L).log().mean()-0.5*torch.mean(inv**2)
//  [python]
def objective(self,L):
        inv=torch.linalg.solve_triangular(L,self.Y)
        return -torch.diag(L).log().mean()-0.5*torch.mean(inv**2)
//  [python]
def _cholesky_covariance(chol_precision: torch.Tensor, covariance_type: str) -> torch.Tensor:
    # For complex covariance types, invert the
    if covariance_type in ("tied", "full"):
        num_features = chol_precision.size(-1)
        target = torch.eye(num_features, dtype=chol_precision.dtype, device=chol_precision.device)
        return torch.triangular_solve(target, chol_precision, upper=True)[0].t()

    # Simple covariance type
    return chol_precision.reciprocal()
//  [python]
def _cholesky_covariance(chol_precision: torch.Tensor, covariance_type: str) -> torch.Tensor:
    # For complex covariance types, invert the
    if covariance_type in ("tied", "full"):
        num_features = chol_precision.size(-1)
        target = torch.eye(num_features, dtype=chol_precision.dtype, device=chol_precision.device)
        return torch.linalg.solve_triangular(chol_precision, target, upper=True).t()

    # Simple covariance type
    return chol_precision.reciprocal()
//  [python]
def weight_inverse(self):
        """Cost:
            inverse = O(D^3)
        where:
            D = num of features
        """
        lower, upper = self._create_lower_upper()
        identity = torch.eye(self.features, self.features, device=self.lower_entries.device)
        lower_inverse = torch.triangular_solve(identity, lower, upper=False, unitriangular=True)[0]
        weight_inverse = torch.triangular_solve(lower_inverse, upper, upper=True, unitriangular=False)[0]
        return weight_inverse
//  [python]
def weight_inverse(self):
        """Cost:
            inverse = O(D^3)
        where:
            D = num of features
        """
        lower, upper = self._create_lower_upper()
        identity = torch.eye(self.features, self.features, device=self.lower_entries.device)
        lower_inverse = torch.linalg.solve_triangular(lower, identity, upper=False, unitriangular=True)
        weight_inverse = torch.linalg.solve_triangular(upper, lower_inverse, upper=True, unitriangular=False)
        return weight_inverse
//  [python]
def linearize_sparse_depth_prior(L_mm):
    B, N_train, _ = L_mm.shape
    device = L_mm.device

    da_dd = torch.eye(N_train, device=device).unsqueeze(0).repeat(B, 1, 1)
    dr_dd = torch.triangular_solve(da_dd, L_mm, upper=False).solution
    # Note: This is K_mm_inv
    H_d_d = torch.einsum("hjk,hjl->hkl", dr_dd, dr_dd)

    return dr_dd, H_d_d
//  [python]
def linearize_sparse_depth_prior(L_mm):
    B, N_train, _ = L_mm.shape
    device = L_mm.device

    da_dd = torch.eye(N_train, device=device).unsqueeze(0).repeat(B, 1, 1)
    dr_dd = torch.linalg.solve_triangular(L_mm, da_dd, upper=False)
    # Note: This is K_mm_inv
    H_d_d = torch.einsum("hjk,hjl->hkl", dr_dd, dr_dd)

    return dr_dd, H_d_d
//  [python]
def backward(ctx, dx):
        # check for None tensors
        if dx is None:
            return None, None

        # unpack cached tensors
        A, b, x, R = ctx.saved_tensors
        B, M, N = A.shape

        if R is None:
            _, R = torch.linalg.qr(A, mode='r')

        dA, db = None, None

        w = torch.triangular_solve(torch.triangular_solve(dx, torch.transpose(R, 2, 1), upper=False).solution, R, upper=True).solution
        Aw = torch.bmm(A, w)

        if ctx.needs_input_grad[0]:
            r = b - torch.bmm(A, x)
            dA = torch.bmm(r.view(B, M, 1), w.view(B, 1, N)) - torch.bmm(Aw.view(B, M, 1), x.view(B, 1, N))
        if ctx.needs_input_grad[1]:
            db = Aw

        # return gradients (None for non-tensor inputs)
        return dA, db, None
//  [python]
def backward(ctx, dx):
        # check for None tensors
        if dx is None:
            return None, None

        # unpack cached tensors
        A, b, x, R = ctx.saved_tensors
        B, M, N = A.shape

        if R is None:
            _, R = torch.linalg.qr(A, mode='r')

        dA, db = None, None

        w = torch.linalg.solve_triangular(R, torch.linalg.solve_triangular(torch.transpose(R, 2, 1), dx, upper=False), upper=True)
        Aw = torch.bmm(A, w)

        if ctx.needs_input_grad[0]:
            r = b - torch.bmm(A, x)
            dA = torch.bmm(r.view(B, M, 1), w.view(B, 1, N)) - torch.bmm(Aw.view(B, M, 1), x.view(B, 1, N))
        if ctx.needs_input_grad[1]:
            db = Aw

        # return gradients (None for non-tensor inputs)
        return dA, db, None
//  [python]
def _forward(a, eye):
        M = a.size(-1) - 1
        K, a = remove_gain(a, return_gain=True)

        U = [a.flip(-1)]
        E = [K**2]
        for m in range(M):
            u0 = U[-1][..., :1]
            u1 = U[-1][..., 1 : M - m]
            t = 1 / (1 - u0**2)
            u = (u1 - u0 * u1.flip(-1)) * t
            u = F.pad(u, (0, m + 2))
            e = E[-1] * t
            U.append(u)
            E.append(e)
        U = torch.stack(U[::-1], dim=-1)
        E = torch.stack(E[::-1], dim=-1)

        V = torch.triangular_solve(eye, U, upper=True, unitriangular=True)[0]
        r = torch.matmul(V[..., :1].transpose(-2, -1) * E, V).squeeze(-2)
        return r
//  [python]
def _forward(a, eye):
        M = a.size(-1) - 1
        K, a = remove_gain(a, return_gain=True)

        U = [a.flip(-1)]
        E = [K**2]
        for m in range(M):
            u0 = U[-1][..., :1]
            u1 = U[-1][..., 1 : M - m]
            t = 1 / (1 - u0**2)
            u = (u1 - u0 * u1.flip(-1)) * t
            u = F.pad(u, (0, m + 2))
            e = E[-1] * t
            U.append(u)
            E.append(e)
        U = torch.stack(U[::-1], dim=-1)
        E = torch.stack(E[::-1], dim=-1)

        V = torch.linalg.solve_triangular(U, eye, upper=True, unitriangular=True)
        r = torch.matmul(V[..., :1].transpose(-2, -1) * E, V).squeeze(-2)
        return r
//  [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)
        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)
        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 predict(self, X):
        if self.x_train is None or self.y_train is None:
            raise RuntimeError("Call forward at least once before calling predict.")
        with torch.autograd.no_grad():
            L, A, AAT, LB, c = self._calc_intermediate(self.x_train, self.y_train)
            kms = self.kernel(self.centers, X)
            tmp1 = torch.triangular_solve(kms, L, upper=False)[0]
            tmp2 = torch.triangular_solve(tmp1, LB, upper=False)[0]
            return tmp2.T @ c
//  [python]
def predict(self, X):
        if self.x_train is None or self.y_train is None:
            raise RuntimeError("Call forward at least once before calling predict.")
        with torch.autograd.no_grad():
            L, A, AAT, LB, c = self._calc_intermediate(self.x_train, self.y_train)
            kms = self.kernel(self.centers, X)
            tmp1 = torch.linalg.solve_triangular(L, kms, upper=False)
            tmp2 = torch.linalg.solve_triangular(LB, tmp1, upper=False)
            return tmp2.T @ c
//  [python]
def _criterion(self, dxs: Tuple[Tensor], dgs: Tuple[Tensor]):
        dx = parameters_to_vector(dxs)
        dg = parameters_to_vector(dgs)
        Q = self.cholesky_factor

        res = dx.dot(torch.solve(dx.unsqueeze(-1), torch.triangular_solve(dx.unsqueeze(-1), Q.T, upper=False).solution).squeeze())
        res += dg.dot(Q.T.mv(Q.mv(dg)))
        return float(res)
//  [python]
def _criterion(self, dxs: Tuple[Tensor], dgs: Tuple[Tensor]):
        dx = parameters_to_vector(dxs)
        dg = parameters_to_vector(dgs)
        Q = self.cholesky_factor

        res = dx.dot(solve_triangular(Q, solve_triangular(Q.T, dx.unsqueeze(-1), upper=False), upper=True).squeeze())
        res += dg.dot(Q.T.mv(Q.mv(dg)))
        return float(res)
//  [python]
def _update_preconditioner(self, dxs: List[Tensor], dgs: List[Tensor], eps=1.2e-38):
        dx = parameters_to_vector(dxs)
        dg = parameters_to_vector(dgs)
        Q = self.cholesky_factor

        a = Q.mv(dg)
        b = torch.solve(dx.unsqueeze(-1), Q.T, upper=False).squeeze()

        grad = torch.triu(torch.outer(a, a) - torch.outer(b, b))
        lr = self.precond_lr / (grad.abs().max() + eps)

        Q.sub_(grad.mm(Q), alpha=float(lr))
//  [python]
def _update_preconditioner(self, dxs: List[Tensor], dgs: List[Tensor], eps=1.2e-38):
        dx = parameters_to_vector(dxs)
        dg = parameters_to_vector(dgs)
        Q = self.cholesky_factor

        a = Q.mv(dg)
        b = solve_triangular(Q.T, dx.unsqueeze(-1), upper=False).squeeze()

        grad = torch.triu(torch.outer(a, a) - torch.outer(b, b))
        lr = self.precond_lr / (grad.abs().max() + eps)

        Q.sub_(grad.mm(Q), alpha=float(lr))
//  [python]
def calc_trace_fwd(
    init_val: torch.Tensor,
    k_mn: Optional[torch.Tensor],
    k_mn_zy: Optional[torch.Tensor],
    kmm_chol: torch.Tensor,
    X: Optional[torch.Tensor],
    t: Optional[int],
    trace_type: str,
):
    """Nystrom kernel trace forward"""
    if trace_type == "ste":
        assert k_mn_zy is not None and t is not None, "Incorrect arguments to trace_fwd"
        solve1 = torch.triangular_solve(k_mn_zy[:, :t], kmm_chol, upper=False)[0]  # m * t
        solve2 = torch.triangular_solve(solve1, kmm_chol.T, upper=True)[0].contiguous()  # m * t
        init_val -= solve1.square_().sum(0).mean()
    elif trace_type == "direct":
        assert k_mn is not None, "Incorrect arguments to trace_fwd"
        solve1 = trsm(k_mn, kmm_chol, 1.0, lower=True, transpose=False)  # (M*N)
        solve2 = trsm(solve1, kmm_chol, 1.0, lower=True, transpose=True)  # (M*N)
        init_val -= solve1.square_().sum()
    elif trace_type == "fast":
        assert k_mn_zy is not None and t is not None, "Incorrect arguments to trace_fwd"
        k_subs = k_mn_zy
        assert k_subs.shape == (kmm_chol.shape[0], t), "Shape incorrect"  # m * t
        solve1 = torch.triangular_solve(k_subs, kmm_chol, upper=False)[0]  # m * t
        solve2 = torch.triangular_solve(solve1, kmm_chol.T, upper=True)[0].contiguous()  # m * t
        norm = X.shape[0] / t
        init_val -= solve1.square_().sum() * norm
    else:
        raise ValueError("Trace-type %s unknown" % (trace_type))
    return init_val, solve2
//  [python]
def calc_trace_fwd(
    init_val: torch.Tensor,
    k_mn: Optional[torch.Tensor],
    k_mn_zy: Optional[torch.Tensor],
    kmm_chol: torch.Tensor,
    X: Optional[torch.Tensor],
    t: Optional[int],
    trace_type: str,
):
    """Nystrom kernel trace forward"""
    if trace_type == "ste":
        assert k_mn_zy is not None and t is not None, "Incorrect arguments to trace_fwd"
        solve1 = torch.linalg.solve_triangular(kmm_chol, k_mn_zy[:, :t], upper=False)  # m * t
        solve2 = torch.linalg.solve_triangular(kmm_chol.T, solve1, upper=True).contiguous()  # m * t
        init_val -= solve1.square_().sum(0).mean()
    elif trace_type == "direct":
        assert k_mn is not None, "Incorrect arguments to trace_fwd"
        solve1 = trsm(k_mn, kmm_chol, 1.0, lower=True, transpose=False)  # (M*N)
        solve2 = trsm(solve1, kmm_chol, 1.0, lower=True, transpose=True)  # (M*N)
        init_val -= solve1.square_().sum()
    elif trace_type == "fast":
        assert k_mn_zy is not None and t is not None, "Incorrect arguments to trace_fwd"
        k_subs = k_mn_zy
        assert k_subs.shape == (kmm_chol.shape[0], t), "Shape incorrect"  # m * t
        solve1 = torch.linalg.solve_triangular(kmm_chol, k_subs, upper=False)  # m * t
        solve2 = torch.linalg.solve_triangular(kmm_chol.T, solve1, upper=True).contiguous()  # m * t
        norm = X.shape[0] / t
        init_val -= solve1.square_().sum() * norm
    else:
        raise ValueError("Trace-type %s unknown" % (trace_type))
    return init_val, solve2
//  [python]
def _calc_intermediate(self, X, Y):
        variance = self.penalty * X.shape[0]
        sqrt_var = torch.sqrt(variance)

        kmn = self.kernel(self.centers, X)
        kmm = self.kernel(self.centers, self.centers)
        L = jittering_cholesky(kmm)

        # A = L^{-1} K_mn / (sqrt(n*pen))
        A = torch.triangular_solve(kmn, L, upper=False).solution / sqrt_var
        AAT = A @ A.T
        # B = A @ A.T + I
        B = AAT + torch.eye(AAT.shape[0], device=X.device, dtype=X.dtype)
        LB = jittering_cholesky(B)  # LB @ LB.T = B
        AY = A @ Y
        c = torch.triangular_solve(AY, LB, upper=False).solution / sqrt_var

        return L, A, AAT, LB, c
//  [python]
def _calc_intermediate(self, X, Y):
        variance = self.penalty * X.shape[0]
        sqrt_var = torch.sqrt(variance)

        kmn = self.kernel(self.centers, X)
        kmm = self.kernel(self.centers, self.centers)
        L = jittering_cholesky(kmm)

        # A = L^{-1} K_mn / (sqrt(n*pen))
        A = torch.linalg.solve_triangular(L, kmn, upper=False) / sqrt_var
        AAT = A @ A.T
        # B = A @ A.T + I
        B = AAT + torch.eye(AAT.shape[0], device=X.device, dtype=X.dtype)
        LB = jittering_cholesky(B)  # LB @ LB.T = B
        AY = A @ Y
        c = torch.linalg.solve_triangular(LB, AY, upper=False) / sqrt_var

        return L, A, AAT, LB, c
//  [python]
def test_triangular_solve_large(self, device, dtype):
        # Repro for https://github.com/pytorch/pytorch/issues/79191
        A = torch.randn(1, 2, 2, device=device, dtype=dtype).tril_()
        B = torch.randn(1, 2, 524281, device=device, dtype=dtype)
        X = torch.triangular_solve(B, A, upper=False)[0]
        self.assertEqual(A @ X, B)
//  [python]
def test_triangular_solve_large(self, device, dtype):
        # Repro for https://github.com/pytorch/pytorch/issues/79191
        A = torch.randn(1, 2, 2, device=device, dtype=dtype).tril_()
        B = torch.randn(1, 2, 524281, device=device, dtype=dtype)
        X = torch.linalg.solve_triangular(A, B, upper=False)
        self.assertEqual(A @ X, B)
//  [python]
def get_obs_info(L, K_mn):
  obs_info = torch.triangular_solve(K_mn, L, upper=False)[0]
  return obs_info
//  [python]
def get_obs_info(L, K_mn):
  obs_info = torch.linalg.solve_triangular(L, K_mn, upper=False)
  return obs_info
//  [python]
def inverse(self, y: TensorType[..., 'dim'], **kwargs) -> TensorType[..., 'dim']:
        x = y - self.bias
        x = torch.triangular_solve(x, self.U, upper=True, transpose=True)[0]
        x = torch.triangular_solve(x, self.L, upper=False, transpose=True)[0]
        return x
//  [python]
def inverse(self, y: TensorType[..., 'dim'], **kwargs) -> TensorType[..., 'dim']:
        x = y - self.bias
        x = torch.linalg.solve_triangular(self.U, x, upper=True, left=False)
        x = torch.linalg.solve_triangular(self.L, x, upper=False, left=False)
        return x
//  [python]
def inverse_no_cache(self, inputs):
        """Cost:
            output = O(D^2N)
            logabsdet = O(D)
        where:
            D = num of features
            N = num of inputs
        """
        lower, upper = self._create_lower_upper()
        outputs = inputs - self.bias
        print("outputs.t().shape ", outputs.t().shape)
        print(lower)
        print("lower.shape ", lower.shape)

        outputs, _ = torch.triangular_solve(
             outputs.t(), lower, upper=False, unitriangular=True
        )
        outputs, _ = torch.triangular_solve(
            outputs, upper, upper=True, unitriangular=False
        )
        # outputs, _ = torch.linalg.solve_triangular(
        #     lower, outputs.t(), upper=False, unitriangular=True
        # )
        # outputs, _ = torch.linalg.solve_triangular(
        #     upper,outputs, upper=True, unitriangular=False
        # )
        outputs = outputs.t()

        logabsdet = -self.logabsdet()
        logabsdet = logabsdet * inputs.new_ones(outputs.shape[0])

        return outputs, logabsdet
//  [python]
def inverse_no_cache(self, inputs):
        """Cost:
            output = O(D^2N)
            logabsdet = O(D)
        where:
            D = num of features
            N = num of inputs
        """
        lower, upper = self._create_lower_upper()
        outputs = inputs - self.bias
        print("outputs.t().shape ", outputs.t().shape)
        print(lower)
        print("lower.shape ", lower.shape)

        outputs, _ = torch.linalg.solve_triangular(
             lower, outputs.t(), upper=False, unitriangular=True
        )
        outputs, _ = torch.linalg.solve_triangular(
            upper,outputs, upper=True, unitriangular=False
        )
        # outputs, _ = torch.triangular_solve(
        #     outputs.t(), lower, upper=False, unitriangular=True
        # )
        # outputs, _ = torch.triangular_solve(
        #     outputs, upper, upper=True, unitriangular=False
        # )
        outputs = outputs.t()

        logabsdet = -self.logabsdet()
        logabsdet = logabsdet * inputs.new_ones(outputs.shape[0])

        return outputs, logabsdet
//  [python]
def forward(
        self,
        x_test: torch.Tensor,
        noise_free: bool = False,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        r"""Make prediction.

        Args:
            x_test (torch.Tensor): Test inputs of shape
                (num_inputs, dim_inputs).
            noise_free (bool, optional): If True, predict the latent function
                values. Otherwise, predict the noisy targets.

        Returns:
            mean (torch.Tensor): Predictive mean of shape (num_inputs, 1)
            std (torch.Tensor): Predictive standard deviation of shape
                (num_inputs, 1).

        ??? note "Difference between `forward` and `predict`"

            The `forward` method uses PyTorch tensors as inputs and outputs.
            The `predict` method uses NumPy arrays as inputs and outputs.
            Users should use `predict` method for model prediction.

        """
        with torch.no_grad():
            L, iK_y = self._compute_common()
            Ksn = self.kernel(x_test, self.x_train)
            Kss_diag = self.kernel.diag(x_test)
            iL_Kns = torch.triangular_solve(Ksn.t(), L, upper=False)[0]
            mean = Ksn @ iK_y
            var = Kss_diag - iL_Kns.square().sum(0).view(-1, 1)
            # Variance might be zero when lengthscale is too large.
            if torch.any(var <= 0.0):
                print(var.ravel().numpy())
                raise ValueError("Predictive variance <= 0.0!")
            var.clamp_(min=self.jitter)
            if not noise_free:
                var += self.noise
            std = var.sqrt()
        return mean, std
//  [python]
def forward(
        self,
        x_test: torch.Tensor,
        noise_free: bool = False,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        r"""Make prediction.

        Args:
            x_test (torch.Tensor): Test inputs of shape
                (num_inputs, dim_inputs).
            noise_free (bool, optional): If True, predict the latent function
                values. Otherwise, predict the noisy targets.

        Returns:
            mean (torch.Tensor): Predictive mean of shape (num_inputs, 1)
            std (torch.Tensor): Predictive standard deviation of shape
                (num_inputs, 1).

        ??? note "Difference between `forward` and `predict`"

            The `forward` method uses PyTorch tensors as inputs and outputs.
            The `predict` method uses NumPy arrays as inputs and outputs.
            Users should use `predict` method for model prediction.

        """
        with torch.no_grad():
            L, iK_y = self._compute_common()
            Ksn = self.kernel(x_test, self.x_train)
            Kss_diag = self.kernel.diag(x_test)
            iL_Kns = torch.linalg.solve_triangular(L, Ksn.t(), upper=False)
            mean = Ksn @ iK_y
            var = Kss_diag - iL_Kns.square().sum(0).view(-1, 1)
            # Variance might be zero when lengthscale is too large.
            if torch.any(var <= 0.0):
                print(var.ravel().numpy())
                raise ValueError("Predictive variance <= 0.0!")
            var.clamp_(min=self.jitter)
            if not noise_free:
                var += self.noise
            std = var.sqrt()
        return mean, std
//  [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)
    step0 = step/(norm_lower_bound(grad) + _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)
    step0 = step/(norm_lower_bound(grad) + _tiny)
              
    return Q - step0*grad.mm(Q)
//  [python]
def solve(self, z):
        return torch.triangular_solve(z.T, self.mat(), upper=False, unitriangular=self.unitriangular)[0].T
//  [python]
def solve(self, z):
        return torch.linalg.solve_triangular(self.mat(), z.T, upper=False, unitriangular=self.unitriangular).T
//  [python]
def _update_precond_dense_dense(Ql, Qr, dX, dG, step=0.01, _tiny=1.2e-38):
    # type: (Tensor, Tensor, Tensor, Tensor, float, float) -> None
    """
    update Kronecker product preconditioner P = kron_prod(Qr^T*Qr, Ql^T*Ql)
    Ql: (left side) Cholesky factor of preconditioner with positive diagonal entries
    Qr: (right side) Cholesky factor of preconditioner with positive diagonal entries
    dX: perturbation of (matrix) parameter
    dG: perturbation of (matrix) gradient
    step: update step size normalized to range [0, 1]
    _tiny: an offset to avoid division by zero
    """
    max_l = torch.max(torch.diag(Ql))
    max_r = torch.max(torch.diag(Qr))

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

    A = Ql.mm(dG.mm(Qr.t()))
    Bt = torch.triangular_solve(torch.triangular_solve(dX.t(), Qr, upper=True).solution, Ql, upper=True).solution.t()

    grad1 = torch.triu(A.mm(A.T) - Bt.mm(Bt.T))
    grad2 = torch.triu(A.T.mm(A) - Bt.T.mm(Bt))

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

    Ql.sub_(grad1.mm(Ql), alpha=float(step1))
    Qr.sub_(grad2.mm(Qr), alpha=float(step2))
//  [python]
def _update_precond_dense_dense(Ql, Qr, dX, dG, step=0.01, _tiny=1.2e-38):
    # type: (Tensor, Tensor, Tensor, Tensor, float, float) -> None
    """
    update Kronecker product preconditioner P = kron_prod(Qr^T*Qr, Ql^T*Ql)
    Ql: (left side) Cholesky factor of preconditioner with positive diagonal entries
    Qr: (right side) Cholesky factor of preconditioner with positive diagonal entries
    dX: perturbation of (matrix) parameter
    dG: perturbation of (matrix) gradient
    step: update step size normalized to range [0, 1]
    _tiny: an offset to avoid division by zero
    """
    max_l = torch.max(torch.diag(Ql))
    max_r = torch.max(torch.diag(Qr))

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

    A = Ql.mm(dG.mm(Qr.t()))
    Bt = solve_triangular(Ql, solve_triangular(Qr, dX.t(), upper=True).t(), upper=True)

    grad1 = torch.triu(A.mm(A.T) - Bt.mm(Bt.T))
    grad2 = torch.triu(A.T.mm(A) - Bt.T.mm(Bt))

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

    Ql.sub_(grad1.mm(Ql), alpha=float(step1))
    Qr.sub_(grad2.mm(Qr), alpha=float(step2))
//  [python]
def forward(
        self,
        x: TensorType[..., 'dim'],
        t: Union[Number, TensorType[..., 1]] = 1.0,
        *,
        reverse: bool = False,
        **kwargs
    ) -> TensorType[..., 'dim']:
        t = self.get_time(t, x.shape)

        if reverse:
            t = -t
            if self.bias is not None:
                x = x - self.bias

        L, U = self.lu()

        x = torch.triangular_solve(x.unsqueeze(-1), L, upper=False, unitriangular=True)[0].squeeze(-1)
        x = torch.triangular_solve(x.unsqueeze(-1), U, upper=True, unitriangular=False)[0].squeeze(-1)

        x = x * (self.diag * t).exp()

        x = F.linear(x, U)
        x = F.linear(x, L)

        if reverse is False and self.bias is not None:
            x = x + self.bias
        return x
//  [python]
def forward(
        self,
        x: TensorType[..., 'dim'],
        t: Union[Number, TensorType[..., 1]] = 1.0,
        *,
        reverse: bool = False,
        **kwargs
    ) -> TensorType[..., 'dim']:
        t = self.get_time(t, x.shape)

        if reverse:
            t = -t
            if self.bias is not None:
                x = x - self.bias

        L, U = self.lu()

        x = torch.linalg.solve_triangular(L, x.unsqueeze(-1), upper=False, unitriangular=True).squeeze(-1)
        x = torch.linalg.solve_triangular(U, x.unsqueeze(-1), upper=True, unitriangular=False).squeeze(-1)

        x = x * (self.diag * t).exp()

        x = F.linear(x, U)
        x = F.linear(x, L)

        if reverse is False and self.bias is not None:
            x = x + self.bias
        return x
//  [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)
    step1 = step/(norm_lower_bound(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)
    step1 = step/(norm_lower_bound(grad1) + _tiny)
    step2 = step/(torch.max(torch.abs(grad2)) + _tiny)
        
    return Ql - step1*grad1.mm(Ql), qr - step2*grad2*qr
//  [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)

        """
        y_flat = y.flatten(start_dim=-2)
        LUx = (y_flat.unsqueeze(-3) * self.permutation.T.unsqueeze(-1)).sum(-2)
        Ux = torch.triangular_solve(LUx, self.L, upper=False)[0]
        x = torch.triangular_solve(Ux, self.U, upper=True)[0]
        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)

        """
        y_flat = y.flatten(start_dim=-2)
        LUx = (y_flat.unsqueeze(-3) * self.permutation.T.unsqueeze(-1)).sum(-2)
        Ux = torch.linalg.solve_triangular(self.L, LUx, upper=False)
        x = torch.linalg.solve_triangular(self.U, Ux, upper=True)
        return x.reshape(x.shape[:-1] + y.shape[-2:])
//  [python]
def solve(
        self: Float[LinearOperator, "... N N"],
        right_tensor: Union[Float[Tensor, "... N P"], Float[Tensor, " N"]],
        left_tensor: Optional[Float[Tensor, "... O N"]] = None,
    ) -> Union[Float[Tensor, "... N P"], Float[Tensor, "... N"], Float[Tensor, "... O P"], Float[Tensor, "... O"]]:
        squeeze = False
        if right_tensor.dim() == 1:
            right_tensor = right_tensor.unsqueeze(-1)
            squeeze = True

        if isinstance(self._tensor, DenseLinearOperator):
            res = torch.solve(right_tensor, self.to_dense())[1]
        elif isinstance(self._tensor, BatchRepeatLinearOperator):
            res = self._tensor.base_linear_op.solve(right_tensor, left_tensor)
            # TODO: Proper broadcasting
            res = res.expand(self._tensor.batch_repeat + res.shape[-2:])
        else:
            # TODO: Can we be smarter here?
            res = self._tensor.solve(right_tensor=right_tensor, left_tensor=left_tensor)

        if squeeze:
            res = res.squeeze(-1)

        if left_tensor is not None:
            res = left_tensor @ res
        return res
//  [python]
def solve(
        self: Float[LinearOperator, "... N N"],
        right_tensor: Union[Float[Tensor, "... N P"], Float[Tensor, " N"]],
        left_tensor: Optional[Float[Tensor, "... O N"]] = None,
    ) -> Union[Float[Tensor, "... N P"], Float[Tensor, "... N"], Float[Tensor, "... O P"], Float[Tensor, "... O"]]:
        squeeze = False
        if right_tensor.dim() == 1:
            right_tensor = right_tensor.unsqueeze(-1)
            squeeze = True

        if isinstance(self._tensor, DenseLinearOperator):
            res = torch.linalg.solve_triangular(self.to_dense(), right_tensor, upper=self.upper)
        elif isinstance(self._tensor, BatchRepeatLinearOperator):
            res = self._tensor.base_linear_op.solve(right_tensor, left_tensor)
            # TODO: Proper broadcasting
            res = res.expand(self._tensor.batch_repeat + res.shape[-2:])
        else:
            # TODO: Can we be smarter here?
            res = self._tensor.solve(right_tensor=right_tensor, left_tensor=left_tensor)

        if squeeze:
            res = res.squeeze(-1)

        if left_tensor is not None:
            res = left_tensor @ res
        return res
//  [python]
def predict(self, X):
        if self.x_train is None or self.y_train is None:
            raise RuntimeError("Call forward at least once before calling predict.")
        with torch.autograd.no_grad():
            L, A, AAT, LB, c = self._calc_intermediate(self.x_train, self.y_train)
            tmp1 = torch.triangular_solve(c, LB.T, upper=True)[0]
            tmp2 = torch.triangular_solve(tmp1, L.T, upper=True)[0]
            kms = self.kernel(self.centers, X)
            return kms.T @ tmp2
//  [python]
def predict(self, X):
        if self.x_train is None or self.y_train is None:
            raise RuntimeError("Call forward at least once before calling predict.")
        with torch.autograd.no_grad():
            L, A, AAT, LB, c = self._calc_intermediate(self.x_train, self.y_train)
            tmp1 = torch.linalg.solve_triangular(LB.T, c, upper=True)
            tmp2 = torch.linalg.solve_triangular(L.T, tmp1, upper=True)
            kms = self.kernel(self.centers, X)
            return kms.T @ tmp2
//  [python]
def forward(self, X, Y):
        self.x_train, self.y_train = X.detach(), Y.detach()
        variance = self.penalty * X.shape[0]
        sqrt_var = torch.sqrt(variance)
        Kdiag = self.kernel(X, X, diag=True).sum()

        L, A, AAT, LB, c = self._calc_intermediate(X, Y)
        C = torch.triangular_solve(A, LB, upper=False)[0]  # m * n

        datafit = torch.square(Y).sum() - torch.square(c / sqrt_var).sum()
        ndeff = (C / sqrt_var).square().sum()
        trace = Kdiag - torch.trace(AAT)
        trace = trace * datafit / (variance * X.shape[0])
        self._save_losses(ndeff, datafit, trace)

        return ndeff + datafit + trace
//  [python]
def forward(self, X, Y):
        self.x_train, self.y_train = X.detach(), Y.detach()
        variance = self.penalty * X.shape[0]
        sqrt_var = torch.sqrt(variance)
        Kdiag = self.kernel(X, X, diag=True).sum()

        L, A, AAT, LB, c = self._calc_intermediate(X, Y)
        C = torch.linalg.solve_triangular(LB, A, upper=False)  # m * n

        datafit = torch.square(Y).sum() - torch.square(c / sqrt_var).sum()
        ndeff = (C / sqrt_var).square().sum()
        trace = Kdiag - torch.trace(AAT)
        trace = trace * datafit / (variance * X.shape[0])
        self._save_losses(ndeff, datafit, trace)

        return ndeff + datafit + trace
//  [python]
def forward(self, x, T=5, eps=1e-5):
        N, G, C, H, W = x.size()
        x_in = x.transpose(0, 1).contiguous().view(G, -1)
        if self.training:
            mean = x_in.mean(-1, keepdim=True)
            xc = x_in - mean
            d, m = x_in.size()
            P = [None] * (T + 1)
            P[0] = torch.eye(G, device=x.device)
            Sigma = (torch.matmul(xc, xc.transpose(0, 1))) / float(m) + P[0] * eps
            rTr = (Sigma * P[0]).sum([0, 1], keepdim=True).reciprocal()
            Sigma_N = Sigma * rTr
            wm = torch.triangular_solve(
                P[0], torch.cholesky(Sigma_N), upper=False
            )[0]
            self.running_wm = self.momentum * self.running_wm + (1 - self.momentum) * wm
        else:
            wm = self.running_wm

        x_out = wm @ x_in
        x_out = x_out.view(G, N, C, H, W).transpose(0, 1).contiguous()

        return x_out