jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def get_alpha(self):
        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(c, LB, upper=False, left=True, transpose=True)
            tmp2 = torch.linalg.solve_triangular(tmp1, L, upper=False, left=True, transpose=True)
            return tmp2
//  [python]
def negative_log_likelihood(self, hyper=None):
        if hyper is not None:
            param_original = self.model.param_to_vec()
            self.cholesky_update(hyper)

        # cholesky is lower triangular matrix
        mean_vec_sol = torch.triangular_solve(self.mean_vec, self.cholesky, upper=False)[0]
        nll = 0.5 * torch.sum(mean_vec_sol ** 2) + torch.sum(torch.log(torch.diag(self.cholesky))) + 0.5 * self.train_y.size(0) * np.log(2 * np.pi)
        if hyper is not None:
            self.cholesky_update(param_original)
        return nll
//  [python]
def negative_log_likelihood(self, hyper=None):
        if hyper is not None:
            param_original = self.model.param_to_vec()
            self.cholesky_update(hyper)

        # cholesky is lower triangular matrix
        mean_vec_sol = torch.linalg.solve_triangular(self.cholesky, self.mean_vec, upper=False)
        nll = 0.5 * torch.sum(mean_vec_sol ** 2) + torch.sum(torch.log(torch.diag(self.cholesky))) + 0.5 * self.train_y.size(0) * np.log(2 * np.pi)
        if hyper is not None:
            self.cholesky_update(param_original)
        return nll
