jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def forward(self, 
                shared_parameters,
                losses: list[torch.Tensor, ...],
               ) -> tuple[torch.Tensor, torch.Tensor]:
        # >>> Loss Balance
        L_t, g, u = [], [], []
        
        for i, l in enumerate(losses):
            if self.method in ['hybrid','loss']:
                # If 'hybrid' or 'loss' method selected, then learnable scale s_t is used
                L_t.append(l * self.e.pow(self.s_t[i]) - self.s_t[i])
            else:
                # Else, the losses are unscaled
                L_t.append(l)

            g[i].append(torch.flatten(torch.autograd.grad(L_t[-1], shared_parameters, retain_graph=True)))
            u[i].append(g[i] / torch.linalg.norm(g[i]))
            
        g_t = torch.stack([_g for _g in g])
        u_t = torch.stack([_u for _u in u])
       
        # >>> Gradient Balance
        D = g_t[0,...] - g_t[1:,...]
        U = u_t[0,...] - u_t[1:,...]

        arg0 = D.matmul(U.t())
        arg1 = torch.eye(arg0.shape[-1].data)
        alpha_2T = g_t[0,...].matmul(U.t()).matmul(torch.linalg.solve(arg0,arg1))
        
        alpha = torch.cat([torch.ones(1, device=self.device) - alpha_2T.sum(dim=0, keepdim=True), 
                           alpha_2T], dim=0).squeeze()
        
        # All methods return two outputs: gradient loss and individual losses
        if self.method=='hybrid':
            return torch.sum(L_t * alpha), L_t
        elif self.method=='gradient':
            return torch.sum(L_t * alpha), None
        elif self.method=='loss':
            return None, L_t
//  [python]
def solve(b: Union[torch.Tensor, ComplexTensor], a: Union[torch.Tensor, ComplexTensor]):
    """Solve the linear equation ax = b."""
    # NOTE: Do not mix ComplexTensor and torch.complex in the input!
    # NOTE (wangyou): Until PyTorch 1.9.0, torch.solve does not support
    # mixed input with complex and real tensors.
    if isinstance(a, ComplexTensor) or isinstance(b, ComplexTensor):
        if isinstance(a, ComplexTensor) and isinstance(b, ComplexTensor):
            return FC.solve(b, a)[0]
        else:
            return matmul(inverse(a), b)
    elif is_torch_1_9_plus and (torch.is_complex(a) or torch.is_complex(b)):
        if torch.is_complex(a) and torch.is_complex(b):
            return torch.solve(b, a)[0]
        else:
            return matmul(inverse(a), b)
    else:
        if is_torch_1_8_plus:
            return torch.solve(b, a)[0]
        else:
            return torch.solve(b, a)[0]
//  [python]
def solve(b: Union[torch.Tensor, ComplexTensor], a: Union[torch.Tensor, ComplexTensor]):
    """Solve the linear equation ax = b."""
    # NOTE: Do not mix ComplexTensor and torch.complex in the input!
    # NOTE (wangyou): Until PyTorch 1.9.0, torch.solve does not support
    # mixed input with complex and real tensors.
    if isinstance(a, ComplexTensor) or isinstance(b, ComplexTensor):
        if isinstance(a, ComplexTensor) and isinstance(b, ComplexTensor):
            return FC.solve(b, a)[0]
        else:
            return matmul(inverse(a), b)
    elif is_torch_1_9_plus and (torch.is_complex(a) or torch.is_complex(b)):
        if torch.is_complex(a) and torch.is_complex(b):
            return torch.linalg.solve(a, b)
        else:
            return matmul(inverse(a), b)
    else:
        if is_torch_1_8_plus:
            return torch.linalg.solve(a, b)
        else:
            return torch.solve(b, a)[0]
//  [python]
def hscic(X, Z, Y, ridge_lambda, kernelX='gaussian', kernelX_params=None,
          kernelZ='gaussian', kernelZ_params=None, kernelY='gaussian', kernelY_params=None):
    '''X ind. Z | Y '''
    # todo:
    #  alternative implementation for RFF
    Kx = eval(f'{kernelX}_kernel(X, **kernelX_params)')
    Kz = eval(f'{kernelZ}_kernel(Z, **kernelZ_params)')
    Ky = eval(f'{kernelY}_kernel(Y, **kernelY_params)')

    # https://arxiv.org/pdf/2207.09768.pdf
    WtKyy = torch.solve(Ky, Ky + ridge_lambda  * torch.eye(Ky.shape[0]).to(Ky.device))[0] # * Ky.shape[0] for ridge_lambda
    # todo:
    #   SVD + LOO here? but that doesn't scale with # params
    #   three LOO for all three regressions?
    #   re-do it every few iters?
    # sum_i A_(i.)B_(.i) = tr(AB) = (A * B^T).sum()
    # A = Kyy^T, B = the other one, so the transposes cancel out
    term_1 = (WtKyy * ((Kx * Kz) @ WtKyy)).sum() # tr(WtKyy.T @ (Kx * Kz) @ WtKyy)
    WkKxWk = WtKyy * (Kx @ WtKyy)
    KzWk = Kz @ WtKyy
    term_2 = (WkKxWk * KzWk).sum()
    # here it's crucial that the first dimension is the batch of other matrices
    term_3 = (WkKxWk.sum(dim=0) * (WtKyy * KzWk).sum(dim=0)).sum()

    return (term_1 - 2 * term_2 + term_3) / Ky.shape[0]
//  [python]
def hscic(X, Z, Y, ridge_lambda, kernelX='gaussian', kernelX_params=None,
          kernelZ='gaussian', kernelZ_params=None, kernelY='gaussian', kernelY_params=None):
    '''X ind. Z | Y '''
    # todo:
    #  alternative implementation for RFF
    Kx = eval(f'{kernelX}_kernel(X, **kernelX_params)')
    Kz = eval(f'{kernelZ}_kernel(Z, **kernelZ_params)')
    Ky = eval(f'{kernelY}_kernel(Y, **kernelY_params)')

    # https://arxiv.org/pdf/2207.09768.pdf
    WtKyy = torch.linalg.solve(Ky + ridge_lambda  * torch.eye(Ky.shape[0]).to(Ky.device), Ky) # * Ky.shape[0] for ridge_lambda
    # todo:
    #   SVD + LOO here? but that doesn't scale with # params
    #   three LOO for all three regressions?
    #   re-do it every few iters?
    # sum_i A_(i.)B_(.i) = tr(AB) = (A * B^T).sum()
    # A = Kyy^T, B = the other one, so the transposes cancel out
    term_1 = (WtKyy * ((Kx * Kz) @ WtKyy)).sum() # tr(WtKyy.T @ (Kx * Kz) @ WtKyy)
    WkKxWk = WtKyy * (Kx @ WtKyy)
    KzWk = Kz @ WtKyy
    term_2 = (WkKxWk * KzWk).sum()
    # here it's crucial that the first dimension is the batch of other matrices
    term_3 = (WkKxWk.sum(dim=0) * (WtKyy * KzWk).sum(dim=0)).sum()

    return (term_1 - 2 * term_2 + term_3) / Ky.shape[0]
