def METHOD_NAME(ctx, representation_tree, rhs, lhs, *matrix_args):
    ctx.representation_tree = representation_tree
    ctx.linear_op = ctx.representation_tree(*matrix_args)
    if lhs is not None:
        terms = torch.cat([rhs, lhs.mT], dim=-1)
        solves, weights, no_shift_solves, shifts = utils.contour_integral_quad(
            ctx.linear_op,
            terms,
            inverse=True,
            num_contour_quadrature=settings.num_contour_quadrature.value(),
        )
        rhs_solves, lhs_solves = solves.split([rhs.size(-1), lhs.size(-2)], dim=-1)
        lhs_no_shift_solves = no_shift_solves[..., -lhs.size(-2) :]
        sqrt_inv_matmul_res = lhs @ (rhs_solves * weights).sum(0)
        inv_quad_res = (lhs_no_shift_solves.mT * lhs).sum(dim=-1).mul_(-1)
    else:
        rhs_solves, weights, _, shifts = utils.contour_integral_quad(
            ctx.linear_op,
            rhs,
            inverse=True,
            num_contour_quadrature=settings.num_contour_quadrature.value(),
        )
        sqrt_inv_matmul_res = (rhs_solves * weights).sum(0)
        lhs_solves = None
        lhs_no_shift_solves = None
        inv_quad_res = torch.zeros(ctx.linear_op.batch_shape, dtype=rhs.dtype, device=rhs.device)
    # Save for backwards
    ctx.save_for_backward(rhs, lhs, rhs_solves, lhs_solves, weights, shifts, lhs_no_shift_solves, *matrix_args)
    return sqrt_inv_matmul_res, inv_quad_res