//  [python]
def inverse_no_cache(self, inputs, context=None):
        """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
        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, context=None):
        """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
        outputs, _ = torch.linalg.solve_triangular(
            outputs.t(), lower, upper=False, unitriangular=True
        )
        outputs, _ = torch.linalg.solve_triangular(
            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 compute_loss(self, X, Y, I0, idx):
        ep = 1e-7 # to be used if a normalized loss is considered (check the variable "res")
        
        p0 = gamma
        p1 = delta
        
        E0l = 3*I0*self.sigmoid_in(self.E0[idx,:,:])
        S0 = 1-I0-E0l
        
        # compute I(t) from R(t)+I(t)
        bIA = torch.cat( ( (1-(self.Ntf-1)*self.dt/2*p0)*I0 - Y[:,0:1,None] , torch.zeros(self.batch_county,self.Nt-2,1).type(self.dtype) ), 1)
        
        IAm = torch.triangular_solve(bIA + torch.matmul( self.expli , Y[:,1:,None] ), self.expli + (self.Ntf-1)*self.dt*p0*self.trap, upper=False)[0]
        If = torch.cat( ( I0 , IAm ), 1)

        beta  = self.forward_pass(X, self.Nt)
        
        # Build RHS of S equation
        bS = torch.cat( ( S0*( 1 - (self.Ntf-1)*self.dt/2*beta[:,0:1,:]*If[:,0:1,:] ) , torch.zeros(self.batch_county,self.Nt-2,1).type(self.dtype) ), 1) # size: batch_county x Nt-1 x 1 
        # Build matrix of S equation
        betaIm = self.matrix_diag ( (beta[:,1:,:]*If[:,1:,:])[:,:,0] ) # size: batch_county x Nt-1 x Nt-1
        betaIm1 = self.matrix_diag( (beta[:,1:-1,:]*If[:,1:-1,:])[:,:,0] ) # size: batch_county x Nt-2 x Nt-2
        betaIm1f = torch.cat( (torch.zeros(self.batch_county,1,self.Nt-2).type(self.dtype),betaIm1) , 1 ) # size: batch_county x Nt-1 x Nt-2
        betaIm1ff = torch.cat( (betaIm1f,torch.zeros(self.batch_county,self.Nt-1,1).type(self.dtype)) , 2 ) # size: batch_county x Nt-1 x Nt-1
        # S prediction
        Sm = torch.triangular_solve(bS, self.expli +(self.Ntf-1)*self.dt/2*(betaIm+betaIm1ff), upper=False)[0] # size: batch_county x Nt-1 x 1 
        
        # Build RHS of E equation
        bE = torch.cat( ( E0l + (self.Ntf-1)*self.dt/2*( S0*beta[:,0:1,:]*If[:,0:1,:] - p1*E0l ) , torch.zeros(self.batch_county,self.Nt-2,1).type(self.dtype) ), 1) # size: batch_county x Nt-1 x 1 
        betaSIm = beta[:,1:,:]*If[:,1:,:]*Sm # size: batch_county x Nt-1 x 1 
        # E prediction
        Em = torch.triangular_solve(bE + (self.Ntf-1)*self.dt*torch.matmul(self.trap,betaSIm), self.expli + (self.Ntf-1)*self.dt*p1*self.trap, upper=False)[0] # size: batch_county x Nt-1 x 1 
        
        # Build RHS of I equation
        bI = torch.cat( ( If[:,0:1,:] + (self.Ntf-1)*self.dt/2* (p1*E0l-p0*If[:,0:1,:]) , torch.zeros(self.batch_county,self.Nt-2,1).type(self.dtype) ), 1)
        if self.is_pred == 1:
            # I prediction
            pred = torch.triangular_solve(bI + (self.Ntf-1)*self.dt*p1*torch.matmul(self.trap,Em), self.expli + (self.Ntf-1)*self.dt*p0*self.trap, upper=False)[0] # size: batch_county x Nt-1 x 1 
            res = (If[:,1:,:] - pred) #/ (torch.mean(If[:,1:,0],dim=1)[:,None,None] + ep) # size: batch_county x Nt-1 x 1 
        else:
            # I equation residual
            res = ( torch.matmul( (self.expli + (self.Ntf-1)*self.dt*p0*self.trap) , If[:,1:,:] ) - ( bI + (self.Ntf-1)*self.dt*p1*torch.matmul(self.trap,Em) ) ) #/ (torch.mean(If[:,1:,0]-If[:,:-1,0],dim=1)[:,None,None] + ep) # size: batch_county x Nt-1 x 1    
            
        loss = torch.mean(res[:,torch.randperm(self.Nt-1)[0:self.batch_size],:]**2)    
        
        return loss
//  [python]
def _get_xk(self, y):
        if y is None:
            return self.x0
        k = self.arnoldi.iter
        if k > 0:
            yy = torch.triangular_solve(y, self.R[:, :k, :k]).solution
            yk = torch.bmm(self.V[:, :, :k], yy)
            # identity = torch.stack([torch.eye((self.x0.shape[1]))]*self.x0.shape[0], dim=0)
            return self.x0 + yk.squeeze()#torch.bmm(identity, yk)
        return self.x0
//  [python]
def _get_xk(self, y):
        if y is None:
            return self.x0
        k = self.arnoldi.iter
        if k > 0:
            yy = torch.linalg.solve_triangular(self.R[:, :k, :k], y)
            yk = torch.bmm(self.V[:, :, :k], yy)
            # identity = torch.stack([torch.eye((self.x0.shape[1]))]*self.x0.shape[0], dim=0)
            return self.x0 + yk.squeeze()#torch.bmm(identity, yk)
        return self.x0
//  [python]
def forward(self, x1):
        v1 = self.conv(x1)
        v2 = torch.transpose(v1, 1, 1)
        v3 = torch.inverse(v2)
        v4 = torch.triangular_solve(v3.contiguous(), torch.cholesky(v1.contiguous()))
        v5 = v4.solution
        v6 = torch.cholesky(v1.contiguous())
        v7 = torch.transpose(v1, 1, 1)
        v8 = torch.inverse(v2)
        v9 = torch.triangular_solve(v8.contiguous(), v6)
        v10 = v9.solution
        v11 = v5 + 1
        v12 = v10 + 1
        v13 = v7 * v11
        v14 = v12 * v13
        v15 = torch.bmm(v14, v14)
        return v14
//  [python]
def forward(self, x1):
        v1 = self.conv(x1)
        v2 = torch.transpose(v1, 1, 1)
        v3 = torch.inverse(v2)
        v4 = torch.linalg.solve_triangular(v3.contiguous(), torch.cholesky(v1.contiguous()))
        v5 = v4
        v6 = torch.cholesky(v1.contiguous())
        v7 = torch.transpose(v1, 1, 1)
        v8 = torch.inverse(v2)
        v9 = torch.linalg.solve_triangular(v8.contiguous(), v6)
        v10 = v9
        v11 = v5 + 1
        v12 = v10 + 1
        v13 = v7 * v11
        v14 = v12 * v13
        v15 = torch.bmm(v14, v14)
        return v14
//  [python]
def predict(self, pred_x, hyper=None, verbose=False, compute_grad = False):
        if hyper is not None:
            param_original = self.model.param_to_vec()
            self.cholesky_update(hyper)

        k_pred_train = self.model.kernel(pred_x, self.train_x)
        k_pred = self.model.kernel(pred_x, diagonal=True)

        # cholesky is lower triangular matrix
        chol_solver = torch.triangular_solve(torch.cat([k_pred_train.t(), self.mean_vec], 1), self.cholesky, upper=False)[0]
        chol_solve_k = chol_solver[:, :-1]
        chol_solve_y = chol_solver[:, -1:]

        pred_mean = torch.mm(chol_solve_k.t(), chol_solve_y) + self.model.mean(pred_x)
        pred_quad = (chol_solve_k ** 2).sum(0).view(-1, 1)
        pred_var = k_pred - pred_quad

        if verbose:
            numerically_stable = (pred_var >= 0).all()
            zero_pred_var = (pred_var <= 0).all()

        if hyper is not None:
            self.cholesky_update(param_original)

        if compute_grad:
            alpha = torch.cholesky_solve(self.mean_vec, self.cholesky, upper=False)
            grad_cross = self.model.kernel.grad(self.train_x, pred_x)
            grad_xp_m = torch.mm(grad_cross, k_pred_train.t()*alpha)
            gamma = torch.triangular_solve(chol_solve_k, self.cholesky.t(), upper=True)[0]
            grad_xp_v = -2 * torch.mm(gamma.t(), (grad_cross * k_pred_train).t()).t()
            return pred_mean, pred_var.clamp(min=1e-8), grad_xp_m, grad_xp_v
        else:
            if verbose:
                return pred_mean, pred_var.clamp(min=1e-8), numerically_stable, zero_pred_var
            else:
                return pred_mean, pred_var.clamp(min=1e-8)
//  [python]
def predict(self, pred_x, hyper=None, verbose=False, compute_grad = False):
        if hyper is not None:
            param_original = self.model.param_to_vec()
            self.cholesky_update(hyper)

        k_pred_train = self.model.kernel(pred_x, self.train_x)
        k_pred = self.model.kernel(pred_x, diagonal=True)

        # cholesky is lower triangular matrix
        chol_solver = torch.linalg.solve_triangular(torch.cat([k_pred_train.t(), self.mean_vec], 1), self.cholesky, upper=False)
        chol_solve_k = chol_solver[:, :-1]
        chol_solve_y = chol_solver[:, -1:]

        pred_mean = torch.mm(chol_solve_k.t(), chol_solve_y) + self.model.mean(pred_x)
        pred_quad = (chol_solve_k ** 2).sum(0).view(-1, 1)
        pred_var = k_pred - pred_quad

        if verbose:
            numerically_stable = (pred_var >= 0).all()
            zero_pred_var = (pred_var <= 0).all()

        if hyper is not None:
            self.cholesky_update(param_original)

        if compute_grad:
            alpha = torch.linalg.solve_triangular(self.mean_vec, self.cholesky, upper=False)
            grad_cross = self.model.kernel.grad(self.train_x, pred_x)
            grad_xp_m = torch.mm(grad_cross, k_pred_train.t()*alpha)
            gamma = torch.linalg.solve_triangular(chol_solve_k, self.cholesky.t(), upper=True)
            grad_xp_v = -2 * torch.mm(gamma.t(), (grad_cross * k_pred_train).t()).t()
            return pred_mean, pred_var.clamp(min=1e-8), grad_xp_m, grad_xp_v
        else:
            if verbose:
                return pred_mean, pred_var.clamp(min=1e-8), numerically_stable, zero_pred_var
            else:
                return pred_mean, pred_var.clamp(min=1e-8)
//  [python]
def _simple_qr_backward(q, r, dq, dr):
    if r.shape[-2] != r.shape[-1]:
        raise NotImplementedError("QrGrad not implemented when ncols > nrows "
                          "or full_matrices is true and ncols != nrows.")

    qdq = q.t() @ dq
    qdq_ = qdq - qdq.t()
    rdr = r @ dr.t()
    rdr_ = rdr - rdr.t()
    tril = torch.tril(qdq_ + rdr_)

    def _TriangularSolve(x, r):
        """Equiv to x @ torch.inverse(r).t() if r is upper-tri."""
        res = torch.triangular_solve(x.t(), r, upper=True, transpose=False)[0].t()
        return res

    grad_a = q @ (dr + _TriangularSolve(tril, r))
    grad_b = _TriangularSolve(dq - q @ qdq, r)
    return grad_a + grad_b
//  [python]
def _simple_qr_backward(q, r, dq, dr):
    if r.shape[-2] != r.shape[-1]:
        raise NotImplementedError("QrGrad not implemented when ncols > nrows "
                          "or full_matrices is true and ncols != nrows.")

    qdq = q.t() @ dq
    qdq_ = qdq - qdq.t()
    rdr = r @ dr.t()
    rdr_ = rdr - rdr.t()
    tril = torch.tril(qdq_ + rdr_)

    def _TriangularSolve(x, r):
        """Equiv to x @ torch.inverse(r).t() if r is upper-tri."""
        res = torch.linalg.solve_triangular(x.t(), r, upper=True, transpose=False).t()
        return res

    grad_a = q @ (dr + _TriangularSolve(tril, r))
    grad_b = _TriangularSolve(dq - q @ qdq, r)
    return grad_a + grad_b
//  [python]
def cho_solve(L, b):
    """Compute tensor $C^{-1} b$ from cholesky factor.

    This function supports batched inputs.

    Args:
        L (torch.Tensor): (B x )N x N lower triangular tensor where $L L^T = C$.
        b (torch.Tensor): (B x )N x L tensor.

    Returns:
        torch.Tensor: $C^{-1} b$
    """

    tmp, _ = torch.triangular_solve(b, L, upper=False)
    tmp2, _ = torch.triangular_solve(tmp, torch.transpose(L, -2, -1), upper=True)
    return tmp2
//  [python]
def cho_solve(L, b):
    """Compute tensor $C^{-1} b$ from cholesky factor.

    This function supports batched inputs.

    Args:
        L (torch.Tensor): (B x )N x N lower triangular tensor where $L L^T = C$.
        b (torch.Tensor): (B x )N x L tensor.

    Returns:
        torch.Tensor: $C^{-1} b$
    """

    tmp = torch.linalg.solve_triangular(L, b, upper=False)
    tmp2 = torch.linalg.solve_triangular(torch.transpose(L, -2, -1), tmp, upper=True)
    return tmp2
//  [python]
def w(self):
        if self.whitened_param:
            return torch.triangular_solve(self.raw_w, self.C.t()).solution
        else:
            return self.raw_w
//  [python]
def w(self):
        if self.whitened_param:
            return torch.linalg.solve_triangular(self.raw_w, self.C.t(), upper=False)
        else:
            return self.raw_w
//  [python]
def sample_full_rank(self, n_samples, mu, tril_elements, as_numpy=True):
        """Sample from a single Gaussian posterior with a full-rank covariance
        matrix

        Parameters
        ----------
        n_samples : int
            how many samples to obtain
        mu : torch.Tensor of shape `[self.batch_size, self.Y_dim]`
            network prediction of the mu (mean parameter) of the BNN posterior
        tril_elements : torch.Tensor of shape `[self.batch_size, tril_len]`
            network prediction of lower-triangular matrix in the log-Cholesky
            decomposition of the precision matrix

        Returns
        -------
        np.array of shape `[self.batch_size, n_samples, self.Y_dim]`
            samples
        """
        eps = torch.randn([self.batch_size, self.Y_dim, n_samples]).to(mu.device)  # [B, Y, N]
        tril = torch.zeros([self.batch_size, self.Y_dim, self.Y_dim]).to(mu.device)  # [B, Y, Y]
        tril[:, self.tril_idx[0], self.tril_idx[1]] = tril_elements
        log_diag_tril = torch.diagonal(tril, offset=0, dim1=1, dim2=2)
        tril[:, torch.eye(self.Y_dim, dtype=bool)] = torch.exp(log_diag_tril)
        scaled_eps, _ = torch.triangular_solve(eps, tril, transpose=True, upper=True)
        samples = mu.unsqueeze(-1) + scaled_eps  # [B, Y, N]
        samples = torch.transpose(samples, dim0=1, dim1=2)  # [B, N, Y]
        samples = torch.nan_to_num(samples, nan=1e5, posinf=1e5, neginf=-1e5)
        samples = torch.clamp(samples, min=-1.e5, max=1.e5)
        samples = self.unwhiten_back(samples)
        if as_numpy:
            return samples.cpu().numpy()
        else:
            return samples
//  [python]
def sample_full_rank(self, n_samples, mu, tril_elements, as_numpy=True):
        """Sample from a single Gaussian posterior with a full-rank covariance
        matrix

        Parameters
        ----------
        n_samples : int
            how many samples to obtain
        mu : torch.Tensor of shape `[self.batch_size, self.Y_dim]`
            network prediction of the mu (mean parameter) of the BNN posterior
        tril_elements : torch.Tensor of shape `[self.batch_size, tril_len]`
            network prediction of lower-triangular matrix in the log-Cholesky
            decomposition of the precision matrix

        Returns
        -------
        np.array of shape `[self.batch_size, n_samples, self.Y_dim]`
            samples
        """
        eps = torch.randn([self.batch_size, self.Y_dim, n_samples]).to(mu.device)  # [B, Y, N]
        tril = torch.zeros([self.batch_size, self.Y_dim, self.Y_dim]).to(mu.device)  # [B, Y, Y]
        tril[:, self.tril_idx[0], self.tril_idx[1]] = tril_elements
        log_diag_tril = torch.diagonal(tril, offset=0, dim1=1, dim2=2)
        tril[:, torch.eye(self.Y_dim, dtype=bool)] = torch.exp(log_diag_tril)
        scaled_eps = torch.linalg.solve_triangular(tril, eps, upper=False, left=True)
        samples = mu.unsqueeze(-1) + scaled_eps  # [B, Y, N]
        samples = torch.transpose(samples, dim0=1, dim1=2)  # [B, N, Y]
        samples = torch.nan_to_num(samples, nan=1e5, posinf=1e5, neginf=-1e5)
        samples = torch.clamp(samples, min=-1.e5, max=1.e5)
        samples = self.unwhiten_back(samples)
        if as_numpy:
            return samples.cpu().numpy()
        else:
            return samples
//  [python]
def cho_solve_AXB(a, L, b):
    """Compute tensor $a C^{-1} b$ from cholesky factor.

    Args:
        a (torch.Tensor): M x N tensor.
        L (torch.Tensor): N x N lower triangular tensor where $L L^T = C$.
        b (torch.Tensor): N x L tensor.

    Returns:
        torch.Tensor: $a C^{-1} b$
    """

    left, _ = torch.triangular_solve(a.t(), L, upper=False)
    right, _ = torch.triangular_solve(b, L, upper=False)
    return torch.mm(left.t(), right)
//  [python]
def cho_solve_AXB(a, L, b):
    """Compute tensor $a C^{-1} b$ from cholesky factor.

    Args:
        a (torch.Tensor): M x N tensor.
        L (torch.Tensor): N x N lower triangular tensor where $L L^T = C$.
        b (torch.Tensor): N x L tensor.

    Returns:
        torch.Tensor: $a C^{-1} b$
    """

    left = torch.linalg.solve_triangular(L, a.t(), upper=False)
    right = torch.linalg.solve_triangular(L, b, upper=False)
    return torch.mm(left.t(), right)
//  [python]
def reset_parameters(self) -> None:
        if self.whitened_param:
            nn.init.normal_(self.raw_w, mean=0.0, std=1.0)
        else:
            w_init = torch.triangular_solve(self.raw_w, self.C.t()).solution
            self.raw_w.data.copy_(w_init)
        nn.init.zeros_(self.b)
//  [python]
def reset_parameters(self) -> None:
        if self.whitened_param:
            nn.init.normal_(self.raw_w, mean=0.0, std=1.0)
        else:
            w_init = torch.linalg.solve_triangular(self.raw_w, self.C.t())
            self.raw_w.data.copy_(w_init)
        nn.init.zeros_(self.b)
//  [python]
def cholesky_forward(self, x):
        """normalized tensor"""
        batch_size, feature_dim = x.size()
        f_cov = torch.mm(x.transpose(0, 1), x) / (batch_size - 1) # d * d
        eye = torch.eye(feature_dim).float().to(f_cov.device)
        f_cov_shrink = (1 - self.eps) * f_cov + self.eps * eye
        inv_sqrt = torch.triangular_solve(eye, torch.linalg.cholesky(f_cov_shrink.float()), upper=False)[0]
        inv_sqrt = inv_sqrt.contiguous().view(feature_dim, feature_dim).detach()
        return torch.mm(inv_sqrt, x.transpose(0, 1)).transpose(0, 1)
//  [python]
def cholesky_forward(self, x):
        """normalized tensor"""
        batch_size, feature_dim = x.size()
        f_cov = torch.mm(x.transpose(0, 1), x) / (batch_size - 1) # d * d
        eye = torch.eye(feature_dim).float().to(f_cov.device)
        f_cov_shrink = (1 - self.eps) * f_cov + self.eps * eye
        inv_sqrt = torch.linalg.solve_triangular(torch.linalg.cholesky(f_cov_shrink.float()), eye, upper=False)
        inv_sqrt = inv_sqrt.contiguous().view(feature_dim, feature_dim).detach()
        return torch.mm(inv_sqrt, x.transpose(0, 1)).transpose(0, 1)
//  [python]
def update_precond_dense(Q, dxs, dgs, step=0.01):
    """
    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: normalized step size in [0, 1]
    """
    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.t(), upper=False)[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):
    """
    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: normalized step size in [0, 1]
    """
    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.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 test_triangular_solve_out(self):
        a = torch.randn(9, 252, 59, 19).npu()
        b = torch.randn(59, 59).npu()
        c = torch.randn(9, 252, 59, 19).npu()
        d = torch.randn(59, 59).npu()
        out = [c, d]
        output = torch.triangular_solve(a, b.tril().npu(), upper=False, transpose=True, unitriangular=False, out=out)