//  [python]
def get_torch_solve(
        self,
        A, b,
        method,
    ):
        if method == 'inverse':
            res = torch.solve(b, A)[0]
        elif method == 'cholesky':
            # method 1
            # old implementation - recheck torch.allclose(res, torch.linalg.solve(A, b))
            # l = torch.linalg.cholesky(A)
            # z = torch.linalg.solve_triangular(l, b, upper=False)
            # res = torch.linalg.solve_triangular(l.mT, z, upper=False)

            # method 2
            # z = torch.triangular_solve(b, l, transpose=False, upper=False)[0]
            # res = torch.triangular_solve(z, l, transpose=True, upper=False)[0]

            # method 3
            l, _ = torch.linalg.cholesky_ex(A)
            res = torch.cholesky_solve(b, l)

        elif method == 'cholesky-sparse':
            # TODO - ATTENTION!
            #  This implementation takes a lot of memory because of the block diagonal construction.
            raise NotImplementedError
            # if self.tensor_args['dtype'] == torch.float32:
            #     cholesky_fn = cholespy.CholeskySolverF
            # elif self.tensor_args['dtype'] == torch.float64:
            #     cholesky_fn = cholespy.CholeskySolverD
            # else:
            #     raise NotImplementedError

            # convert batch to block diagonal - https://github.com/rgl-epfl/cholespy/issues/26
            # TODO - creating a big block diagonal matrix leads to memory problems
            A_block_diag = torch.block_diag(*A)
            b_stacked = einops.rearrange(b, "b d 1 -> (b d) 1")

            A_sparse = A_block_diag.to_sparse(layout=torch.sparse_coo)

            cholesky_fn = cholespy.CholeskySolverF
            solver = cholesky_fn(
                A_sparse.size()[0],
                A_sparse.indices()[0], A_sparse.indices()[1], A_sparse.values(),
                cholespy.MatrixType.COO
            )
            res_ = torch.zeros_like(b_stacked)

            solver.solve(b_stacked, res_)

            # batchify results
            res = einops.rearrange(res_, "(b d) 1 -> b d 1", b=A.shape[0])

        elif method == 'lstq':
            # empirically slower
            res = torch.linalg.lstsq(A, b)[0]
        else:
            raise NotImplementedError

        return res
//  [python]
def get_torch_solve(
        self,
        A, b,
        method,
    ):
        if method == 'inverse':
            res = torch.linalg.solve(A, b)
        elif method == 'cholesky':
            # method 1
            # old implementation - recheck torch.allclose(res, torch.linalg.solve(A, b))
            # l = torch.linalg.cholesky(A)
            # z = torch.linalg.solve_triangular(l, b, upper=False)
            # res = torch.linalg.solve_triangular(l.mT, z, upper=False)

            # method 2
            # z = torch.triangular_solve(b, l, transpose=False, upper=False)[0]
            # res = torch.triangular_solve(z, l, transpose=True, upper=False)[0]

            # method 3
            l, _ = torch.linalg.cholesky_ex(A)
            res = torch.cholesky_solve(b, l)

        elif method == 'cholesky-sparse':
            # TODO - ATTENTION!
            #  This implementation takes a lot of memory because of the block diagonal construction.
            raise NotImplementedError
            # if self.tensor_args['dtype'] == torch.float32:
            #     cholesky_fn = cholespy.CholeskySolverF
            # elif self.tensor_args['dtype'] == torch.float64:
            #     cholesky_fn = cholespy.CholeskySolverD
            # else:
            #     raise NotImplementedError

            # convert batch to block diagonal - https://github.com/rgl-epfl/cholespy/issues/26
            # TODO - creating a big block diagonal matrix leads to memory problems
            A_block_diag = torch.block_diag(*A)
            b_stacked = einops.rearrange(b, "b d 1 -> (b d) 1")

            A_sparse = A_block_diag.to_sparse(layout=torch.sparse_coo)

            cholesky_fn = cholespy.CholeskySolverF
            solver = cholesky_fn(
                A_sparse.size()[0],
                A_sparse.indices()[0], A_sparse.indices()[1], A_sparse.values(),
                cholespy.MatrixType.COO
            )
            res_ = torch.zeros_like(b_stacked)

            solver.solve(b_stacked, res_)

            # batchify results
            res = einops.rearrange(res_, "(b d) 1 -> b d 1", b=A.shape[0])

        elif method == 'lstq':
            # empirically slower
            res = torch.linalg.lstsq(A, b)[0]
        else:
            raise NotImplementedError

        return res
//  [python]
def _root_find(self, z0, x, *args, **kwargs):

        # Compute forward pass: find root of function outside autograd tape.
        with torch.no_grad():
            z_root = self.solver(
                lambda z: self.f(z, x, *args, **remove_kwargs(kwargs, 'solver_')),
                z0,
                **filter_kwargs(kwargs, 'solver_fwd_'),
            )['result']
            new_z_root = z_root

        if self.training:
            # Re-engage autograd tape (no-op in terms of value of z).
            new_z_root = z_root.requires_grad_() - self.f(z_root.requires_grad_(), x, *args, **remove_kwargs(kwargs, 'solver_'))

            # Set up backward hook for solving of linear system in backward pass.
            z_root_bwd = new_z_root.clone().detach().requires_grad_()

            if kwargs.get('solver_fwd_jac_f') is not None:
                def backward_hook(grad):
                    return torch.solve(kwargs['solver_fwd_jac_f'](z_root_bwd), grad)[0]
            else:
                f_bwd = -self.f(z_root_bwd, x, *args, **remove_kwargs(kwargs, 'solver_'))

                def backward_hook(grad):
                    return self.solver(
                        lambda y: autograd.grad(f_bwd, z_root_bwd, y, retain_graph=True, create_graph=True)[0] + grad,
                        torch.zeros_like(grad), **filter_kwargs(kwargs, 'solver_bwd_')
                    )['result']

            new_z_root.register_hook(backward_hook)

        return new_z_root
