jiahuimbzuai/CoQuIR huggingface dataset:
//  [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.linalg.solve_triangular(
                torch.linalg.cholesky(Sigma_N), P[0], upper=False
            )
            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
//  [python]
def _calc_intermediate(self, X, Y):
        variance = self.penalty * X.shape[0]

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

        L = jittering_cholesky(kmm)
        A = torch.triangular_solve(kmn, L, upper=False)[0]
        AAT = A @ A.T  # m*n @ n*m = m*m in O(n * m^2), equivalent to kmn @ knm.
        # B = A @ A.T + I
        B = AAT / variance + torch.eye(AAT.shape[0], device=X.device, dtype=X.dtype)
        LB = jittering_cholesky(B)  # LB @ LB.T = B

        AY = A @ Y  # m*1
        c = torch.triangular_solve(AY, LB, upper=False)[0]  # m * p

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

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

        L = jittering_cholesky(kmm)
        A = torch.linalg.solve_triangular(L, kmn, upper=False)
        AAT = A @ A.T  # m*n @ n*m = m*m in O(n * m^2), equivalent to kmn @ knm.
        # B = A @ A.T + I
        B = AAT / variance + torch.eye(AAT.shape[0], device=X.device, dtype=X.dtype)
        LB = jittering_cholesky(B)  # LB @ LB.T = B

        AY = A @ Y  # m*1
        c = torch.linalg.solve_triangular(LB, AY, upper=False)  # m * p

        return L, A, AAT, LB, 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, LB, C, c = self._calc_intermediate(self.x_train, self.y_train)
            # Predictions are handled directly.
            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, LB, C, c = self._calc_intermediate(self.x_train, self.y_train)
            # Predictions are handled directly.
            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 _calc_intermediate(self, X, Y):
        # - 1 due to LOOCV effect.
        variance = self.penalty * (X.shape[0] - 1)
        sqrt_var = torch.sqrt(variance)

        kmn = self.kernel(self.centers, X)
        kmm = self.kernel(self.centers, self.centers)
        L = jittering_cholesky(kmm)  # L @ L.T = kmm
        A = torch.triangular_solve(kmn, L, upper=False).solution / sqrt_var  # m, n
        AAT = A @ A.T  # m, m
        B = AAT + torch.eye(AAT.shape[0], device=X.device, dtype=X.dtype)  # m, m
        LB = jittering_cholesky(B)  # LB @ LB.T = B  m, m

        # C = LB^{-1} A
        C = torch.triangular_solve(A, LB, upper=False).solution  # m, n

        c = C @ Y / sqrt_var
        return L, A, LB, C, c
//  [python]
def _calc_intermediate(self, X, Y):
        # - 1 due to LOOCV effect.
        variance = self.penalty * (X.shape[0] - 1)
        sqrt_var = torch.sqrt(variance)

        kmn = self.kernel(self.centers, X)
        kmm = self.kernel(self.centers, self.centers)
        L = jittering_cholesky(kmm)  # L @ L.T = kmm
        A = torch.linalg.solve_triangular(L, kmn, upper=False) / sqrt_var  # m, n
        AAT = A @ A.T  # m, m
        B = AAT + torch.eye(AAT.shape[0], device=X.device, dtype=X.dtype)  # m, m
        LB = jittering_cholesky(B)  # LB @ LB.T = B  m, m

        # C = LB^{-1} A
        C = torch.linalg.solve_triangular(LB, A, upper=False)  # m, n

        c = C @ Y / sqrt_var
        return L, A, LB, C, c
//  [python]
def test_solve_triangular(self):
        linear_op = self.create_linear_op()
        is_upper = linear_op.upper
        rhs = torch.randn(linear_op.size(-1), 1)
        res = torch.triangular_solve(rhs, linear_op, upper=is_upper)[0]
        res_actual = torch.linalg.solve_triangular(linear_op.to_dense(), rhs, upper=is_upper)
        self.assertAllClose(res, res_actual)
        with self.assertRaisesRegex(RuntimeError, "Incompatible argument"):
            torch.triangular_solve(rhs, linear_op, upper=not is_upper)[0]
//  [python]
def test_solve_triangular(self):
        linear_op = self.create_linear_op()
        is_upper = linear_op.upper
        rhs = torch.randn(linear_op.size(-1), 1)
        res = torch.linalg.solve_triangular(linear_op, rhs, upper=is_upper)
        res_actual = torch.linalg.solve_triangular(linear_op.to_dense(), rhs, upper=is_upper)
        self.assertAllClose(res, res_actual)
        with self.assertRaisesRegex(RuntimeError, "Incompatible argument"):
            torch.linalg.solve_triangular(linear_op, rhs, upper=not is_upper)
//  [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)

        L, A, LB, c = self._calc_intermediate(X, Y)
        C = torch.solve(A / sqrt_var, LB)[0]  # m*n

        ndeff = C.square().sum()
        datafit = torch.square(Y).sum() - torch.square(c * sqrt_var).sum()
        self._save_losses(ndeff, datafit)

        return ndeff + datafit
//  [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)

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

        ndeff = C.square().sum()
        datafit = torch.square(Y).sum() - torch.square(c * sqrt_var).sum()
        self._save_losses(ndeff, datafit)

        return ndeff + datafit
//  [python]
def solve_triangular(mat, rhs, upper=False):
    return torch.triangular_solve(rhs, mat, upper=upper)[0]
//  [python]
def solve_triangular(mat, rhs, upper=False):
    return torch.linalg.solve_triangular(mat, rhs, upper=upper)