//  [python]
def test_triangular_solve_out(self):
        a = torch.randn(9, 252, 59, 19).npu()
        b = torch.randn(59, 59).npu()
        c = torch.randn(9, 252, 59, 19).npu()
        d = torch.randn(59, 59).npu()
        out = [c, d]
        output = torch.linalg.solve_triangular(a, b.tril().npu(), upper=False, transpose=True, unitriangular=False, out=out)
//  [python]
def tr_solve(x,T,freeGPU,transpose = False):
    m = T.size(0)
    useGPU = not(isinstance(freeGPU,type(None)))
    if useGPU:
        if m*(2*m+2) < 0.98*freeGPU:
            download,upload = bp.produceDU(useGPU)
            xg = upload(x)
            Tg = upload(T)
            resg = torch.triangular_solve(xg,Tg,upper \
                                                   = True,transpose = transpose)
            res = download(resg[0])
            del xg
            del Tg
            del resg
            torch.cuda.empty_cache()
            return res
        else:
            return torch.triangular_solve(x,T,upper \
                                                   = True,transpose = transpose)[0]
    else:
        return torch.triangular_solve(x,T,upper \
                                                   = True,transpose = transpose)[0]
//  [python]
def tr_solve(x,T,freeGPU,transpose = False):
    m = T.size(0)
    useGPU = not(isinstance(freeGPU,type(None)))
    if useGPU:
        if m*(2*m+2) < 0.98*freeGPU:
            download,upload = bp.produceDU(useGPU)
            xg = upload(x)
            Tg = upload(T)
            resg = torch.linalg.solve_triangular(xg,Tg,upper \
                                                   = True,transpose = transpose)
            res = download(resg)
            del xg
            del Tg
            del resg
            torch.cuda.empty_cache()
            return res
        else:
            return torch.linalg.solve_triangular(x,T,upper \
                                                   = True,transpose = transpose)
    else:
        return torch.linalg.solve_triangular(x,T,upper \
                                                   = True,transpose = transpose)