//  [python]
def _root_find(self, z0, x, *args, **kwargs):

        # Compute forward pass: find root of function outside autograd tape.
        with torch.no_grad():
            z_root = self.solver(
                lambda z: self.f(z, x, *args, **remove_kwargs(kwargs, 'solver_')),
                z0,
                **filter_kwargs(kwargs, 'solver_fwd_'),
            )['result']
            new_z_root = z_root

        if self.training:
            # Re-engage autograd tape (no-op in terms of value of z).
            new_z_root = z_root.requires_grad_() - self.f(z_root.requires_grad_(), x, *args, **remove_kwargs(kwargs, 'solver_'))

            # Set up backward hook for solving of linear system in backward pass.
            z_root_bwd = new_z_root.clone().detach().requires_grad_()

            if kwargs.get('solver_fwd_jac_f') is not None:
                def backward_hook(grad):
                    return torch.linalg.solve(kwargs['solver_fwd_jac_f'](z_root_bwd), grad)
            else:
                f_bwd = -self.f(z_root_bwd, x, *args, **remove_kwargs(kwargs, 'solver_'))

                def backward_hook(grad):
                    return self.solver(
                        lambda y: autograd.grad(f_bwd, z_root_bwd, y, retain_graph=True, create_graph=True)[0] + grad,
                        torch.zeros_like(grad), **filter_kwargs(kwargs, 'solver_bwd_')
                    )['result']

            new_z_root.register_hook(backward_hook)

        return new_z_root
//  [python]
def solve(a: ArrayLike, b: ArrayLike):
    a, b = _atleast_float_2(a, b)
    return torch.solve(b, a)[0]
//  [python]
def solve(a: ArrayLike, b: ArrayLike):
    a, b = _atleast_float_2(a, b)
    return torch.linalg.solve(a, b)
//  [python]
def matrix_pade_approximant(p, I):
    p_sqrt = pade_p[0] * I
    q_sqrt = pade_q[0] * I
    p_app = I - p
    p_hat = p_app
    for i in range(5):
        p_sqrt += pade_p[i + 1] * p_hat
        q_sqrt += pade_q[i + 1] * p_hat
        p_hat = p_hat.bmm(p_app)
    # There are 4 options to compute the MPA: comput Matrix Inverse or Matrix Linear System on CPU/GPU;
    # It seems that single matrix is faster on CPU and batched matrices are faster on GPU
    # Please check which one is faster befores running the code;
    # out=torch.linalg.solve(q_sqrt.cpu(), p_sqrt.cpu()).cuda()
    # return out

    # return torch.linalg.solve(q_sqrt, p_sqrt)
    return torch.solve(p_sqrt.cpu(), q_sqrt.cpu())[0].cuda()
//  [python]
def matrix_pade_approximant(p, I):
    p_sqrt = pade_p[0] * I
    q_sqrt = pade_q[0] * I
    p_app = I - p
    p_hat = p_app
    for i in range(5):
        p_sqrt += pade_p[i + 1] * p_hat
        q_sqrt += pade_q[i + 1] * p_hat
        p_hat = p_hat.bmm(p_app)
    # There are 4 options to compute the MPA: comput Matrix Inverse or Matrix Linear System on CPU/GPU;
    # It seems that single matrix is faster on CPU and batched matrices are faster on GPU
    # Please check which one is faster befores running the code;
    # out=torch.linalg.solve(q_sqrt.cpu(), p_sqrt.cpu()).cuda()
    # return out

    # return torch.linalg.solve(q_sqrt, p_sqrt)
    return torch.linalg.solve(q_sqrt.cpu(), p_sqrt.cpu()).cuda()
//  [python]
def matrix_pade_approximant_inverse(p, I):
    p_sqrt = pade_p[0] * I
    q_sqrt = pade_q[0] * I
    p_app = I - p
    p_hat = p_app
    for i in range(5):
        p_sqrt += pade_p[i + 1] * p_hat
        q_sqrt += pade_q[i + 1] * p_hat
        p_hat = p_hat.bmm(p_app)
    # There are 4 options to compute the MPA_inverse: comput Matrix Inverse or Matrix Linear System on CPU/GPU;
    # It seems that single matrix is faster on CPU and batched matrices are faster on GPU
    # Please check which one is faster before running the code;
    # return torch.linalg.solve(p_sqrt, q_sqrt)
    return torch.solve(q_sqrt.cpu(), p_sqrt.cpu()).solution.cuda()
//  [python]
def matrix_pade_approximant_inverse(p, I):
    p_sqrt = pade_p[0] * I
    q_sqrt = pade_q[0] * I
    p_app = I - p
    p_hat = p_app
    for i in range(5):
        p_sqrt += pade_p[i + 1] * p_hat
        q_sqrt += pade_q[i + 1] * p_hat
        p_hat = p_hat.bmm(p_app)
    # There are 4 options to compute the MPA_inverse: comput Matrix Inverse or Matrix Linear System on CPU/GPU;
    # It seems that single matrix is faster on CPU and batched matrices are faster on GPU
    # Please check which one is faster before running the code;
    # return torch.linalg.solve(p_sqrt, q_sqrt)
    return torch.linalg.solve(p_sqrt.cpu(), q_sqrt.cpu()).cuda()
//  [python]
def solve(self, a, b):
        return torch.solve(b, a)[0]
//  [python]
def solve(self, a, b):
        return torch.linalg.solve(a, b)
//  [python]
def backward(ctx, grad_p):
        p, a, b = ctx.saved_tensors

        device = grad_p.device

        m, n = p.shape[-2:]
        batch_shape = list(p.shape[:-2])

        grad_p *= -1 / ctx.lambd_sink * p
        K = torch.cat(
            (
                torch.cat((torch.diag_embed(a), p), dim=-1),
                torch.cat((p.transpose(-2, -1), torch.diag_embed(b)), dim=-1),
            ),
            dim=-2,
        )[..., :-1, :-1]
        t = torch.cat(
            (grad_p.sum(dim=-1), grad_p[..., :, :-1].sum(dim=-2)), dim=-1
        ).unsqueeze(-1)
        grad_ab = torch.solve(t, K)[0]
        grad_a = grad_ab[..., :m, :]
        grad_b = torch.cat(
            (
                grad_ab[..., m:, :],
                torch.zeros(batch_shape + [1, 1], device=device, dtype=torch.float32),
            ),
            dim=-2,
        )
        U = grad_a + grad_b.transpose(-2, -1)
        grad_p -= p * U
        grad_a = -ctx.lambd_sink * grad_a.squeeze(dim=-1)
        grad_b = -ctx.lambd_sink * grad_b.squeeze(dim=-1)
        return grad_p, grad_a, grad_b, None, None, None