//  [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, LB, c = self._calc_intermediate(self.x_train, self.y_train)
            tmp1 = torch.solve(c, torch.triu(LB.T))[0]
            tmp2 = torch.solve(tmp1, torch.triu(L.T))[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, 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 get_chol_features(K_mn, L):
    nystrom = torch.triangular_solve(K_mn, L, upper=False).solution
    return nystrom
//  [python]
def get_chol_features(K_mn, L):
    nystrom = torch.linalg.solve_triangular(L, K_mn, upper=False)
    return nystrom
//  [python]
def stable_pinverse(A: Tensor) -> Tensor:
    """Compute a pseudoinverse of a matrix. Employs a stabilized QR decomposition."""
    if A.shape[-2] >= A.shape[-1]:
        # skinny (or square) matrix
        Q, R = stable_qr(A)
        return torch.triangular_solve(Q.mT, R, upper=True)[0]
    else:
        # fat matrix
        Q, R = stable_qr(A.mT)
        return torch.triangular_solve(Q.mT, R, upper=True)[0].mT
//  [python]
def stable_pinverse(A: Tensor) -> Tensor:
    """Compute a pseudoinverse of a matrix. Employs a stabilized QR decomposition."""
    if A.shape[-2] >= A.shape[-1]:
        # skinny (or square) matrix
        Q, R = stable_qr(A)
        return torch.linalg.solve_triangular(R, Q.mT, upper=True)
    else:
        # fat matrix
        Q, R = stable_qr(A.mT)
        return torch.linalg.solve_triangular(R, Q.mT, upper=True).mT
//  [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)  # L @ L.T = kmm
        # A = L^{-1} K_mn / (sqrt(n*pen))
        A = torch.triangular_solve(kmn, L, upper=False)[0]
        AAT = A @ A.T  # m*n @ n*m = m*m in O(n * m^2), equivalent to kmn @ knm.
        # B = A @ A.T + I
        B = AAT / variance + torch.eye(AAT.shape[0], device=X.device, dtype=X.dtype)
        LB = jittering_cholesky(B)  # LB @ LB.T = B
        AY = A @ Y / sqrt_var  # m*1
        c = torch.triangular_solve(AY, LB, upper=False)[0] / sqrt_var  # m*1

        return L, A, 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)  # L @ L.T = kmm
        # A = L^{-1} K_mn / (sqrt(n*pen))
        A = torch.linalg.solve_triangular(L, kmn, upper=False)
        AAT = A @ A.T  # m*n @ n*m = m*m in O(n * m^2), equivalent to kmn @ knm.
        # B = A @ A.T + I
        B = AAT / variance + torch.eye(AAT.shape[0], device=X.device, dtype=X.dtype)
        LB = jittering_cholesky(B)  # LB @ LB.T = B
        AY = A @ Y / sqrt_var  # m*1
        c = torch.linalg.solve_triangular(LB, AY, upper=False) / sqrt_var  # m*1

        return L, A, LB, c
