def _while_loop_scan(cond_fun, body_fun, init_val, max_iter):
  def _iter(val):
  def _fun(tup, it):
def _while_loop_python(cond_fun, body_fun, init_val, maxiter):
def _while_loop_lax(cond_fun, body_fun, init_val, maxiter):
  def _cond_fun(_val):
  def _body_fun(_val):
def while_loop(cond_fun, body_fun, init_val, maxiter, unroll=False, jit=False): class OptStep(NamedTuple):
class KKTSolution(NamedTuple):
class Solver(abc.ABC):
class IterativeSolver(Solver):
class StochasticSolver(IterativeSolver):
class LinearOperator(object):
class LineSearchStep(NamedTuple):
class IterativeLineSearch(IterativeSolver):
  def run(self,
          init_params: Any,
          *args,
          **kwargs) -> OptStep:
  def l2_optimality_error(self, params, *args, **kwargs):
  def attribute_names(self):
  def attribute_values(self):
  def _get_loop_options(self):
  def _cond_fun(self, inputs):
  def _body_fun(self, inputs):
  def _make_zero_step(self, init_params, state) -> OptStep:
  def _run(self,
           init_params: Any,
           *args,
           **kwargs) -> OptStep:
  def run(self,
          init_params: Any,
          *args,
          **kwargs) -> OptStep:
def _where(cond, x, y):
  def run_iterator(self,
                   init_params: Any,
                   iterator,
                   *args,
                   **kwargs) -> OptStep:
  def __init__(self, A):
  def shape(self):
  def matvec(self, x):
  def matvec_element(self, x, idx):
  def rmatvec(self, x):
  def rmatvec_element(self, x, idx):
  def update_matvec(self, Ax, delta, idx):
  def update_rmatvec(self, ATx, delta, idx):
  def column_l2_norms(self, squared=False):
  def tree_flatten(self):
  def tree_unflatten(cls, aux_data, children):
  def _make_zero_step(self, init_stepsize, state) -> LineSearchStep:
  def run(self,
          init_stepsize: float,
          params: Any,
          value: Optional[float] = None,
          grad: Optional[Any] = None,
          descent_direction: Optional[Any] = None,
          *args,
          **kwargs) -> LineSearchStep: def root_vjp(optimality_fun: Callable,
             sol: Any,
             args: Tuple,
             cotangent: Any,
             solve: Callable = linear_solve.solve_normal_cg) -> Any:
  def fun_sol(sol):
  def fun_args(*args):
def _jvp_sol(optimality_fun, sol, args, tangent):
def _jvp_args(optimality_fun, sol, args, tangents):
def root_jvp(optimality_fun: Callable,
             sol: Any,
             args: Tuple,
             tangents: Tuple,
             solve:Callable = linear_solve.solve_normal_cg) -> Any:
def _extract_kwargs(kwarg_keys, flat_args):
def _signature_bind(signature, *args, **kwargs):
def _signature_bind_and_match(signature, *args, **kwargs):
  def map_back(out_args):
def _custom_root(solver_fun, optimality_fun, solve, has_aux,
                 reference_signature=None):
  def make_custom_vjp_solver_fun(solver_fun, kwarg_keys):
    def solver_fun_flat(*flat_args):
    def solver_fun_fwd(*flat_args):
    def solver_fun_bwd(tup, cotangent):
  def wrapped_solver_fun(*args, **kwargs):
def custom_root(optimality_fun: Callable,
                has_aux: bool = False,
                solve: Callable = linear_solve.solve_normal_cg,
                reference_signature: Optional[Callable] = None):
  def wrapper(solver_fun):
def custom_fixed_point(fixed_point_fun: Callable,
                       has_aux: bool = False,
                       solve: Callable = linear_solve.solve_normal_cg,
                       reference_signature: Optional[Callable] = None):
  def optimality_fun(params, *args):
def make_kkt_optimality_fun(obj_fun, eq_fun, ineq_fun=None):
  def optimality_fun(params, params_obj, params_eq, params_ineq): def broadcast_pytrees(*trees):
def tree_scalar_mul(scalar, tree_x):
def tree_add_scalar_mul(tree_x, scalar, tree_y):
def _vdot_safe(a, b):
def tree_vdot(tree_x, tree_y):
def tree_dot(tree_x, tree_y):
def tree_sum(tree_x):
def tree_l2_norm(tree_x, squared=False):
def tree_zeros_like(tree_x):
def tree_ones_like(tree_x):
def tree_average(trees, weights):
def tree_gram(a):
def tree_inf_norm(tree_x):
def tree_where(cond, a, b):
def tree_negative(tree):
def tree_reciproqual(tree):
def tree_mean(tree): def tree_add_scalar_mul(tree_x, scalar, tree_y):
  """Compute tree_x + scalar * tree_y."""
  return tree_multimap(lambda x, y: x + scalar * y, tree_x, tree_y) def tree_scalar_mul(scalar, tree_x):
  """Compute scalar * tree_x."""
  return tree_map(lambda x: scalar * x, tree_x) def tree_vdot(tree_x, tree_y):
  """Compute the inner product <tree_x, tree_y>."""
  vdots = tree_multimap(_vdot_safe, tree_x, tree_y)
  return tree_reduce(operator.add, vdots) def tree_negative(tree):
  """Computes elementwise negation -x."""
  return tree_scalar_mul(-1, tree) def tree_l2_norm(tree_x, squared=False):
  """Compute the l2 norm ||tree_x||."""
  squared_tree = tree_map(jnp.square, tree_x)
  sqnorm = tree_sum(squared_tree)
  if squared:
    return sqnorm
  else:
    return jnp.sqrt(sqnorm) def _make_linear_operator(matvec):
  if matvec is None:
    return DenseLinearOperator
  else:
    return functools.partial(FunctionalLinearOperator, matvec)