//  [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).to(lower.device)
        lower_inverse, _ = torch.triangular_solve(identity, lower, upper=False, unitriangular=True)
        weight_inverse, _ = torch.triangular_solve(lower_inverse, upper, upper=True, unitriangular=False)
        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).to(lower.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 test_triangular_solve(self, device, dtype):
        ks = [0, 1, 3]
        ns = [0, 5]
        for k, n, (upper, unitriangular, transpose) in itertools.product(ks, ns,
                                                                         itertools.product([True, False], repeat=3)):
            b, A = self.triangular_solve_test_helper((n, n), (n, k), upper,
                                                     unitriangular, device, dtype)
            with pytorch_op_timer():
                x = torch.triangular_solve(
                    b, A, upper=upper, unitriangular=unitriangular, transpose=transpose)[0]
            if transpose:
                self.assertEqual(b, np.matmul(A.t().cpu(), x.cpu()))
            else:
                self.assertEqual(b, np.matmul(A.cpu(), x.cpu()))
//  [python]
def test_triangular_solve(self, device, dtype):
        ks = [0, 1, 3]
        ns = [0, 5]
        for k, n, (upper, unitriangular, transpose) in itertools.product(ks, ns,
                                                                         itertools.product([True, False], repeat=3)):
            b, A = self.triangular_solve_test_helper((n, n), (n, k), upper,
                                                     unitriangular, device, dtype)
            with pytorch_op_timer():
                x = torch.linalg.solve_triangular(
                    b, A, upper=upper, unitriangular=unitriangular, left=not transpose)
            if transpose:
                self.assertEqual(b, np.matmul(A.t().cpu(), x.cpu()))
            else:
                self.assertEqual(b, np.matmul(A.cpu(), x.cpu()))
//  [python]
def update_precond_kron(Ql, Qr, dX, dG, step=0.01):
    """
    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: normalized step size in range [0, 1] 
    """
    max_l = torch.max(torch.abs(Ql))
    max_r = torch.max(torch.abs(Qr))
    
    rho = torch.sqrt(max_l/max_r)
    Ql = Ql/rho
    Qr = rho*Qr
    
    A = Ql.mm( dG.mm( Qr.t() ) )
    Bt = torch.triangular_solve((torch.triangular_solve(dX.t(), Qr.t(), upper=False))[0].t(), 
                     Ql.t(), upper=False)[0]
    
    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)
        
    return Ql - step1*grad1.mm(Ql), Qr - step2*grad2.mm(Qr)