//  [python]
def inverse(self, y: torch.Tensor, h: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        lower, upper, upper_log_diag = self.extract_matrices(h)

        # Flatten inputs
        y_flat = flatten_event(y, self.event_shape)[..., None]  # (*batch_shape, event_size)
        z_flat = torch.triangular_solve(y_flat, lower, upper=False, unitriangular=True).solution  # y = Lz => z = L^{-1}y
        x_flat = torch.triangular_solve(z_flat, upper, upper=True, unitriangular=False).solution  # z = Ux => x = U^{-1}z
        x_flat = x_flat.squeeze(-1)

        output = unflatten_event(x_flat, self.event_shape)
        return output, -self.log_determinant(upper_log_diag)
//  [python]
def inverse(self, y: torch.Tensor, h: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        lower, upper, upper_log_diag = self.extract_matrices(h)

        # Flatten inputs
        y_flat = flatten_event(y, self.event_shape)[..., None]  # (*batch_shape, event_size)
        z_flat = torch.linalg.solve_triangular(lower, y_flat, upper=False, unitriangular=True)  # y = Lz => z = L^{-1}y
        x_flat = torch.linalg.solve_triangular(upper, z_flat, upper=True, unitriangular=False)  # z = Ux => x = U^{-1}z
        x_flat = x_flat.squeeze(-1)

        output = unflatten_event(x_flat, self.event_shape)
        return output, -self.log_determinant(upper_log_diag)
//  [python]
def torchqr_grad(
    a: Array,
    q: Array,
    r: Array,
    dq: Array,
    dr: Array,
) -> Array:
    """Get the gradient for Qr."""

    if r.shape[-2] > r.shape[-1] and q.shape[-2] == q.shape[-1]:
        raise NotImplementedError(
            "QrGrad not implemented when nrows > ncols "
            "and full_matrices is true. Received r.shape="
            f"{r.shape} with nrows={r.shape[-2]}"
            f"and ncols={r.shape[-1]}."
        )

    def _TriangularSolve(x: Array, r: Array) -> Array:
        """Equivalent to matmul(x, adjoint(matrix_inverse(r))) if r is upper-tri."""
        return torch.triangular_solve(
            x.adjoint(), r, upper=True, unitriangular=False
        )[0].adjoint()

    def _QrGradSquareAndDeepMatrices(q: Array, r: Array, dq: Array, dr: Array) -> Array:
        """
        Get the gradient for matrix orders num_rows >= num_cols and full_matrices is false.
        """

        # Modification begins
        rdiag = torch.linalg.diagonal(r)
        # if abs(rdiag[i]) < qr_epsilon then rdiag[i] = qr_epsilon otherwise keep the old value
        qr_epsilon_diag = torch.ones_like(rdiag) * qr_epsilon
        rdiag = torch.where(rdiag.abs() < qr_epsilon, qr_epsilon_diag, rdiag)
        r = torch.diagonal_scatter(r, rdiag, dim1=-2, dim2=-1)
        # delta_dq = math_ops.matmul(q, math_ops.matmul(dr, tf.linalg.adjoint(delta_r)))
        # dq = dq + delta_dq
        # Modification ends

        qdq = torch.matmul(q.adjoint(), dq)
        qdq_ = qdq - qdq.adjoint()
        rdr = torch.matmul(r, dr.adjoint())
        rdr_ = rdr - rdr.adjoint()
        tril = torch.tril(qdq_ + rdr_)

        grad_a = torch.matmul(q, dr + _TriangularSolve(tril, r))
        grad_b = _TriangularSolve(dq - torch.matmul(q, qdq), r)
        ret = grad_a + grad_b

        if q.is_complex():
            m = rdr - qdq.adjoint()
            eyem = torch.diagonal_scatter(
                torch.zeros_like(m), torch.linalg.diagonal(m), dim1=-2, dim2=-1
            )
            correction = eyem - torch.real(eyem).to(dtype=q.dtype)
            ret = ret + _TriangularSolve(torch.matmul(q, correction.adjoint()), r)

        return ret

    num_rows, num_cols = q.shape[-2], r.shape[-1]

    if num_rows >= num_cols:
        return _QrGradSquareAndDeepMatrices(q, r, dq, dr)

    y = a[..., :, num_rows:]
    u = r[..., :, :num_rows]
    dv = dr[..., :, num_rows:]
    du = dr[..., :, :num_rows]
    dy = torch.matmul(q, dv)
    dx = _QrGradSquareAndDeepMatrices(q, u, dq + torch.matmul(y, dv.adjoint()), du)
    return torch.cat([dx, dy], dim=-1)
//  [python]
def torchqr_grad(
    a: Array,
    q: Array,
    r: Array,
    dq: Array,
    dr: Array,
) -> Array:
    """Get the gradient for Qr."""

    if r.shape[-2] > r.shape[-1] and q.shape[-2] == q.shape[-1]:
        raise NotImplementedError(
            "QrGrad not implemented when nrows > ncols "
            "and full_matrices is true. Received r.shape="
            f"{r.shape} with nrows={r.shape[-2]}"
            f"and ncols={r.shape[-1]}."
        )

    def _TriangularSolve(x: Array, r: Array) -> Array:
        """Equivalent to matmul(x, adjoint(matrix_inverse(r))) if r is upper-tri."""
        return torch.linalg.solve_triangular(
            r, x.adjoint(), upper=True, unitriangular=False
        ).adjoint()

    def _QrGradSquareAndDeepMatrices(q: Array, r: Array, dq: Array, dr: Array) -> Array:
        """
        Get the gradient for matrix orders num_rows >= num_cols and full_matrices is false.
        """

        # Modification begins
        rdiag = torch.linalg.diagonal(r)
        # if abs(rdiag[i]) < qr_epsilon then rdiag[i] = qr_epsilon otherwise keep the old value
        qr_epsilon_diag = torch.ones_like(rdiag) * qr_epsilon
        rdiag = torch.where(rdiag.abs() < qr_epsilon, qr_epsilon_diag, rdiag)
        r = torch.diagonal_scatter(r, rdiag, dim1=-2, dim2=-1)
        # delta_dq = math_ops.matmul(q, math_ops.matmul(dr, tf.linalg.adjoint(delta_r)))
        # dq = dq + delta_dq
        # Modification ends

        qdq = torch.matmul(q.adjoint(), dq)
        qdq_ = qdq - qdq.adjoint()
        rdr = torch.matmul(r, dr.adjoint())
        rdr_ = rdr - rdr.adjoint()
        tril = torch.tril(qdq_ + rdr_)

        grad_a = torch.matmul(q, dr + _TriangularSolve(tril, r))
        grad_b = _TriangularSolve(dq - torch.matmul(q, qdq), r)
        ret = grad_a + grad_b

        if q.is_complex():
            m = rdr - qdq.adjoint()
            eyem = torch.diagonal_scatter(
                torch.zeros_like(m), torch.linalg.diagonal(m), dim1=-2, dim2=-1
            )
            correction = eyem - torch.real(eyem).to(dtype=q.dtype)
            ret = ret + _TriangularSolve(torch.matmul(q, correction.adjoint()), r)

        return ret

    num_rows, num_cols = q.shape[-2], r.shape[-1]

    if num_rows >= num_cols:
        return _QrGradSquareAndDeepMatrices(q, r, dq, dr)

    y = a[..., :, num_rows:]
    u = r[..., :, :num_rows]
    dv = dr[..., :, num_rows:]
    du = dr[..., :, :num_rows]
    dy = torch.matmul(q, dv)
    dx = _QrGradSquareAndDeepMatrices(q, u, dq + torch.matmul(y, dv.adjoint()), du)
    return torch.cat([dx, dy], dim=-1)
//  [python]
def _inverse(self, y, **kwargs):
        """Inverse transform.

        Attributes
        ----------
        y : torch.tensor
            The input vector. The transform is applied to the last dimension.
        kwargs : dict
            keyword arguments to satisfy the interface

        Returns
        -------
         x : torch.tensor
            W^T*(y-b)
        dlogp : torch.tensor
            natural log of the Jacobian determinant
        """
        R = self._make_r()
        dlogp = torch.ones_like(y[...,0,None])*(-torch.sum(torch.log(torch.abs(torch.diagonal(R)))))
        try:
            x, _ = torch.triangular_solve((y-self.b)[...,None], R)
        except AttributeError:
            # legacy call for torch < 1.11
            x, _ = torch.triangular_solve((y-self.b)[...,None], R)
        return x[...,0], dlogp
//  [python]
def _inverse(self, y, **kwargs):
        """Inverse transform.

        Attributes
        ----------
        y : torch.tensor
            The input vector. The transform is applied to the last dimension.
        kwargs : dict
            keyword arguments to satisfy the interface

        Returns
        -------
         x : torch.tensor
            W^T*(y-b)
        dlogp : torch.tensor
            natural log of the Jacobian determinant
        """
        R = self._make_r()
        dlogp = torch.ones_like(y[...,0,None])*(-torch.sum(torch.log(torch.abs(torch.diagonal(R)))))
        try:
            x = torch.linalg.solve_triangular(R, (y-self.b)[...,None], upper=True)
        except AttributeError:
            # legacy call for torch < 1.11
            x, _ = torch.triangular_solve((y-self.b)[...,None], R)
        return x[...,0], dlogp
//  [python]
def tril_solve(L, B):
    """
    Solve a linear system L @ X = B with a lower-triangular matrix L.
    """
    return torch.triangular_solve(B, L, upper=False)[0]
//  [python]
def tril_solve(L, B):
    """
    Solve a linear system L @ X = B with a lower-triangular matrix L.
    """
    return torch.linalg.solve_triangular(L, B, upper=False)
//  [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: Tensor) -> Tensor:
        return torch.triangular_solve(
            torch.triangular_solve(
                y.unsqueeze(-1),
                self.L,
                upper=False,
                unitriangular=True,
            )[0],
            self.U,
            upper=True,
            unitriangular=True,
        )[0].squeeze(-1)
//  [python]
def _inverse(self, y: Tensor) -> Tensor:
        return torch.linalg.solve_triangular(
            self.U,
            torch.linalg.solve_triangular(
                self.L,
                y.unsqueeze(-1),
                upper=False,
                unitriangular=True,
            ),
            upper=True,
            unitriangular=True,
        ).squeeze(-1)
//  [python]
def test_dtrsv(results):
    N = 32_768
    A = (sp.eye(N) * 2 - sp.eye(N, k=-1)).to_dense()
    x = torch.ones(N)

    A_c = A.to(gpu)
    x_c = x.to(gpu)

    A_np = A.numpy()
    x_np = x.numpy()
    A_cp = cupy.asarray(A_c)
    x_cp = cupy.asarray(x_c)

    cupyx.scipy.linalg.solve_triangular(A_cp, x_cp, lower=True)

    it = 20
    results.add(TestResult(
        name='DTRSV Forward',
        cpu_numml=time_op(lambda: torch.solve(x.reshape((-1, 1)), A, upper=False).solution, it),
        cpu_sci=time_op(lambda: scipy.linalg.solve_triangular(A_np, x_np, lower=True), it),
        gpu_numml=time_op(lambda: torch.solve(x_c.reshape((-1, 1)), A_c, upper=False).solution, it),
        gpu_cupy=time_op(lambda: cupyx.scipy.linalg.solve_triangular(A_cp, x_cp, lower=True), it),
        iterations=it,
        n=N
    ))

    A.requires_grad = True
    A_c.requires_grad = True

    it = 10
    results.add(TestResult(
        name='DTRSV Backward',
        cpu_numml=time_op(lambda: torch.solve(x.reshape((-1, 1)), A, upper=False).solution.sum().backward(), it),
        gpu_numml=time_op(lambda: torch.solve(x_c.reshape((-1, 1)), A_c, upper=False).solution.sum().backward(), it),
        iterations=it,
        n=N
    ))
    results.add(TestResultsMidrule())
//  [python]
def test_dtrsv(results):
    N = 32_768
    A = (sp.eye(N) * 2 - sp.eye(N, k=-1)).to_dense()
    x = torch.ones(N)

    A_c = A.to(gpu)
    x_c = x.to(gpu)

    A_np = A.numpy()
    x_np = x.numpy()
    A_cp = cupy.asarray(A_c)
    x_cp = cupy.asarray(x_c)

    cupyx.scipy.linalg.solve_triangular(A_cp, x_cp, lower=True)

    it = 20
    results.add(TestResult(
        name='DTRSV Forward',
        cpu_numml=time_op(lambda: torch.linalg.solve_triangular(A, x.reshape((-1, 1)), upper=False, unitriangular=False), it),
        cpu_sci=time_op(lambda: scipy.linalg.solve_triangular(A_np, x_np, lower=True), it),
        gpu_numml=time_op(lambda: torch.linalg.solve_triangular(A_c, x_c.reshape((-1, 1)), upper=False, unitriangular=False), it),
        gpu_cupy=time_op(lambda: cupyx.scipy.linalg.solve_triangular(A_cp, x_cp, lower=True), it),
        iterations=it,
        n=N
    ))

    A.requires_grad = True
    A_c.requires_grad = True

    it = 10
    results.add(TestResult(
        name='DTRSV Backward',
        cpu_numml=time_op(lambda: torch.linalg.solve_triangular(A, x.reshape((-1, 1)), upper=False, unitriangular=False).sum().backward(), it),
        gpu_numml=time_op(lambda: torch.linalg.solve_triangular(A_c, x_c.reshape((-1, 1)), upper=False, unitriangular=False).sum().backward(), it),
        iterations=it,
        n=N
    ))
    results.add(TestResultsMidrule())
//  [python]
def solve_compact_depth_hierarchical(L_mm, Knm_Kmminv, mean_log_depth, log_depth_obs, log_depth_stdev_inv):
    batch_size, n, m = Knm_Kmminv.shape
    device = Knm_Kmminv.device

    A = torch.empty((batch_size, m+n, m), device=device)
    identity = torch.eye(m, device=device).reshape((1,m,m)).repeat(batch_size,1,1)
    L_inv = torch.triangular_solve(identity, L_mm, upper=False).solution
    A[:,:m,:] = L_inv
    A[:,m:,:] = log_depth_stdev_inv*Knm_Kmminv

    b = torch.empty((batch_size, m+n, 1), device=device)
    b[:,:m,:] = torch.triangular_solve(mean_log_depth*torch.ones((batch_size, m, 1), device=device), L_mm, upper=False).solution
    b[:,m:,:] = log_depth_stdev_inv*(log_depth_obs + torch.sum(Knm_Kmminv * mean_log_depth, dim=(2), keepdim=True) - mean_log_depth)

    compact_log_depth, _, _, _ = torch.linalg.lstsq(A, b)

    residuals = torch.matmul(A, compact_log_depth) - b
    total_err = torch.sum(torch.square(residuals), dim=(1,2))

    return compact_log_depth, total_err
//  [python]
def solve_compact_depth_hierarchical(L_mm, Knm_Kmminv, mean_log_depth, log_depth_obs, log_depth_stdev_inv):
    batch_size, n, m = Knm_Kmminv.shape
    device = Knm_Kmminv.device

    A = torch.empty((batch_size, m+n, m), device=device)
    identity = torch.eye(m, device=device).reshape((1,m,m)).repeat(batch_size,1,1)
    L_inv = torch.linalg.solve_triangular(L_mm, identity, upper=False)
    A[:,:m,:] = L_inv
    A[:,m:,:] = log_depth_stdev_inv*Knm_Kmminv

    b = torch.empty((batch_size, m+n, 1), device=device)
    b[:,:m,:] = torch.linalg.solve_triangular(L_mm, mean_log_depth*torch.ones((batch_size, m, 1), device=device), upper=False)
    b[:,m:,:] = log_depth_stdev_inv*(log_depth_obs + torch.sum(Knm_Kmminv * mean_log_depth, dim=(2), keepdim=True) - mean_log_depth)

    compact_log_depth, _, _, _ = torch.linalg.lstsq(A, b)

    residuals = torch.matmul(A, compact_log_depth) - b
    total_err = torch.sum(torch.square(residuals), dim=(1,2))

    return compact_log_depth, total_err
//  [python]
def sample_gpu_dense(self, sample_shape):
        normal_samples = torch.randn(*sample_shape, device=self.device)
        res = torch.triangular_solve(
            torch.transpose(normal_samples, 1, 2), self.dense_upper_matrix, upper=True
        ).solution
        return torch.transpose(res, 1, 2)
//  [python]
def sample_gpu_dense(self, sample_shape):
        normal_samples = torch.randn(*sample_shape, device=self.device)
        res = torch.linalg.solve_triangular(
            self.dense_upper_matrix, torch.transpose(normal_samples, 1, 2), upper=True
        )
        return torch.transpose(res, 1, 2)
//  [python]
def _iLyy_Kyx(self) -> Tensor:
        r"""Cov(Y, Y)^{-1/2}Cov(Y, X)`."""
        return torch.triangular_solve(
            self.cross_covariance_matrix.transpose(-1, -2),
            self.gauss.scale_tril,
            upper=False,
        )[0]
//  [python]
def _iLyy_Kyx(self) -> Tensor:
        r"""Cov(Y, Y)^{-1/2}Cov(Y, X)`."""
        return torch.linalg.solve_triangular(
            self.gauss.scale_tril,
            self.cross_covariance_matrix.transpose(-1, -2),
            upper=False,
        )
//  [python]
def _iKyy_Kyx(self) -> Tensor:
        r"""Cov(Y, Y)^{-1}Cov(Y, X)`."""
        return torch.triangular_solve(
            self._iLyy_Kyx,
            self.gauss.scale_tril.transpose(-1, -2),
            upper=True,
        )[0]
//  [python]
def _iKyy_Kyx(self) -> Tensor:
        r"""Cov(Y, Y)^{-1}Cov(Y, X)`."""
        return torch.linalg.solve_triangular(
            self.gauss.scale_tril.transpose(-1, -2),
            self._iLyy_Kyx,
            upper=True,
        )
//  [python]
def _iLxx_Kxy(self) -> Tensor:
        r"""Cov(X, X)^{-1/2}Cov(X, Y)`."""
        return torch.triangular_solve(self.cross_covariance_matrix, self.trunc.scale_tril, upper=False)[0]
//  [python]
def _iLxx_Kxy(self) -> Tensor:
        r"""Cov(X, X)^{-1/2}Cov(X, Y)`."""
        return torch.linalg.solve_triangular(
            self.trunc.scale_tril, self.cross_covariance_matrix, upper=False
        )
//  [python]
def init_leg_s(N, max_length=1024, measure="legs", discretization="bilinear"):
    A, B = transition(measure, N)
    A = torch.from_numpy(A.copy()).float().cuda()
    B = torch.from_numpy(B.copy().squeeze(-1)).float().cuda()
    A_stacked = torch.empty((max_length, N, N), device="cuda")
    B_stacked = torch.empty((max_length, N), device="cuda")
    for t in range(1, max_length + 1):
        At = A / t
        Bt = B / t
        if discretization == "forward":
            A_stacked[t - 1] = torch.eye(N, device="cuda") + At
            B_stacked[t - 1] = Bt
        elif discretization == "backward":
            A_stacked[t - 1] = torch.triangular_solve(
                torch.eye(N, device="cuda"), torch.eye(N, device="cuda") - At, upper=False
            ).solution
            B_stacked[t - 1] = torch.solve(Bt, torch.eye(N, device="cuda") - At)[0]
        elif discretization == "bilinear":
            A_stacked[t - 1] = torch.triangular_solve(
                torch.eye(N, device="cuda") + At / 2, torch.eye(N, device="cuda") - At / 2, upper=False
            ).solution
            B_stacked[t - 1] = torch.solve(Bt, torch.eye(N, device="cuda") - At / 2)[0]
        else:  # ZOH
            A_stacked[t - 1] = torch.linalg.matrix_exp(A * (math.log(t + 1) - math.log(t)))
            B_stacked[t - 1] = torch.linalg.solve(A, A_stacked[t - 1] @ B - B)

    E = (B[:, None].cpu() * ss.eval_legendre(np.arange(N)[:, None], 2 * np.linspace(0.0, 1.0, max_length) - 1)).T

    return A_stacked, B_stacked, E.float().cuda()
//  [python]
def init_leg_s(N, max_length=1024, measure="legs", discretization="bilinear"):
    A, B = transition(measure, N)
    A = torch.from_numpy(A.copy()).float().cuda()
    B = torch.from_numpy(B.copy().squeeze(-1)).float().cuda()
    A_stacked = torch.empty((max_length, N, N), device="cuda")
    B_stacked = torch.empty((max_length, N), device="cuda")
    for t in range(1, max_length + 1):
        At = A / t
        Bt = B / t
        if discretization == "forward":
            A_stacked[t - 1] = torch.eye(N, device="cuda") + At
            B_stacked[t - 1] = Bt
        elif discretization == "backward":
            A_stacked[t - 1] = torch.linalg.solve_triangular(
                torch.eye(N, device="cuda") - At, torch.eye(N, device="cuda"), upper=False
            )
            B_stacked[t - 1] = torch.linalg.solve(torch.eye(N, device="cuda") - At, Bt)
        elif discretization == "bilinear":
            A_stacked[t - 1] = torch.linalg.solve_triangular(
                torch.eye(N, device="cuda") - At / 2,
                torch.eye(N, device="cuda") + At / 2,
                upper=False,
            )
            B_stacked[t - 1] = torch.linalg.solve(torch.eye(N, device="cuda") - At / 2, Bt)
        else:  # ZOH
            A_stacked[t - 1] = torch.linalg.matrix_exp(A * (math.log(t + 1) - math.log(t)))
            B_stacked[t - 1] = torch.linalg.solve(A, A_stacked[t - 1] @ B - B)

    E = (B[:, None].cpu() * ss.eval_legendre(np.arange(N)[:, None], 2 * np.linspace(0.0, 1.0, max_length) - 1)).T

    return A_stacked, B_stacked, E.float().cuda()
//  [python]
def gp_ml_cost(
    logzm,
    log_median_depths,
    L_mm,
    dlogzm_dPw,
    dlogzm_dTwc,
    p_inds_batched,
    pose_inds_batched,
    H,
    g,
    sigma,
):
    b, m, _ = L_mm.shape
    device = L_mm.device

    da_dlogzm = torch.eye(m, device=device).reshape((1, m, m)).repeat(b, 1, 1)
    dr_dlogzm = torch.triangular_solve(da_dlogzm, L_mm, upper=False).solution  # L_inv

    # h(x) - z with z = 0
    r = dr_dlogzm @ (logzm - log_median_depths)  # L_inv @ logzm

    # Point Jacobians
    dr_dPw = (
        dr_dlogzm[:, :, :, None] * dlogzm_dPw[:, None, :, 0, :]
    )  # (B.M,M,3) <- (B,M,M) x (B,M,1,3)
    dr_dPw = torch.reshape(dr_dPw, (b, m, 1, 3 * m))  # (B.M,1,3M)
    # Pose Jacobians
    dr_dTwc = (dr_dlogzm @ dlogzm_dTwc[:, :, 0, :]).unsqueeze(
        2
    )  # (B,M,1,6) <- (B,M,M) x (B,M,1,6)

    info_scale = 1.0 / (sigma**2)

    # Gradient g = -Jt @ r
    grad_Pw = info_scale * lin_sys.get_gradient(dr_dPw, r)
    grad_Twc = info_scale * lin_sys.get_gradient(dr_dTwc, r)
    H_Pw = info_scale * lin_sys.get_hessian_diag_block(dr_dPw)
    H_Twc = info_scale * lin_sys.get_hessian_diag_block(dr_dTwc)
    H_Twc_Pw = info_scale * lin_sys.get_hessian_off_diag_block(dr_dTwc, dr_dPw)

    grads_flat = torch.cat((grad_Pw.flatten(), grad_Twc.flatten()))
    grad_inds_flat = torch.cat((p_inds_batched.flatten(), pose_inds_batched.flatten()))
    lin_sys.accumulate_gradient_scatter(grads_flat, g, grad_inds_flat)

    W = H.shape[1]
    H_inds_Pw = lin_sys.row_col_to_lin_index(
        p_inds_batched[:, :, None], p_inds_batched[:, None, :], W
    )
    H_inds_Twc = lin_sys.row_col_to_lin_index(
        pose_inds_batched[:, :, None], pose_inds_batched[:, None, :], W
    )
    H_inds_Twc_Pw = lin_sys.row_col_to_lin_index(
        pose_inds_batched[:, :, None], p_inds_batched[:, None, :], W
    )
    H_inds_Twc_Pw2 = lin_sys.row_col_to_lin_index(
        p_inds_batched[:, :, None], pose_inds_batched[:, None, :], W
    )

    H_flat = torch.cat(
        (H_Pw.flatten(), H_Twc.flatten(), H_Twc_Pw.flatten(), H_Twc_Pw.mT.flatten())
    )
    H_inds_flat = torch.cat(
        (
            H_inds_Pw.flatten(),
            H_inds_Twc.flatten(),
            H_inds_Twc_Pw.flatten(),
            H_inds_Twc_Pw2.flatten(),
        )
    )
    lin_sys.accumulate_hessian_scatter(H_flat, H, H_inds_flat)

    total_err = torch.sum(info_scale * torch.square(r))

    return total_err
//  [python]
def gp_ml_cost(
    logzm,
    log_median_depths,
    L_mm,
    dlogzm_dPw,
    dlogzm_dTwc,
    p_inds_batched,
    pose_inds_batched,
    H,
    g,
    sigma,
):
    b, m, _ = L_mm.shape
    device = L_mm.device

    da_dlogzm = torch.eye(m, device=device).reshape((1, m, m)).repeat(b, 1, 1)
    dr_dlogzm = torch.linalg.solve_triangular(L_mm, da_dlogzm, upper=False)  # L_inv

    # h(x) - z with z = 0
    r = dr_dlogzm @ (logzm - log_median_depths)  # L_inv @ logzm

    # Point Jacobians
    dr_dPw = (
        dr_dlogzm[:, :, :, None] * dlogzm_dPw[:, None, :, 0, :]
    )  # (B.M,M,3) <- (B,M,M) x (B,M,1,3)
    dr_dPw = torch.reshape(dr_dPw, (b, m, 1, 3 * m))  # (B.M,1,3M)
    # Pose Jacobians
    dr_dTwc = (dr_dlogzm @ dlogzm_dTwc[:, :, 0, :]).unsqueeze(
        2
    )  # (B,M,1,6) <- (B,M,M) x (B,M,1,6)

    info_scale = 1.0 / (sigma**2)

    # Gradient g = -Jt @ r
    grad_Pw = info_scale * lin_sys.get_gradient(dr_dPw, r)
    grad_Twc = info_scale * lin_sys.get_gradient(dr_dTwc, r)
    H_Pw = info_scale * lin_sys.get_hessian_diag_block(dr_dPw)
    H_Twc = info_scale * lin_sys.get_hessian_diag_block(dr_dTwc)
    H_Twc_Pw = info_scale * lin_sys.get_hessian_off_diag_block(dr_dTwc, dr_dPw)

    grads_flat = torch.cat((grad_Pw.flatten(), grad_Twc.flatten()))
    grad_inds_flat = torch.cat((p_inds_batched.flatten(), pose_inds_batched.flatten()))
    lin_sys.accumulate_gradient_scatter(grads_flat, g, grad_inds_flat)

    W = H.shape[1]
    H_inds_Pw = lin_sys.row_col_to_lin_index(
        p_inds_batched[:, :, None], p_inds_batched[:, None, :], W
    )
    H_inds_Twc = lin_sys.row_col_to_lin_index(
        pose_inds_batched[:, :, None], pose_inds_batched[:, None, :], W
    )
    H_inds_Twc_Pw = lin_sys.row_col_to_lin_index(
        pose_inds_batched[:, :, None], p_inds_batched[:, None, :], W
    )
    H_inds_Twc_Pw2 = lin_sys.row_col_to_lin_index(
        p_inds_batched[:, :, None], pose_inds_batched[:, None, :], W
    )

    H_flat = torch.cat(
        (H_Pw.flatten(), H_Twc.flatten(), H_Twc_Pw.flatten(), H_Twc_Pw.mT.flatten())
    )
    H_inds_flat = torch.cat(
        (
            H_inds_Pw.flatten(),
            H_inds_Twc.flatten(),
            H_inds_Twc_Pw.flatten(),
            H_inds_Twc_Pw2.flatten(),
        )
    )
    lin_sys.accumulate_hessian_scatter(H_flat, H, H_inds_flat)

    total_err = torch.sum(info_scale * torch.square(r))

    return total_err
//  [python]
def test_triangular_solve(batch_shape, m, n):
    import torch

    b = torch.randn(batch_shape + (n, m))
    L = torch.randn(batch_shape + (n, n)).tril_()
    L.diagonal(dim1=-2, dim2=-1).abs_().add_(0.2)
    U = L.transpose(-1, -2)

    actual = ops.triangular_solve(b, L, upper=False)
    expected = torch.solve(b, L, upper=False)
    assert_close(actual, expected)
    assert_close(L @ actual, b, atol=1e-4, rtol=1e-3)

    actual = ops.triangular_solve(b, U, upper=True)
    expected = torch.solve(b, U, upper=True)
    assert_close(actual, expected)
    assert_close(U @ actual, b, atol=1e-4, rtol=1e-3)

    actual = ops.triangular_solve(b, L, upper=False, transpose=True)
    expected = torch.solve(b, L.transpose(-1, -2), upper=True)
    assert_close(actual, expected)
    assert_close(L.transpose(-1, -2) @ actual, b, atol=1e-4, rtol=1e-3)

    actual = ops.triangular_solve(b, U, upper=True, transpose=True)
    expected = torch.solve(b, U.transpose(-1, -2), upper=False)
    assert_close(actual, expected)
    assert_close(U.transpose(-1, -2) @ actual, b, atol=1e-4, rtol=1e-3)
//  [python]
def test_triangular_solve(batch_shape, m, n):
    import torch

    b = torch.randn(batch_shape + (n, m))
    L = torch.randn(batch_shape + (n, n)).tril_()
    L.diagonal(dim1=-2, dim2=-1).abs_().add_(0.2)
    U = L.transpose(-1, -2)

    actual = ops.triangular_solve(b, L, upper=False)
    expected = torch.linalg.solve_triangular(L, b, upper=False)
    assert_close(actual, expected)
    assert_close(L @ actual, b, atol=1e-4, rtol=1e-3)

    actual = ops.triangular_solve(b, U, upper=True)
    expected = torch.linalg.solve_triangular(U, b, upper=True)
    assert_close(actual, expected)
    assert_close(U @ actual, b, atol=1e-4, rtol=1e-3)

    actual = ops.triangular_solve(b, L, upper=False, transpose=True)
    expected = torch.linalg.solve_triangular(L.transpose(-1, -2), b, upper=True)
    assert_close(actual, expected)
    assert_close(L.transpose(-1, -2) @ actual, b, atol=1e-4, rtol=1e-3)

    actual = ops.triangular_solve(b, U, upper=True, transpose=True)
    expected = torch.linalg.solve_triangular(U.transpose(-1, -2), b, upper=False)
    assert_close(actual, expected)
    assert_close(U.transpose(-1, -2) @ actual, b, atol=1e-4, rtol=1e-3)
//  [python]
def backward(ctx, grad_output):
        jitter = 1.0e-8  # do i really need this?
        z, epsilon, L = ctx.saved_tensors

        dim = L.shape[0]
        g = grad_output
        loc_grad = sum_leftmost(grad_output, -1)

        identity = eye_like(g, dim)
        R_inv = torch.triangular_solve(identity, L.t(), upper=True)[0]

        z_ja = z.unsqueeze(-1)
        g_R_inv = torch.matmul(g, R_inv).unsqueeze(-2)
        epsilon_jb = epsilon.unsqueeze(-2)
        g_ja = g.unsqueeze(-1)
        diff_L_ab = 0.5 * sum_leftmost(g_ja * epsilon_jb + g_R_inv * z_ja, -2)

        Sigma_inv = torch.mm(R_inv, R_inv.t())
        V, D, _ = torch.svd(Sigma_inv + jitter)
        D_outer = D.unsqueeze(-1) + D.unsqueeze(0)

        expand_tuple = tuple([-1] * (z.dim() - 1) + [dim, dim])
        z_tilde = identity * torch.matmul(z, V).unsqueeze(-1).expand(*expand_tuple)
        g_tilde = identity * torch.matmul(g, V).unsqueeze(-1).expand(*expand_tuple)

        Y = sum_leftmost(
            torch.matmul(z_tilde, torch.matmul(1.0 / D_outer, g_tilde)), -2
        )
        Y = torch.mm(V, torch.mm(Y, V.t()))
        Y = Y + Y.t()

        Tr_xi_Y = torch.mm(torch.mm(Sigma_inv, Y), R_inv) - torch.mm(
            Y, torch.mm(Sigma_inv, R_inv)
        )
        diff_L_ab += 0.5 * Tr_xi_Y
        L_grad = torch.tril(diff_L_ab)

        return loc_grad, L_grad, None
//  [python]
def backward(ctx, grad_output):
        jitter = 1.0e-8  # do i really need this?
        z, epsilon, L = ctx.saved_tensors

        dim = L.shape[0]
        g = grad_output
        loc_grad = sum_leftmost(grad_output, -1)

        identity = eye_like(g, dim)
        R_inv = torch.linalg.solve_triangular(L.t(), identity, upper=True)

        z_ja = z.unsqueeze(-1)
        g_R_inv = torch.matmul(g, R_inv).unsqueeze(-2)
        epsilon_jb = epsilon.unsqueeze(-2)
        g_ja = g.unsqueeze(-1)
        diff_L_ab = 0.5 * sum_leftmost(g_ja * epsilon_jb + g_R_inv * z_ja, -2)

        Sigma_inv = torch.mm(R_inv, R_inv.t())
        V, D, _ = torch.svd(Sigma_inv + jitter)
        D_outer = D.unsqueeze(-1) + D.unsqueeze(0)

        expand_tuple = tuple([-1] * (z.dim() - 1) + [dim, dim])
        z_tilde = identity * torch.matmul(z, V).unsqueeze(-1).expand(*expand_tuple)
        g_tilde = identity * torch.matmul(g, V).unsqueeze(-1).expand(*expand_tuple)

        Y = sum_leftmost(
            torch.matmul(z_tilde, torch.matmul(1.0 / D_outer, g_tilde)), -2
        )
        Y = torch.mm(V, torch.mm(Y, V.t()))
        Y = Y + Y.t()

        Tr_xi_Y = torch.mm(torch.mm(Sigma_inv, Y), R_inv) - torch.mm(
            Y, torch.mm(Sigma_inv, R_inv)
        )
        diff_L_ab += 0.5 * Tr_xi_Y
        L_grad = torch.tril(diff_L_ab)

        return loc_grad, L_grad, None
//  [python]
def _iKxx_Kxy(self) -> Tensor:
        r"""Cov(X, X)^{-1}Cov(X, Y)`."""
        return torch.triangular_solve(
            self._iLxx_Kxy,
            self.trunc.scale_tril.transpose(-1, -2),
            upper=True,
        ).solution
//  [python]
def _iKxx_Kxy(self) -> Tensor:
        r"""Cov(X, X)^{-1}Cov(X, Y)`."""
        return torch.linalg.solve_triangular(
            self.trunc.scale_tril.transpose(-1, -2),
            self._iLxx_Kxy,
            upper=True,
        )
//  [python]
def covar_cache(self):
        # Here, the covar_cache is going to be K_{UU}^{-1/2} K_{UX}( K_{XX} + \sigma^2 I )^{-1} K_{XU} K_{UU}^{-1/2}
        # This is easily computed using Woodbury
        # K_{XX} + \sigma^2 I = R R^T + \sigma^2 I
        #                     = \sigma^{-2} ( I - \sigma^{-2} R (I + \sigma^{-2} R^T R)^{-1} R^T  )
        train_train_covar = self.lik_train_train_covar.evaluate_kernel()

        # Get terms needed for woodbury
        root = train_train_covar._linear_op.root_decomposition().root.to_dense()  # R
        inv_diag = train_train_covar._diag_tensor.inverse()  # \sigma^{-2}

        # Form LT using woodbury
        ones = torch.tensor(1.0, dtype=root.dtype, device=root.device)
        chol_factor = to_linear_operator(root.transpose(-1, -2) @ (inv_diag @ root)).add_diagonal(
            ones
        )  # (I + \sigma^{-2} R^T R)^{-1}
        woodbury_term = inv_diag @ torch.solve(root.transpose(-1, -2), chol_factor.cholesky().to_dense())[0].transpose(-1, -2)
        # woodbury_term @ woodbury_term^T = \sigma^{-2} R (I + \sigma^{-2} R^T R)^{-1} R^T \sigma^{-2}

        inverse = AddedDiagLinearOperator(
            inv_diag, MatmulLinearOperator(-woodbury_term, woodbury_term.transpose(-1, -2))
        )
        # \sigma^{-2} ( I - \sigma^{-2} R (I + \sigma^{-2} R^T R)^{-1} R^T  )

        return root.transpose(-1, -2) @ (inverse @ root)
//  [python]
def covar_cache(self):
        # Here, the covar_cache is going to be K_{UU}^{-1/2} K_{UX}( K_{XX} + \sigma^2 I )^{-1} K_{XU} K_{UU}^{-1/2}
        # This is easily computed using Woodbury
        # K_{XX} + \sigma^2 I = R R^T + \sigma^2 I
        #                     = \sigma^{-2} ( I - \sigma^{-2} R (I + \sigma^{-2} R^T R)^{-1} R^T  )
        train_train_covar = self.lik_train_train_covar.evaluate_kernel()

        # Get terms needed for woodbury
        root = train_train_covar._linear_op.root_decomposition().root.to_dense()  # R
        inv_diag = train_train_covar._diag_tensor.inverse()  # \sigma^{-2}

        # Form LT using woodbury
        ones = torch.tensor(1.0, dtype=root.dtype, device=root.device)
        chol_factor = to_linear_operator(root.transpose(-1, -2) @ (inv_diag @ root)).add_diagonal(
            ones
        )  # (I + \sigma^{-2} R^T R)^{-1}
        woodbury_term = inv_diag @ torch.linalg.solve_triangular(
            chol_factor.cholesky().to_dense(), root.transpose(-1, -2), upper=False
        ).transpose(-1, -2)
        # woodbury_term @ woodbury_term^T = \sigma^{-2} R (I + \sigma^{-2} R^T R)^{-1} R^T \sigma^{-2}

        inverse = AddedDiagLinearOperator(
            inv_diag, MatmulLinearOperator(-woodbury_term, woodbury_term.transpose(-1, -2))
        )
        # \sigma^{-2} ( I - \sigma^{-2} R (I + \sigma^{-2} R^T R)^{-1} R^T  )

        return root.transpose(-1, -2) @ (inverse @ root)
//  [python]
def test_triangular_solve(upper):
    b = torch.randn(5, 6)
    A = torch.randn(5, 5)
    expected = torch.triangular_solve(b, A, upper=upper).solution
    actual = torch.solve(b, A.triu() if upper else A.tril()).solution
    assert_close(actual, expected)
    A = A.triu() if upper else A.tril()
    assert_close(A @ actual, b)
//  [python]
def test_triangular_solve(upper):
    b = torch.randn(5, 6)
    A = torch.randn(5, 5)
    expected = torch.triangular_solve(b, A, upper=upper).solution
    actual = torch.linalg.solve_triangular(A, b, upper=upper)
    assert_close(actual, expected)
    A = A.triu() if upper else A.tril()
    assert_close(A @ actual, b)
//  [python]
def weight_inverse(self):
        lower, upper = self._create_lower_upper()
        identity = torch.eye(self.features, self.features)
        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)
        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_solver_result_coo_triu(self):
        x = self.solve(self.As_coo_triu, self.Bd, upper=True, unitriangular=self.unitriangular)
        x2 = torch.triangular_solve(self.Bd, self.Ad_triu, upper=True, unitriangular=self.unitriangular).solution
        self.assertTrue(torch.isclose(x, x2, rtol=self.RTOL).all())
//  [python]
def test_solver_result_coo_triu(self):
        x = self.solve(self.As_coo_triu, self.Bd, upper=True, unitriangular=self.unitriangular)
        x2 = torch.linalg.solve_triangular(self.Ad_triu, self.Bd, upper=True, unitriangular=self.unitriangular)
        self.assertTrue(torch.isclose(x, x2, rtol=self.RTOL).all())
//  [python]
def test_triangular_solve_transpose(upper):
    b = torch.randn(5, 6)
    A = torch.randn(5, 5)
    expected = torch.triangular_solve(b, A, upper=upper, transpose=True).solution
    actual = torch.triangular_solve(b, A.T, upper=not upper).solution
    assert_close(actual, expected)
    A = A.triu() if upper else A.tril()
    assert_close(A.T @ actual, b)
//  [python]
def test_triangular_solve_transpose(upper):
    b = torch.randn(5, 6)
    A = torch.randn(5, 5)
    expected = torch.triangular_solve(b, A, upper=upper, transpose=True).solution
    actual = torch.linalg.solve_triangular(A.T, b, upper=not upper)
    assert_close(actual, expected)
    A = A.triu() if upper else A.tril()
    assert_close(A.T @ actual, b)
//  [python]
def test_solver_result_csr_triu(self):
        x = self.solve(self.As_csr_triu, self.Bd, upper=True, unitriangular=self.unitriangular)
        x2 = torch.solve(self.Ad_triu, self.Bd, upper=True, unitriangular=self.unitriangular)
        self.assertTrue(torch.isclose(x, x2, rtol=self.RTOL).all())
//  [python]
def test_solver_result_csr_triu(self):
        x = self.solve(self.As_csr_triu, self.Bd, upper=True, unitriangular=self.unitriangular)
        x2 = torch.linalg.solve_triangular(self.Ad_triu, self.Bd, upper=True, unitriangular=self.unitriangular)
        self.assertTrue(torch.isclose(x, x2, rtol=self.RTOL).all())
//  [python]
def test_solver_result_coo_tril(self):
        x = self.solve(self.As_coo_tril, self.Bd, upper=False, unitriangular=self.unitriangular)
        x2 = torch.triangular_solve(self.Bd, self.Ad_tril, upper=False, unitriangular=self.unitriangular)[0]
        self.assertTrue(torch.isclose(x, x2, rtol=self.RTOL).all())
//  [python]
def test_solver_result_coo_tril(self):
        x = self.solve(self.As_coo_tril, self.Bd, upper=False, unitriangular=self.unitriangular)
        x2 = torch.linalg.solve_triangular(self.Ad_tril, self.Bd, upper=False, unitriangular=self.unitriangular)
        self.assertTrue(torch.isclose(x, x2, rtol=self.RTOL).all())
//  [python]
def test_solver_result_csr_tril(self):
        x = self.solve(self.As_csr_tril, self.Bd, upper=False, unitriangular=self.unitriangular)
        x2 = torch.triangular_solve(self.Ad_tril, self.Bd, upper=False, unitriangular=self.unitriangular).solution
        self.assertTrue(torch.isclose(x, x2, rtol=self.RTOL).all())
//  [python]
def test_solver_result_csr_tril(self):
        x = self.solve(self.As_csr_tril, self.Bd, upper=False, unitriangular=self.unitriangular)
        x2 = torch.linalg.solve_triangular(self.Ad_tril, self.Bd, upper=False, unitriangular=self.unitriangular)
        self.assertTrue(torch.isclose(x, x2, rtol=self.RTOL).all())
//  [python]
def _triangular_solve(x, y, upper=False, transpose=False):
    if y.size(-1) == 1:
        return x / y
    if transpose:
        y = y.transpose(-1, -2)
        upper = not upper
    return torch.triangular_solve(x, y, upper=upper)[0]
//  [python]
def _triangular_solve(x, y, upper=False, transpose=False):
    if y.size(-1) == 1:
        return x / y
    if transpose:
        y = y.transpose(-1, -2)
        upper = not upper
    return torch.linalg.solve_triangular(y, x, upper=upper)