from typing import Any
from typing import Callable
from typing import NamedTuple
from typing import Optional
from typing import Union
from dataclasses import dataclass
from functools import partial
from jaxopt._src import loop
from jaxopt._src import base
from jaxopt._src import implicit_diff as idf
from jaxopt._src.tree_util import tree_zeros_like, tree_add, tree_sub
from jaxopt._src.tree_util import tree_add_scalar_mul, tree_scalar_mul
from jaxopt._src.tree_util import tree_vdot, tree_negative, tree_l2_norm
from jaxopt._src.linear_operator import _make_linear_operator
import jax
import jax.numpy as jnp
import jaxopt._src.linear_solve as linear_solve
      and a warm start ``init`` as third argument.
      This solver can be inaccurate and run with low precision.
    maxiter: maximum number of iterations (default: 10).
    tol: absolute tolerance for stoping criterion (default: 1e-7).
    verbose: If verbose=1, print error at each iteration.
    implicit_diff: whether to enable implicit diff or autodiff of unrolled
      iterations.
    implicit_diff_solve: the linear system solver to use.
    jit: whether to JIT-compile the optimization loop (default: "auto").
    unroll: whether to unroll the optimization loop (default: "auto")

  References:

    [1] J. H. Wilkinson. Rounding Errors in Algebraic Processes. Prentice Hall, Englewood Cliffs, NJ, 1963.

    [2] Moler, C.B., 1967. Iterative refinement in floating point. Journal of the ACM (JACM), 14(2), pp.316-321.

    [3] https://en.wikipedia.org/wiki/Iterative_refinement.
  """
  matvec_A: Optional[Callable] = None
  matvec_A_bar: Optional[Callable] = None
  solve: Callable = partial(linear_solve.solve_gmres, ridge=1e-6)
  maxiter: int = 10
  tol: float = 1e-7
  verbose: int = 0
  implicit_diff_solve: Optional[Callable] = None
  jit: base.AutoOrBoolean = "auto"
  unroll: base.AutoOrBoolean = "auto"

  def init_state(self,
                 init_params,
                 A: Any,
                 b: Any,
                 A_bar: Any = None):
    return IterativeRefinementState(
      iter_num=jnp.asarray(0),
      error=jnp.asarray(jnp.inf),
      target_residuals=b,
      init=init_params)

  def init_params(self,
                  A: Any,
                  b: Any,
                  A_bar: Any = None):
    return tree_zeros_like(b)

  def update(self,
             params: Any,
             state: IterativeRefinementState,
             A: Any,
             b: Any,
             A_bar: Optional[Any] = None):

    if self._copy_A:
      A_bar = A

    matvec_A = self.matvec_A(A)
    matvec_A_bar = self.matvec_A_bar(A_bar)

    # TODO(lbethune): support preconditioners ?
    # Could it be done by user with partial(solver, M=precond) ?
    residual_sol = self.solve(matvec_A_bar, state.target_residuals, init=state.init)

    params = tree_add(params, residual_sol)

    target_residuals = tree_sub(b, matvec_A(params))
    error = tree_l2_norm(target_residuals)

    state = IterativeRefinementState(
      iter_num=state.iter_num+1,
      error=error,
      target_residuals=target_residuals,
      init=self.init_params(A, b, A_bar))

    return base.OptStep(params, state)

  def run(self,
          init_params,
          A: Any,
          b: Any,
          A_bar: Optional[Any] = None):
    """Runs the iterative refinement.

    Args:
      init_params: init_params for warm start.
      A: params for ``self.matvec_A``.
      b: vector ``b`` in ``Ax=b``.
      A_bar: optional parameters for ``matvec_A_bar``.
    Returns:
      (params, state), ``params = (primal_var, dual_var_eq, dual_var_ineq)``
    """
    if init_params is None:
      init_params = self.init_params(A, b, A_bar)

    return super().run(init_params, A, b, A_bar)

  def optimality_fun(self,
                     params: Any,
                     A: Any,
                     b: Any,
                     A_bar: Optional[Any] = None):
    del A_bar  # unused
    A = self.matvec_A(A)
    return tree_sub(b, A(params))

  def l2_optimality_error(self,
                          params: Any,
                          A: Any,
                          b: Any,
                          A_bar: Optional[Any] = None):
    del A_bar  # unused
    return tree_l2_norm(self.optimality_fun(params, A, b))

  def __post_init__(self):
    self._copy_A = False

    if self.matvec_A_bar is None:
      self.matvec_A_bar = self.matvec_A
      self._copy_A = True


    self.matvec_A = _make_linear_operator(self.matvec_A)