//  [python]
def update_precond_kron(Ql, Qr, dX, dG, step=0.01):
    """
    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: normalized step size in range [0, 1] 
    """
    max_l = torch.max(torch.abs(Ql))
    max_r = torch.max(torch.abs(Qr))
    
    rho = torch.sqrt(max_l/max_r)
    Ql = Ql/rho
    Qr = rho*Qr
    
    A = Ql.mm( dG.mm( Qr.t() ) )
    Bt = torch.linalg.solve_triangular((torch.linalg.solve_triangular(dX.t(), Qr.t(), upper=False)).t(), 
                     Ql.t(), upper=False)
    
    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)
        
    return Ql - step1*grad1.mm(Ql), Qr - step2*grad2.mm(Qr)
//  [python]
def update_precond_scaw(Ql, qr, dX, dG, step=0.01):
    """
    update scaling-and-whitening preconditioner
    """
    max_l = torch.max(torch.abs(Ql))
    max_r = torch.max(torch.abs(qr))
    
    rho = torch.sqrt(max_l/max_r)
    Ql = Ql/rho
    qr = rho*qr
    
    A = Ql.mm( dG*qr )
    Bt = torch.triangular_solve(dX/qr, Ql.t(), upper=False)[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_scaw(Ql, qr, dX, dG, step=0.01):
    """
    update scaling-and-whitening preconditioner
    """
    max_l = torch.max(torch.abs(Ql))
    max_r = torch.max(torch.abs(qr))
    
    rho = torch.sqrt(max_l/max_r)
    Ql = Ql/rho
    qr = rho*qr
    
    A = Ql.mm( dG*qr )
    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 lstsq(b, A):
    m, n = A.shape[-2:]
    if m < n:
        # solve least-norm problem
        Q, R = qr(A.transpose(-1,-2))
        d = torch.triangular_solve(b, R.transpose(-1,-2), upper=False)[0]
        x = torch.matmul(Q, d)
    else:
        # solve least-squares problem
        Q, R = qr(A)
        d = torch.matmul(Q.transpose(-1,-2), b)
        x = torch.triangular_solve(d, R)[0]
    return x
//  [python]
def lstsq(b, A):
    m, n = A.shape[-2:]
    if m < n:
        # solve least-norm problem
        Q, R = qr(A.transpose(-1,-2))
        d = torch.linalg.solve_triangular(R.transpose(-1,-2), b, upper=False)
        x = torch.matmul(Q, d)
    else:
        # solve least-squares problem
        Q, R = qr(A)
        d = torch.matmul(Q.transpose(-1,-2), b)
        x = torch.linalg.solve_triangular(R, d)
    return x
//  [python]
def _inverse(self, y):
        return torch.triangular_solve((y - self.loc)[..., :, None],
                                      self.L,
                                      upper=False,
                                      transpose=True)[0].squeeze(-1)
//  [python]
def _inverse(self, y):
        return torch.linalg.solve_triangular((y - self.loc)[..., :, None],
                                      self.L,
                                      upper=False,
                                      left=True)[0].squeeze(-1)
//  [python]
def get_white_from_latents(latents, const_mean, cholesky_gram):
    """
    latents: (B,D)
    const_mean: scalar
    cholesky_gram: (D,D)
    """
    white = th.triangular_solve(
        latents.unsqueeze(-1) - const_mean, cholesky_gram, upper=False
    ).solution  # (B,D, 1)
    return white.squeeze(dim=-1)
//  [python]
def get_white_from_latents(latents, const_mean, cholesky_gram):
    """
    latents: (B,D)
    const_mean: scalar
    cholesky_gram: (D,D)
    """
    white = torch.linalg.solve_triangular(
        latents.unsqueeze(-1) - const_mean, cholesky_gram, upper=False
    )  # (B,D, 1)
    return white.squeeze(dim=-1)
//  [python]
def _default_invchol(self, x, ret_chol=False):
        l = self.chol(x)
        eye = torch.eye(self.n, out=x.new(self.n, self.n))
        l_inv = torch.triangular_solve(eye, l, upper=False).solution
        if not ret_chol:
            return l_inv, None
        return l_inv, l
//  [python]
def _default_invchol(self, x, ret_chol=False):
        l = self.chol(x)
        eye = torch.eye(self.n, out=x.new(self.n, self.n))
        l_inv = torch.linalg.solve_triangular(eye, l, upper=False)
        if not ret_chol:
            return l_inv, None
        return l_inv, l
//  [python]
def retr(self, x, u):
        # Computing :math:`X + U + \frac{1}{2} U X^{-1} U` as:
        # :math:`U X^{-1} U = U (L L^T)^{-1} U = (L^{-1} U)^T (L^{-1} U)`.
        l = self.chol(x)
        l_inv_u = torch.triangular_solve(u, l, upper=False).solution
        u_xinv_u = tb.xty(l_inv_u, l_inv_u)
        y = tb.sym(x + u + 0.5 * u_xinv_u)

        return y
//  [python]
def retr(self, x, u):
        # Computing :math:`X + U + \frac{1}{2} U X^{-1} U` as:
        # :math:`U X^{-1} U = U (L L^T)^{-1} U = (L^{-1} U)^T (L^{-1} U)`.
        l = self.chol(x)
        l_inv_u = torch.linalg.solve_triangular(l, u, upper=False)
        u_xinv_u = tb.xty(l_inv_u, l_inv_u)
        y = tb.sym(x + u + 0.5 * u_xinv_u)

        return y
//  [python]
def solve(self, nsteps=None, source_func=None):
        """Integrates the model for a given number of steps starting from the
        initial conditions.

        Parameters
        ----------
        nsteps : int, optional
            number of time steps to take in the integration, by default None
        source_func : function, optional
            The source term as an explicit function of the unknown only, by default None

        Returns
        -------
        tensor
            Tensor of shape (n_steps, m), where m is the number of levels, whose
            entries are the zonal wind profiles at each time.
        """

        if nsteps is None:
            nsteps, = self.time.shape

        if source_func is None:
            # source_func, _, _ = utils.make_source_func(self)
            source_func = utils.load_model(self)

        u = torch.zeros((nsteps, self.nlev))
        u[0] = self.initial_condition(self.z)

        # a single forward Euler step
        source = source_func(u[0])
        u[1] = (torch.matmul(torch.eye(self.nlev) - self.D, u[0]) -
        self.dt * source)

        for n in range(1, nsteps - 1):

            self.current_time += self.dt
            source = source_func(u[n])

            # RHS multiplied by QT on the left
            b = torch.matmul(self.QT, (
                torch.matmul(torch.eye(self.nlev) - self.D, u[n - 1]) -
                2 * self.dt * source
            )).reshape(-1, 1)

            u[n + 1] = torch.triangular_solve(b, self.R).solution.flatten()

        return u
//  [python]
def solve(self, nsteps=None, source_func=None):
        """Integrates the model for a given number of steps starting from the
        initial conditions.

        Parameters
        ----------
        nsteps : int, optional
            number of time steps to take in the integration, by default None
        source_func : function, optional
            The source term as an explicit function of the unknown only, by default None

        Returns
        -------
        tensor
            Tensor of shape (n_steps, m), where m is the number of levels, whose
            entries are the zonal wind profiles at each time.
        """

        if nsteps is None:
            nsteps, = self.time.shape

        if source_func is None:
            # source_func, _, _ = utils.make_source_func(self)
            source_func = utils.load_model(self)

        u = torch.zeros((nsteps, self.nlev))
        u[0] = self.initial_condition(self.z)

        # a single forward Euler step
        source = source_func(u[0])
        u[1] = (torch.matmul(torch.eye(self.nlev) - self.D, u[0]) -
        self.dt * source)

        for n in range(1, nsteps - 1):

            self.current_time += self.dt
            source = source_func(u[n])

            # RHS multiplied by QT on the left
            b = torch.matmul(self.QT, (
                torch.matmul(torch.eye(self.nlev) - self.D, u[n - 1]) -
                2 * self.dt * source
            )).reshape(-1, 1)

            u[n + 1] = torch.linalg.solve_triangular(self.R, b, upper=True).flatten()

        return u
//  [python]
def solve_triangular(A, b, **kwargs):
    return torch.triangular_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).squeeze(1)