//  [python]
def backward(ctx, grad_p):
        p, a, b = ctx.saved_tensors

        device = grad_p.device

        m, n = p.shape[-2:]
        batch_shape = list(p.shape[:-2])

        grad_p *= -1 / ctx.lambd_sink * p
        K = torch.cat(
            (
                torch.cat((torch.diag_embed(a), p), dim=-1),
                torch.cat((p.transpose(-2, -1), torch.diag_embed(b)), dim=-1),
            ),
            dim=-2,
        )[..., :-1, :-1]
        t = torch.cat(
            (grad_p.sum(dim=-1), grad_p[..., :, :-1].sum(dim=-2)), dim=-1
        ).unsqueeze(-1)
        grad_ab = torch.linalg.solve(K, t)
        grad_a = grad_ab[..., :m, :]
        grad_b = torch.cat(
            (
                grad_ab[..., m:, :],
                torch.zeros(batch_shape + [1, 1], device=device, dtype=torch.float32),
            ),
            dim=-2,
        )
        U = grad_a + grad_b.transpose(-2, -1)
        grad_p -= p * U
        grad_a = -ctx.lambd_sink * grad_a.squeeze(dim=-1)
        grad_b = -ctx.lambd_sink * grad_b.squeeze(dim=-1)
        return grad_p, grad_a, grad_b, None, None, None
//  [python]
def bilinear(dt, A, B=None):
    """
    dt: (...) timescales
    A: (... N N)
    B: (... N)
    """
    N = A.shape[-1]
    I = torch.eye(N).to(A)
    A_backwards = I - dt[:, None, None] / 2 * A
    A_forwards = I + dt[:, None, None] / 2 * A

    if B is None:
        dB = None
    else:
        dB = dt[..., None] * torch.solve(B.unsqueeze(-1), A_backwards)[0].squeeze(
            -1
        )  # (... N)

    dA = torch.solve(A_forwards, A_backwards)[0]  # (... N N)
    return dA, dB
//  [python]
def bilinear(dt, A, B=None):
    """
    dt: (...) timescales
    A: (... N N)
    B: (... N)
    """
    N = A.shape[-1]
    I = torch.eye(N).to(A)
    A_backwards = I - dt[:, None, None] / 2 * A
    A_forwards = I + dt[:, None, None] / 2 * A

    if B is None:
        dB = None
    else:
        dB = dt[..., None] * torch.linalg.solve(A_backwards, B.unsqueeze(-1)).squeeze(
            -1
        )  # (... N)

    dA = torch.linalg.solve(A_backwards, A_forwards)  # (... N N)
    return dA, dB
//  [python]
def test_root_inv_decomposition(self):
        linear_op = self.create_linear_op()
        root_approx = linear_operator.root_inv_decomposition(linear_op)

        test_mat = torch.randn(*linear_op.batch_shape, linear_op.size(-1), 5)

        res = root_approx.matmul(test_mat)
        actual = torch.solve(test_mat, linear_op)[0]
        self.assertAllClose(res, actual, **self.tolerances["root_inv_decomposition"])
//  [python]
def test_root_inv_decomposition(self):
        linear_op = self.create_linear_op()
        root_approx = linear_operator.root_inv_decomposition(linear_op)

        test_mat = torch.randn(*linear_op.batch_shape, linear_op.size(-1), 5)

        res = root_approx.matmul(test_mat)
        actual = torch.linalg.solve(linear_op, test_mat)
        self.assertAllClose(res, actual, **self.tolerances["root_inv_decomposition"])
//  [python]
def test_solve(real_vec):
    if is_torch_1_9_plus:
        wrappers = [ComplexTensor, torch.complex]
        modules = [FC, torch]
    else:
        wrappers = [ComplexTensor]
        modules = [FC]

    for complex_wrapper, complex_module in zip(wrappers, modules):
        mat = complex_wrapper(
            torch.from_numpy(mat_np.real), torch.from_numpy(mat_np.imag)
        )
        if not real_vec or complex_wrapper is ComplexTensor:
            vec = complex_wrapper(torch.rand(2, 3, 1), torch.rand(2, 3, 1))
            vec2 = vec
        else:
            vec = torch.rand(2, 3, 1)
            vec2 = complex_wrapper(vec, torch.zeros_like(vec))
        ret = solve(vec, mat)
        if isinstance(vec2, ComplexTensor):
            ret2 = FC.solve(vec2, mat, return_LU=False)
        else:
            return torch.solve(vec2, mat)[0]
        assert complex_module.allclose(ret, ret2)
//  [python]
def test_solve(real_vec):
    if is_torch_1_9_plus:
        wrappers = [ComplexTensor, torch.complex]
        modules = [FC, torch]
    else:
        wrappers = [ComplexTensor]
        modules = [FC]

    for complex_wrapper, complex_module in zip(wrappers, modules):
        mat = complex_wrapper(
            torch.from_numpy(mat_np.real), torch.from_numpy(mat_np.imag)
        )
        if not real_vec or complex_wrapper is ComplexTensor:
            vec = complex_wrapper(torch.rand(2, 3, 1), torch.rand(2, 3, 1))
            vec2 = vec
        else:
            vec = torch.rand(2, 3, 1)
            vec2 = complex_wrapper(vec, torch.zeros_like(vec))
        ret = solve(vec, mat)
        if isinstance(vec2, ComplexTensor):
            ret2 = FC.solve(vec2, mat, return_LU=False)
        else:
            return torch.linalg.solve(mat, vec2)
        assert complex_module.allclose(ret, ret2)
//  [python]
def test_add_low_rank(self):
        linear_op = self.create_linear_op()
        linear_op = self.create_linear_op()
        evaluated = self.evaluate_linear_op(linear_op)
        new_rows = torch.randn(*linear_op.shape[:-1], 3)

        summed_lt = evaluated + new_rows.matmul(new_rows.mT)
        new_lt = linear_op.add_low_rank(new_rows)

        # check that the concatenation is okay
        self.assertAllClose(new_lt.to_dense(), summed_lt)

        # check that the root approximation is close
        rhs = torch.randn(linear_op.size(-1))
        summed_rhs = summed_lt.matmul(rhs)
        root_rhs = linear_operator.root_decomposition(new_lt).matmul(rhs)
        self.assertAllClose(root_rhs, summed_rhs, **self.tolerances["root_decomposition"])

        # check that the inverse root decomposition is close
        summed_solve = torch.solve(summed_lt, rhs.unsqueeze(-1)).solution.squeeze(-1)
        root_inv_solve = linear_operator.root_inv_decomposition(new_lt).matmul(rhs)
        self.assertAllClose(root_inv_solve, summed_solve, **self.tolerances["root_inv_decomposition"])
//  [python]
def test_add_low_rank(self):
        linear_op = self.create_linear_op()
        linear_op = self.create_linear_op()
        evaluated = self.evaluate_linear_op(linear_op)
        new_rows = torch.randn(*linear_op.shape[:-1], 3)

        summed_lt = evaluated + new_rows.matmul(new_rows.mT)
        new_lt = linear_op.add_low_rank(new_rows)

        # check that the concatenation is okay
        self.assertAllClose(new_lt.to_dense(), summed_lt)

        # check that the root approximation is close
        rhs = torch.randn(linear_op.size(-1))
        summed_rhs = summed_lt.matmul(rhs)
        root_rhs = linear_operator.root_decomposition(new_lt).matmul(rhs)
        self.assertAllClose(root_rhs, summed_rhs, **self.tolerances["root_decomposition"])

        # check that the inverse root decomposition is close
        summed_solve = torch.linalg.solve(summed_lt, rhs.unsqueeze(-1)).squeeze(-1)
        root_inv_solve = linear_operator.root_inv_decomposition(new_lt).matmul(rhs)
        self.assertAllClose(root_inv_solve, summed_solve, **self.tolerances["root_inv_decomposition"])
//  [python]
def get_result_of_compare_method(h, p) -> TEN:  # MMSE beamformer
        h = h[0] + 1j * h[1]
        k, n = h.shape
        eye_mat = th.eye(n, dtype=h.dtype, device=h.device)
        w = th.solve(eye_mat * k / p + th.conj(th.transpose(h, 0, 1)) @ h, th.conj(th.transpose(h, 0, 1)))[0]
        w = w / (th.norm(w, dim=0, keepdim=True) * k ** 0.5)  # w.shape == [K, N]
        return th.stack((w.real, w.imag), dim=0)
//  [python]
def get_result_of_compare_method(h, p) -> TEN:  # MMSE beamformer
        h = h[0] + 1j * h[1]
        k, n = h.shape
        eye_mat = th.eye(n, dtype=h.dtype, device=h.device)
        w = th.linalg.solve(eye_mat * k / p + th.conj(th.transpose(h, 0, 1)) @ h, th.conj(th.transpose(h, 0, 1)))
        w = w / (th.norm(w, dim=0, keepdim=True) * k ** 0.5)  # w.shape == [K, N]
        return th.stack((w.real, w.imag), dim=0)
//  [python]
def forward(self, matA):
        assert matA.ndim >= 2
        shape = matA.shape
        N0 = shape[-1]
        matA = matA.reshape(-1, N0, N0)
        torch1 = _PSDMatrixSqrtmRepeat.apply(matA, self.num_sqrtm)
        eye0 = torch.eye(N0, device=matA.device)
        # ret = sum(torch.linalg.solve((1-b)*eye0+b*torch1, a*torch1-a*eye0) for a,b in zip(self.alpha,self.beta))
        tmp0 = (1-self.beta)*eye0+self.beta*torch1
        tmp1 = self.alpha*torch1 - self.alpha*eye0
        ret = torch.solve(tmp1, tmp0)[0].sum(dim=0).reshape(*shape)
        return ret
//  [python]
def forward(self, matA):
        assert matA.ndim >= 2
        shape = matA.shape
        N0 = shape[-1]
        matA = matA.reshape(-1, N0, N0)
        torch1 = _PSDMatrixSqrtmRepeat.apply(matA, self.num_sqrtm)
        eye0 = torch.eye(N0, device=matA.device)
        # ret = sum(torch.linalg.solve((1-b)*eye0+b*torch1, a*torch1-a*eye0) for a,b in zip(self.alpha,self.beta))
        tmp0 = (1-self.beta)*eye0+self.beta*torch1
        tmp1 = self.alpha*torch1 - self.alpha*eye0
        ret = torch.linalg.solve(tmp0, tmp1).sum(dim=0).reshape(*shape)
        return ret
//  [python]
def _test_solve(self, rhs, lhs=None, cholesky=False):
        linear_op = self.create_linear_op().detach().requires_grad_(True)
        linear_op_copy = torch.clone(linear_op).detach().requires_grad_(True)
        evaluated = self.evaluate_linear_op(linear_op_copy)
        evaluated.register_hook(self._ensure_symmetric_grad)

        # Create a test right hand side and left hand side
        rhs.requires_grad_(True)
        rhs_copy = rhs.clone().detach().requires_grad_(True)
        if lhs is not None:
            lhs.requires_grad_(True)
            lhs_copy = lhs.clone().detach().requires_grad_(True)

        _wrapped_cg = MagicMock(wraps=linear_operator.utils.linear_cg)
        with patch("linear_operator.utils.linear_cg", new=_wrapped_cg) as linear_cg_mock:
            with linear_operator.settings.max_cholesky_size(
                math.inf if cholesky else 0
            ), linear_operator.settings.cg_tolerance(1e-4):
                # Perform the solve
                if lhs is not None:
                    res = linear_operator.solve(linear_op, rhs, lhs)
                    actual = lhs_copy @ evaluated.inverse() @ rhs_copy
                else:
                    res = torch.solve(linear_op, rhs)[0]
                    actual = evaluated.inverse().matmul(rhs_copy)
                self.assertAllClose(res, actual, **self.tolerances["solve"])

                # Perform backward pass
                grad = torch.randn_like(res)
                res.backward(gradient=grad)
                actual.backward(gradient=grad)
                for arg, arg_copy in zip(linear_op.representation(), linear_op_copy.representation()):
                    if arg_copy.requires_grad and arg_copy.is_leaf and arg_copy.grad is not None:
                        self.assertAllClose(arg.grad, arg_copy.grad, **self.tolerances["grad"])
                self.assertAllClose(rhs.grad, rhs_copy.grad, **self.tolerances["grad"])
                if lhs is not None:
                    self.assertAllClose(lhs.grad, lhs_copy.grad, **self.tolerances["grad"])

            # Determine if we've called CG or not
            if not cholesky and self.__class__.should_call_cg:
                self.assertTrue(linear_cg_mock.called)
            else:
                self.assertFalse(linear_cg_mock.called)
//  [python]
def _test_solve(self, rhs, lhs=None, cholesky=False):
        linear_op = self.create_linear_op().detach().requires_grad_(True)
        linear_op_copy = torch.clone(linear_op).detach().requires_grad_(True)
        evaluated = self.evaluate_linear_op(linear_op_copy)
        evaluated.register_hook(self._ensure_symmetric_grad)

        # Create a test right hand side and left hand side
        rhs.requires_grad_(True)
        rhs_copy = rhs.clone().detach().requires_grad_(True)
        if lhs is not None:
            lhs.requires_grad_(True)
            lhs_copy = lhs.clone().detach().requires_grad_(True)

        _wrapped_cg = MagicMock(wraps=linear_operator.utils.linear_cg)
        with patch("linear_operator.utils.linear_cg", new=_wrapped_cg) as linear_cg_mock:
            with linear_operator.settings.max_cholesky_size(
                math.inf if cholesky else 0
            ), linear_operator.settings.cg_tolerance(1e-4):
                # Perform the solve
                if lhs is not None:
                    res = linear_operator.solve(linear_op, rhs, lhs)
                    actual = lhs_copy @ evaluated.inverse() @ rhs_copy
                else:
                    res = torch.linalg.solve(linear_op, rhs)
                    actual = evaluated.inverse().matmul(rhs_copy)
                self.assertAllClose(res, actual, **self.tolerances["solve"])

                # Perform backward pass
                grad = torch.randn_like(res)
                res.backward(gradient=grad)
                actual.backward(gradient=grad)
                for arg, arg_copy in zip(linear_op.representation(), linear_op_copy.representation()):
                    if arg_copy.requires_grad and arg_copy.is_leaf and arg_copy.grad is not None:
                        self.assertAllClose(arg.grad, arg_copy.grad, **self.tolerances["grad"])
                self.assertAllClose(rhs.grad, rhs_copy.grad, **self.tolerances["grad"])
                if lhs is not None:
                    self.assertAllClose(lhs.grad, lhs_copy.grad, **self.tolerances["grad"])

            # Determine if we've called CG or not
            if not cholesky and self.__class__.should_call_cg:
                self.assertTrue(linear_cg_mock.called)
            else:
                self.assertFalse(linear_cg_mock.called)