//  [python]
def forward(ctx, A, B, upper, unitriangular, transpose):
        ctx.batch_size = B.size()[0] if B.dim() == 3 else None
        ctx.A_shape = A.size()  # (b), m, m
        ctx.B_shape = B.size()  # (b), m, p
        ctx.csr = True
        ctx.upper = upper
        ctx.unitriangular = unitriangular
        ctx.transpose = transpose

        grad_flag = A.requires_grad or B.requires_grad

        if ctx.batch_size is not None:
            A = sparse_block_diag(*A)
            B = torch.cat([*B])

        if A.layout == torch.sparse_coo:
            A = convert_coo_to_csr(A)  # triangular solve doesn't work with sparse coo
            ctx.csr = False

        # Check if a workaround for https://github.com/pytorch/pytorch/issues/88890 is needed
        workaround88890 = (
            A.device == torch.device("cpu")
            and (not ctx.upper)
            and ctx.unitriangular
            and (torch.__version__ < (2,))
            and ctx.transpose
        )
        if not workaround88890:
            x = torch.triangular_solve(
                B.detach(), A.detach(), upper=upper, unitriangular=unitriangular, transpose=transpose
            ).solution
        else:
            n = A.shape[0]
            id_csr = torch.sparse_csr_tensor(
                torch.arange(n + 1),
                torch.arange(n),
                torch.ones(n, device=A.device, dtype=A.dtype),
                (n, n),
                device=A.device,
            )
            x = torch.triangular_solve(B.detach(), A.detach() + id_csr, upper=ctx.upper, transpose=transpose).solution

        x.requires_grad = grad_flag
        ctx.save_for_backward(A, x.detach())

        if ctx.batch_size is not None:
            x = x.view(ctx.batch_size, ctx.A_shape[-2], ctx.B_shape[-1])

        return x