//  [python]
def log_likelihood(self, x: torch.Tensor, u: torch.Tensor, extra_info=None, theta=None, mu_r=0, lam_r=0):
        if theta is None:
            theta = torch.cat([torch.tensor(self.gt_theta), torch.tensor([-np.inf])])
        
        g_hat = -2 * x
        g_tilde = -2 * u

        g = theta[1] * g_tilde.unsqueeze(-1) + self.J.unsqueeze(0) @ (theta[0] * g_hat.unsqueeze(-1))
        H = self.hessian(x, u, theta)

        ll_val = 0.5 * torch.transpose(g, -1, -2) @ torch.solve(g, H)[0] + 0.5 * torch.logdet(-H)
        if np.isinf(mu_r):
            # We've switched off the constrained variable (auglagcost.m:17)
            return ll_val.mean()
        else:
            return ll_val.mean() - 0.5*mu_r*(torch.exp(theta[-1])**2) + lam_r*torch.exp(theta[-1])
//  [python]
def log_likelihood(self, x: torch.Tensor, u: torch.Tensor, extra_info=None, theta=None, mu_r=0, lam_r=0):
        if theta is None:
            theta = torch.cat([torch.tensor(self.gt_theta), torch.tensor([-np.inf])])
        
        g_hat = -2 * x
        g_tilde = -2 * u

        g = theta[1] * g_tilde.unsqueeze(-1) + self.J.unsqueeze(0) @ (theta[0] * g_hat.unsqueeze(-1))
        H = self.hessian(x, u, theta)

        ll_val = 0.5 * torch.transpose(g, -1, -2) @ torch.linalg.solve(H, g) + 0.5 * torch.logdet(-H)
        if np.isinf(mu_r):
            # We've switched off the constrained variable (auglagcost.m:17)
            return ll_val.mean()
        else:
            return ll_val.mean() - 0.5*mu_r*(torch.exp(theta[-1])**2) + lam_r*torch.exp(theta[-1])
//  [python]
def matrix_pade_approximant(p,I):
    p_sqrt = pade_p[0]*I
    q_sqrt = pade_q[0]*I
    p_app = I - p
    p_hat = p_app
    for i in range(5):
        p_sqrt += pade_p[i+1]*p_hat
        q_sqrt += pade_q[i+1]*p_hat
        p_hat = p_hat.bmm(p_app)
    #There are 4 options to compute the MPA: comput Matrix Inverse or Matrix Linear System on CPU/GPU;
    #It seems that single matrix is faster on CPU and batched matrices are faster on GPU
    #Please check which one is faster before running the code;
    return torch.solve(p_sqrt, q_sqrt)[0]
//  [python]
def matrix_pade_approximant(p,I):
    p_sqrt = pade_p[0]*I
    q_sqrt = pade_q[0]*I
    p_app = I - p
    p_hat = p_app
    for i in range(5):
        p_sqrt += pade_p[i+1]*p_hat
        q_sqrt += pade_q[i+1]*p_hat
        p_hat = p_hat.bmm(p_app)
    #There are 4 options to compute the MPA: comput Matrix Inverse or Matrix Linear System on CPU/GPU;
    #It seems that single matrix is faster on CPU and batched matrices are faster on GPU
    #Please check which one is faster before running the code;
    return torch.linalg.solve(q_sqrt, p_sqrt)
//  [python]
def test_cat_rows(self):
        linear_op = self.create_linear_op()
        evaluated = self.evaluate_linear_op(linear_op)

        for batch_shape in (torch.Size(), torch.Size([2])):
            new_rows = 1e-4 * torch.randn(*batch_shape, *linear_op.shape[:-2], 1, linear_op.shape[-1])
            new_point = torch.rand(*batch_shape, *linear_op.shape[:-2], 1, 1)

            # we need to expand here to be able to concat (this happens automatically in cat_rows)
            cat_col1 = torch.cat((evaluated.expand(*batch_shape, *evaluated.shape), new_rows), dim=-2)
            cat_col2 = torch.cat((new_rows.mT, new_point), dim=-2)

            concatenated_lt = torch.cat((cat_col1, cat_col2), dim=-1)
            try:
                new_lt = linear_op.cat_rows(new_rows, new_point)
            except Exception:
                msg = traceback.format_exc()
                logging.warning(msg)

            # check that the concatenation is okay
            self.assertAllClose(new_lt.to_dense(), concatenated_lt)

            # check that the root approximation is close
            rhs = torch.randn(linear_op.size(-1) + 1)
            concat_rhs = concatenated_lt.matmul(rhs)
            root_rhs = linear_operator.root_decomposition(new_lt).matmul(rhs)
            self.assertAllClose(root_rhs, concat_rhs, **self.tolerances["root_decomposition"])

            # check that root inv is cached
            root_inv = get_from_cache(new_lt, "root_inv_decomposition")
            # check that the inverse root decomposition is close
            concat_solve = torch.solve(rhs.unsqueeze(-1), concatenated_lt)[0].squeeze(-1)
            root_inv_solve = root_inv.matmul(rhs)
            self.assertLess(
                (root_inv_solve - concat_solve).norm() / concat_solve.norm(),
                self.tolerances["root_inv_decomposition"]["rtol"],
            )
            # test generate_inv_roots=False
            new_lt = linear_op.cat_rows(new_rows, new_point, generate_inv_roots=False)
            with self.assertRaises(CachingError):
                get_from_cache(new_lt, "root_inv_decomposition")
//  [python]
def test_cat_rows(self):
        linear_op = self.create_linear_op()
        evaluated = self.evaluate_linear_op(linear_op)

        for batch_shape in (torch.Size(), torch.Size([2])):
            new_rows = 1e-4 * torch.randn(*batch_shape, *linear_op.shape[:-2], 1, linear_op.shape[-1])
            new_point = torch.rand(*batch_shape, *linear_op.shape[:-2], 1, 1)

            # we need to expand here to be able to concat (this happens automatically in cat_rows)
            cat_col1 = torch.cat((evaluated.expand(*batch_shape, *evaluated.shape), new_rows), dim=-2)
            cat_col2 = torch.cat((new_rows.mT, new_point), dim=-2)

            concatenated_lt = torch.cat((cat_col1, cat_col2), dim=-1)
            try:
                new_lt = linear_op.cat_rows(new_rows, new_point)
            except Exception:
                msg = traceback.format_exc()
                logging.warning(msg)

            # check that the concatenation is okay
            self.assertAllClose(new_lt.to_dense(), concatenated_lt)

            # check that the root approximation is close
            rhs = torch.randn(linear_op.size(-1) + 1)
            concat_rhs = concatenated_lt.matmul(rhs)
            root_rhs = linear_operator.root_decomposition(new_lt).matmul(rhs)
            self.assertAllClose(root_rhs, concat_rhs, **self.tolerances["root_decomposition"])

            # check that root inv is cached
            root_inv = get_from_cache(new_lt, "root_inv_decomposition")
            # check that the inverse root decomposition is close
            concat_solve = torch.linalg.solve(concatenated_lt, rhs.unsqueeze(-1)).squeeze(-1)
            root_inv_solve = root_inv.matmul(rhs)
            self.assertLess(
                (root_inv_solve - concat_solve).norm() / concat_solve.norm(),
                self.tolerances["root_inv_decomposition"]["rtol"],
            )
            # test generate_inv_roots=False
            new_lt = linear_op.cat_rows(new_rows, new_point, generate_inv_roots=False)
            with self.assertRaises(CachingError):
                get_from_cache(new_lt, "root_inv_decomposition")
//  [python]
def solve(a, b):
    a = convert_to_tensor(a)
    b = convert_to_tensor(b)
    return torch.solve(b, a)[0]
//  [python]
def solve(a, b):
    a = convert_to_tensor(a)
    b = convert_to_tensor(b)
    return torch.linalg.solve(a, b)
//  [python]
def bilinear(cls, dt, A, B=None, separate=False):
        """
        dt: (...) timescales
        A: (... N N)
        B: (... N)
        """
        N = A.shape[-1]
        I = torch.eye(N).to(A)
        A_backwards = I - dt[:, None, None] / 2 * A
        A_forwards = I + dt[:, None, None] / 2 * A

        if B is None:
            dB = None
        else:
            dB = dt[..., None] * torch.solve(
                A_backwards, B.unsqueeze(-1)
            ).solution.squeeze(
                -1
            )  # (... N)

        if separate:
            A_b = torch.solve(A_backwards, I).solution  # (... N N)
            return A_forwards, A_b, dB
        else:
            dA = torch.solve(A_backwards, A_forwards).solution  # (... N N)
            return dA, dB
//  [python]
def bilinear(cls, dt, A, B=None, separate=False):
        """
        dt: (...) timescales
        A: (... N N)
        B: (... N)
        """
        N = A.shape[-1]
        I = torch.eye(N).to(A)
        A_backwards = I - dt[:, None, None] / 2 * A
        A_forwards = I + dt[:, None, None] / 2 * A

        if B is None:
            dB = None
        else:
            dB = dt[..., None] * torch.linalg.solve(
                A_backwards, B.unsqueeze(-1)
            ).squeeze(
                -1
            )  # (... N)

        if separate:
            A_b = torch.linalg.solve(A_backwards, I)  # (... N N)
            return A_forwards, A_b, dB
        else:
            dA = torch.linalg.solve(A_backwards, A_forwards)  # (... N N)
            return dA, dB
//  [python]
def solve(
    x1: torch.Tensor,
    x2: torch.Tensor,
    /,
    *,
    adjoint: bool = False,
    out: Optional[torch.Tensor] = None,
) -> torch.Tensor:
    if adjoint:
        x1 = torch.adjoint(x1)
    x1, x2 = ivy.promote_types_of_inputs(x1, x2)
    expanded_last = False
    if len(x2.shape) <= 1:
        if x2.shape[-1] == x1.shape[-1]:
            expanded_last = True
            x2 = torch.unsqueeze(x2, dim=1)

    is_empty_x1 = x1.nelement() == 0
    is_empty_x2 = x2.nelement() == 0
    if is_empty_x1 or is_empty_x2:
        for i in range(len(x1.shape) - 2):
            x2 = torch.unsqueeze(x2, dim=0)
        output_shape = list(torch.broadcast_shapes(x1.shape[:-2], x2.shape[:-2]))
        output_shape.append(x2.shape[-2])
        output_shape.append(x2.shape[-1])
        ret = torch.Tensor([])
        ret = torch.reshape(ret, output_shape)
    else:
        ret = torch.solve(x2, x1)[0]

    if expanded_last:
        ret = torch.squeeze(ret, dim=-1)
    return ret
//  [python]
def solve(
    x1: torch.Tensor,
    x2: torch.Tensor,
    /,
    *,
    adjoint: bool = False,
    out: Optional[torch.Tensor] = None,
) -> torch.Tensor:
    if adjoint:
        x1 = torch.adjoint(x1)
    x1, x2 = ivy.promote_types_of_inputs(x1, x2)
    expanded_last = False
    if len(x2.shape) <= 1:
        if x2.shape[-1] == x1.shape[-1]:
            expanded_last = True
            x2 = torch.unsqueeze(x2, dim=1)

    is_empty_x1 = x1.nelement() == 0
    is_empty_x2 = x2.nelement() == 0
    if is_empty_x1 or is_empty_x2:
        for i in range(len(x1.shape) - 2):
            x2 = torch.unsqueeze(x2, dim=0)
        output_shape = list(torch.broadcast_shapes(x1.shape[:-2], x2.shape[:-2]))
        output_shape.append(x2.shape[-2])
        output_shape.append(x2.shape[-1])
        ret = torch.Tensor([])
        ret = torch.reshape(ret, output_shape)
    else:
        ret = torch.linalg.solve(x1, x2)

    if expanded_last:
        ret = torch.squeeze(ret, dim=-1)
    return ret