//  [python]
def forward(ctx, A, B, upper, unitriangular, transpose):
        ctx.batch_size = B.size()[0] if B.dim() == 3 else None
        ctx.A_shape = A.size()  # (b), m, m
        ctx.B_shape = B.size()  # (b), m, p
        ctx.csr = True
        ctx.upper = upper
        ctx.unitriangular = unitriangular
        ctx.transpose = transpose

        grad_flag = A.requires_grad or B.requires_grad

        if ctx.batch_size is not None:
            A = sparse_block_diag(*A)
            B = torch.cat([*B])

        if A.layout == torch.sparse_coo:
            A = convert_coo_to_csr(A)  # triangular solve doesn't work with sparse coo
            ctx.csr = False

        # Check if a workaround for https://github.com/pytorch/pytorch/issues/88890 is needed
        workaround88890 = (
            A.device == torch.device("cpu")
            and (not ctx.upper)
            and ctx.unitriangular
            and (torch.__version__ < (2,))
            and ctx.transpose
        )
        if not workaround88890:
            x = torch.linalg.solve_triangular(
                B.detach(), A.detach(), upper=upper, unitriangular=unitriangular, transpose=transpose
            )
        else:
            n = A.shape[0]
            id_csr = torch.sparse_csr_tensor(
                torch.arange(n + 1),
                torch.arange(n),
                torch.ones(n, device=A.device, dtype=A.dtype),
                (n, n),
                device=A.device,
            )
            x = torch.linalg.solve_triangular(B.detach(), A.detach() + id_csr, upper=ctx.upper, transpose=transpose)

        x.requires_grad = grad_flag
        ctx.save_for_backward(A, x.detach())

        if ctx.batch_size is not None:
            x = x.view(ctx.batch_size, ctx.A_shape[-2], ctx.B_shape[-1])

        return x
//  [python]
def cholesky_kl(p0: CholeskyModule, p1: CholeskyModule, decompose=False, nograd_logdet=False):
    z0, L0 = p0.joint_parameters()
    z1, L1 = p1.joint_parameters()
    V, _ = torch.triangular_solve(L1, L0, upper=False)
    diff = V @ z0 - z1
    vecnorm = .5 * torch.sum(diff ** 2)
    matnorm = .5 * (torch.sum(V ** 2) - z0.shape[0])
    logdet = - torch.sum(torch.log(torch.diag(V)))
    if nograd_logdet:
        logdet.detach_()

    matdivergence = matnorm + logdet
    total = vecnorm + matdivergence
    if decompose:
        return {'vector': vecnorm.item(), 'matrix': matdivergence.item(),
                'total': total.item(), 'v/t': vecnorm.item() / total.item(),
                'm/t': matdivergence.item() / total.item()}
    else:
        return total
//  [python]
def cholesky_kl(p0: CholeskyModule, p1: CholeskyModule, decompose=False, nograd_logdet=False):
    z0, L0 = p0.joint_parameters()
    z1, L1 = p1.joint_parameters()
    V = torch.linalg.solve_triangular(L0, L1, upper=False)
    diff = V @ z0 - z1
    vecnorm = .5 * torch.sum(diff ** 2)
    matnorm = .5 * (torch.sum(V ** 2) - z0.shape[0])
    logdet = - torch.sum(torch.log(torch.diag(V)))
    if nograd_logdet:
        logdet.detach_()

    matdivergence = matnorm + logdet
    total = vecnorm + matdivergence
    if decompose:
        return {'vector': vecnorm.item(), 'matrix': matdivergence.item(),
                'total': total.item(), 'v/t': vecnorm.item() / total.item(),
                'm/t': matdivergence.item() / total.item()}
    else:
        return total
//  [python]
def inv(self, x):
        return t.triangular_solve(x, self.A, upper=False, transpose=False)[0]
//  [python]
def inv(self, x):
        return torch.linalg.solve_triangular(self.A, x, upper=False, left=True)[0]
//  [python]
def _kl_divergence_state(inv_state1, inv_state2, temperature_scaling: Union[float, Tensor] = 1.) -> Tensor:
        temperature_scaling = torch.as_tensor(temperature_scaling).to(inv_state1)
        inv_state2 = inv_state2.to(inv_state1)
        half_term1 = (inv_state2.diagonal(dim1=-2, dim2=-1).log().sum(-1) -
                      inv_state1.diagonal(dim1=-2, dim2=-1).log().sum(-1))
        mm = torch.triangular_solve(inv_state1, inv_state2, upper=False)[0]
        term2 = temperature_scaling * mm.pow(2).sum()
        return half_term1 + .5 * (term2 - (temperature_scaling.log() + 1) * inv_state1.shape[-1])