//  [python]
def proj_tangent(x, u):
    assert x.shape[-2:] == u.shape[-2:], "Wrong shapes"
    x, u = torch.broadcast_tensors(x, u)
    x_shape = x.shape
    x = x.reshape(-1, x_shape[-2], x_shape[-1])
    u = u.reshape(-1, x_shape[-2], x_shape[-1])
    xt = x.transpose(-1, -2)
    batch_size, n = x.shape[0:2]

    I = torch.eye(n, dtype=x.dtype, device=x.device)
    I = I.expand_as(x)

    mu = x * u

    A = linalg.block_matrix([[I, x], [xt, I]])

    B = A[:, :, 1:]

    z1 = mu.sum(dim=-1).unsqueeze(-1)
    zt1 = mu.sum(dim=-2).unsqueeze(-1)

    b = torch.cat(
        [z1, zt1],
        dim=1,
    )
    rhs = B.transpose(1, 2) @ (b - A[:, :, 0:1])
    lhs = B.transpose(1, 2) @ B
    zeta = torch.solve(rhs, lhs)[0]
    alpha = torch.cat(
        [
            torch.ones(batch_size, 1, 1, device=x.device, dtype=x.dtype),
            zeta[:, 0 : n - 1],
        ],
        dim=1,
    )
    beta = zeta[:, n - 1 : 2 * n - 1]
    rgrad = mu - (alpha + beta.transpose(-1, -2)) * x

    rgrad = rgrad.reshape(x_shape)
    return rgrad
//  [python]
def proj_tangent(x, u):
    assert x.shape[-2:] == u.shape[-2:], "Wrong shapes"
    x, u = torch.broadcast_tensors(x, u)
    x_shape = x.shape
    x = x.reshape(-1, x_shape[-2], x_shape[-1])
    u = u.reshape(-1, x_shape[-2], x_shape[-1])
    xt = x.transpose(-1, -2)
    batch_size, n = x.shape[0:2]

    I = torch.eye(n, dtype=x.dtype, device=x.device)
    I = I.expand_as(x)

    mu = x * u

    A = linalg.block_matrix([[I, x], [xt, I]])

    B = A[:, :, 1:]

    z1 = mu.sum(dim=-1).unsqueeze(-1)
    zt1 = mu.sum(dim=-2).unsqueeze(-1)

    b = torch.cat(
        [z1, zt1],
        dim=1,
    )
    rhs = B.transpose(1, 2) @ (b - A[:, :, 0:1])
    lhs = B.transpose(1, 2) @ B
    zeta = torch.linalg.solve(lhs, rhs)
    alpha = torch.cat(
        [
            torch.ones(batch_size, 1, 1, device=x.device, dtype=x.dtype),
            zeta[:, 0 : n - 1],
        ],
        dim=1,
    )
    beta = zeta[:, n - 1 : 2 * n - 1]
    rgrad = mu - (alpha + beta.transpose(-1, -2)) * x

    rgrad = rgrad.reshape(x_shape)
    return rgrad
//  [python]
def solve(x1: array, x2: array, /, **kwargs) -> array:
    x1, x2 = _fix_promotion(x1, x2, only_scalar=False)
    return torch.solve(x2, x1, **kwargs)[0]
//  [python]
def solve(x1: array, x2: array, /, **kwargs) -> array:
    x1, x2 = _fix_promotion(x1, x2, only_scalar=False)
    return torch.linalg.solve(x1, x2, **kwargs)
//  [python]
def solve(a: Numeric, b: Numeric):
    return torch.solve(b, a)[0]
//  [python]
def solve(a: Numeric, b: Numeric):
    return torch.linalg.solve(a, b)
//  [python]
def _h(L, grad, hess) -> torch.Tensor:
        I = torch.eye(len(grad), dtype=grad.dtype, device=grad.device)
        D = torch.ones_like(hess) - I
        # Alternate damping scheme
        # (hess + 1e-2 * L**2 * I) * (1 + L**2 * I) ** 2 / (1 + L**2),
        h = torch.solve(
            grad,
            hess * (I + D / (1 + L)) + L * I * (1 + torch.diag(hess)),
        )[0]

        return h
//  [python]
def _h(L, grad, hess) -> torch.Tensor:
        I = torch.eye(len(grad), dtype=grad.dtype, device=grad.device)
        D = torch.ones_like(hess) - I
        # Alternate damping scheme
        # (hess + 1e-2 * L**2 * I) * (1 + L**2 * I) ** 2 / (1 + L**2),
        h = torch.linalg.solve(
            hess * (I + D / (1 + L)) + L * I * (1 + torch.diag(hess)),
            grad,
        )

        return h
//  [python]
def exp(self, x, u, t=1):
        if self.metric == "aff_inv":
            return x @ torch.matrix_exp(t * torch.solve(u, x)[1])
        elif self.metric == "log_euc":
            return torch.matrix_exp(self.log(x) + t * u)
//  [python]
def exp(self, x, u, t=1):
        if self.metric == "aff_inv":
            return x @ torch.matrix_exp(t * torch.linalg.solve(x, u))
        elif self.metric == "log_euc":
            return torch.matrix_exp(self.log(x) + t * u)
//  [python]
def torch_solve_qp_eqcon(Q, p, A, b,):
    #######################################################################
    # Solve a QP in form:
    #   x_star =   argmin_x 0.5*x^TQx + p^Tx
    #             subject to Ax =  b
    # Q:  A (n_batch,n_x,n_x) SPD tensor
    # p:  A (n_batch,n_x,1) tensor.
    # A:  A (n_batch,n_eq, n_x) tensor.
    # b:  A (n_batch,n_eq) tensorr.
    # Returns: x_star:  A (n_batch,n_x,1) tensor
    #######################################################################
    # --- prep:
    n_eq = get_ncon(A, dim=1)
    any_eq = n_eq > 0
    if any_eq:
        n_x = p.shape[1]
        # --- setup and solve system
        rhs = torch.cat((-p, b), 1)
        lhs = torch_qp_eqcon_mat(Q=Q, A=A)
        xv = torch.solve(rhs, lhs)[0]

        # --- unpack solution:
        x = xv[:, :n_x, :]
        nus = xv[:, -n_eq:, :]
        sol = {"x": x, "nus": nus}
    else:
        sol = torch_solve_qp_uncon(Q=Q, p=p)

    return sol
//  [python]
def torch_solve_qp_eqcon(Q, p, A, b,):
    #######################################################################
    # Solve a QP in form:
    #   x_star =   argmin_x 0.5*x^TQx + p^Tx
    #             subject to Ax =  b
    # Q:  A (n_batch,n_x,n_x) SPD tensor
    # p:  A (n_batch,n_x,1) tensor.
    # A:  A (n_batch,n_eq, n_x) tensor.
    # b:  A (n_batch,n_eq) tensorr.
    # Returns: x_star:  A (n_batch,n_x,1) tensor
    #######################################################################
    # --- prep:
    n_eq = get_ncon(A, dim=1)
    any_eq = n_eq > 0
    if any_eq:
        n_x = p.shape[1]
        # --- setup and solve system
        rhs = torch.cat((-p, b), 1)
        lhs = torch_qp_eqcon_mat(Q=Q, A=A)
        xv = torch.linalg.solve(lhs, rhs)

        # --- unpack solution:
        x = xv[:, :n_x, :]
        nus = xv[:, -n_eq:, :]
        sol = {"x": x, "nus": nus}
    else:
        sol = torch_solve_qp_uncon(Q=Q, p=p)

    return sol