//  [python]
def _kl_divergence_state(inv_state1, inv_state2, temperature_scaling: Union[float, Tensor] = 1.) -> Tensor:
        temperature_scaling = torch.as_tensor(temperature_scaling).to(inv_state1)
        inv_state2 = inv_state2.to(inv_state1)
        half_term1 = (inv_state2.diagonal(dim1=-2, dim2=-1).log().sum(-1) -
                      inv_state1.diagonal(dim1=-2, dim2=-1).log().sum(-1))
        mm = torch.linalg.solve_triangular(inv_state1, inv_state2, upper=False)
        term2 = temperature_scaling * mm.pow(2).sum()
        return half_term1 + .5 * (term2 - (temperature_scaling.log() + 1) * inv_state1.shape[-1])
//  [python]
def test_triangular_solve_singular(self, device, dtype):
        b = torch.rand(3, 1, device=device)
        A = torch.eye(3, 3, device=device)
        A[-1, -1] = 0  # Now A is singular
        err_str = r"triangular_solve_cpu: U\(3,3\) is zero, singular U\."
        with self.assertRaisesRegex(RuntimeError, err_str):
            torch.triangular_solve(b, A)
//  [python]
def test_triangular_solve_singular(self, device, dtype):
        b = torch.rand(3, 1, device=device)
        A = torch.eye(3, 3, device=device)
        A[-1, -1] = 0  # Now A is singular
        err_str = r"triangular_solve_cpu: U\(3,3\) is zero, singular U\."
        with self.assertRaisesRegex(RuntimeError, err_str):
            torch.linalg.solve_triangular(b, A)
//  [python]
def trtrs(b, A, upper=True, transpose=False, unitriangular=False, out=None):
    r"""Solves a system of equations with a triangular coefficient matrix :math:`A`
    and multiple right-hand sides :attr:`b`.

    In particular, solves :math:`AX = b` and assumes :math:`A` is upper-triangular
    with the default keyword arguments.

    For more information regarding :func:`torch.trtrs`, please check :func:`torch.triangular_solve`.

    .. warning::
        :func:`torch.trtrs` is deprecated in favour of :func:`torch.triangular_solve` and will be
        removed in the next release. Please use :func:`torch.triangular_solve` instead.
    """
    warnings.warn("torch.trtrs is deprecated in favour of torch.triangular_solve and will be "
                  "removed in the next release. Please use torch.triangular_solve instead.", stacklevel=2)
    return torch.triangular_solve(b, A, upper=upper, transpose=transpose, unitriangular=unitriangular, out=out)
//  [python]
def trtrs(b, A, upper=True, transpose=False, unitriangular=False, out=None):
    r"""Solves a system of equations with a triangular coefficient matrix :math:`A`
    and multiple right-hand sides :attr:`b`.

    In particular, solves :math:`AX = b` and assumes :math:`A` is upper-triangular
    with the default keyword arguments.

    For more information regarding :func:`torch.trtrs`, please check :func:`torch.triangular_solve`.

    .. warning::
        :func:`torch.trtrs` is deprecated in favour of :func:`torch.triangular_solve` and will be
        removed in the next release. Please use :func:`torch.triangular_solve` instead.
    """
    warnings.warn("torch.trtrs is deprecated in favour of torch.triangular_solve and will be "
                  "removed in the next release. Please use torch.triangular_solve instead.", stacklevel=2)
    return torch.linalg.solve_triangular(b, A, upper=upper, transpose=transpose, unitriangular=unitriangular, out=out)
//  [python]
def predict(self, pred_x, hyper=None, verbose=False):
		if hyper is not None:
			param_original = self.model.param_to_vec()
			self.cholesky_update(hyper)

		k_pred_train = self.model.kernel(pred_x, self.train_x)
		k_pred = self.model.kernel(pred_x, diagonal=True)

		# cholesky is lower triangular matrix
		chol_solver = torch.triangular_solve(torch.cat([k_pred_train.t(), self.mean_vec], 1), self.cholesky, upper=False)[0]
		chol_solve_k = chol_solver[:, :-1]
		chol_solve_y = chol_solver[:, -1:]

		pred_mean = torch.mm(chol_solve_k.t(), chol_solve_y) + self.model.mean(pred_x)
		pred_quad = (chol_solve_k ** 2).sum(0).view(-1, 1)
		pred_var = k_pred - pred_quad

		if verbose:
			numerically_stable = (pred_var >= 0).all()
			zero_pred_var = (pred_var <= 0).all()

		if hyper is not None:
			self.cholesky_update(param_original)

		if verbose:
			return pred_mean, pred_var.clamp(min=1e-8), numerically_stable, zero_pred_var
		else:
			return pred_mean, pred_var.clamp(min=1e-8)
//  [python]
def predict(self, pred_x, hyper=None, verbose=False):
		if hyper is not None:
			param_original = self.model.param_to_vec()
			self.cholesky_update(hyper)

		k_pred_train = self.model.kernel(pred_x, self.train_x)
		k_pred = self.model.kernel(pred_x, diagonal=True)

		# cholesky is lower triangular matrix
		chol_solver = torch.linalg.solve_triangular(torch.cat([k_pred_train.t(), self.mean_vec], 1), self.cholesky, upper=False)
		chol_solve_k = chol_solver[:, :-1]
		chol_solve_y = chol_solver[:, -1:]

		pred_mean = torch.mm(chol_solve_k.t(), chol_solve_y) + self.model.mean(pred_x)
		pred_quad = (chol_solve_k ** 2).sum(0).view(-1, 1)
		pred_var = k_pred - pred_quad

		if verbose:
			numerically_stable = (pred_var >= 0).all()
			zero_pred_var = (pred_var <= 0).all()

		if hyper is not None:
			self.cholesky_update(param_original)

		if verbose:
			return pred_mean, pred_var.clamp(min=1e-8), numerically_stable, zero_pred_var
		else:
			return pred_mean, pred_var.clamp(min=1e-8)