# Path: jaxopt/_src/base.py
# 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:
# 
# Path: jaxopt/_src/linear_solve.py
# def _materialize_array(matvec, shape, dtype=None):
# def _make_ridge_matvec(matvec: Callable, ridge: float = 0.0):
#   def ridge_matvec(v: Any) -> Any:
# def solve_lu(matvec: Callable, b: jnp.ndarray) -> jnp.ndarray:
# def solve_cholesky(matvec: Callable, b: jnp.ndarray) -> jnp.ndarray:
# def solve_cg(matvec: Callable,
#              b: Any,
#              ridge: Optional[float] = None,
#              init: Optional[Any] = None,
#              **kwargs) -> Any:
# def _rmatvec(matvec, x):
# def _normal_matvec(matvec, x):
# def solve_normal_cg(matvec: Callable,
#                     b: Any,
#                     ridge: Optional[float] = None,
#                     init: Optional[Any] = None,
#                     **kwargs) -> Any:
#   def _matvec(x):
# def solve_gmres(matvec: Callable,
#                 b: Any,
#                 ridge: Optional[float] = None,
#                 init: Optional[Any] = None,
#                 tol: float = 1e-5,
#                 **kwargs) -> Any:
# def solve_bicgstab(matvec: Callable,
#                    b: Any,
#                    ridge: Optional[float] = None,
#                    init: Optional[Any] = None,
#                    **kwargs) -> Any:
#     A = _materialize_array(matvec, b.shape, b.dtype)
#     A = _materialize_array(matvec, b.shape, b.dtype)  # 4d array (tensor)
#     A = A.reshape(-1, b.shape[0] * b.shape[1])  # 2d array (matrix)
#     A = _materialize_array(matvec, b.shape)
#     A = _materialize_array(matvec, b.shape)
# 
# Path: jaxopt/tree_util.py



from typing import Any
from typing import Callable
from typing import NamedTuple
from typing import Optional
from dataclasses import dataclass
from jaxopt._src import base
from jaxopt._src import linear_solve
from jaxopt.tree_util import tree_l2_norm, tree_sub
import jax
import jax.numpy as jnp
    # Compute actual values of state variables at init_param
    residual, aux = self._fun_with_aux(init_params, *args, **kwargs)
    matvec = lambda v: self._jtj_op(init_params, v, *args, **kwargs)
    gradient = self._jt_op(init_params, residual, *args, **kwargs)

    return GaussNewtonState(iter_num=jnp.asarray(0),
                            error=jnp.asarray(jnp.inf),
                            residual=residual,
                            delta=init_params,
                            gradient=gradient,
                            aux=aux)

  def update(self,
             params,
             state: NamedTuple,
             *args,
             **kwargs) -> base.OptStep:
    """Performs one iteration of the least-squares solver.

    Args:
      params: pytree containing the parameters.
      state: named tuple containing the solver state.
    Returns:
      (params, state)
    """
    residual, aux = self._fun_with_aux(params, *args, **kwargs)

    matvec = lambda v: self._jtj_op(params, v, *args, **kwargs)
    gradient = self._jt_op(params, residual, *args, **kwargs)


# Path: jaxopt/_src/base.py
# 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:
# 
# Path: jaxopt/_src/linear_solve.py
# def _materialize_array(matvec, shape, dtype=None):
# def _make_ridge_matvec(matvec: Callable, ridge: float = 0.0):
#   def ridge_matvec(v: Any) -> Any:
# def solve_lu(matvec: Callable, b: jnp.ndarray) -> jnp.ndarray:
# def solve_cholesky(matvec: Callable, b: jnp.ndarray) -> jnp.ndarray:
# def solve_cg(matvec: Callable,
#              b: Any,
#              ridge: Optional[float] = None,
#              init: Optional[Any] = None,
#              **kwargs) -> Any:
# def _rmatvec(matvec, x):
# def _normal_matvec(matvec, x):
# def solve_normal_cg(matvec: Callable,
#                     b: Any,
#                     ridge: Optional[float] = None,
#                     init: Optional[Any] = None,
#                     **kwargs) -> Any:
#   def _matvec(x):
# def solve_gmres(matvec: Callable,
#                 b: Any,
#                 ridge: Optional[float] = None,
#                 init: Optional[Any] = None,
#                 tol: float = 1e-5,
#                 **kwargs) -> Any:
# def solve_bicgstab(matvec: Callable,
#                    b: Any,
#                    ridge: Optional[float] = None,
#                    init: Optional[Any] = None,
#                    **kwargs) -> Any:
#     A = _materialize_array(matvec, b.shape, b.dtype)
#     A = _materialize_array(matvec, b.shape, b.dtype)  # 4d array (tensor)
#     A = A.reshape(-1, b.shape[0] * b.shape[1])  # 2d array (matrix)
#     A = _materialize_array(matvec, b.shape)
#     A = _materialize_array(matvec, b.shape)
# 
# Path: jaxopt/tree_util.py



# Path: jaxopt/_src/gauss_newton.py
from typing import Any
from typing import Callable
from typing import NamedTuple
from typing import Optional
from dataclasses import dataclass
from jaxopt._src import base
from jaxopt._src import linear_solve
from jaxopt.tree_util import tree_l2_norm, tree_sub
import jax
import jax.numpy as jnp

    # Compute actual values of state variables at init_param
    residual, aux = self._fun_with_aux(init_params, *args, **kwargs)
    matvec = lambda v: self._jtj_op(init_params, v, *args, **kwargs)
    gradient = self._jt_op(init_params, residual, *args, **kwargs)

    return GaussNewtonState(iter_num=jnp.asarray(0),
                            error=jnp.asarray(jnp.inf),
                            residual=residual,
                            delta=init_params,
                            gradient=gradient,
                            aux=aux)

  def update(self,
             params,
             state: NamedTuple,
             *args,
             **kwargs) -> base.OptStep:
    """Performs one iteration of the least-squares solver.

    Args:
      params: pytree containing the parameters.
      state: named tuple containing the solver state.
    Returns:
      (params, state)
    """
    residual, aux = self._fun_with_aux(params, *args, **kwargs)

    matvec = lambda v: self._jtj_op(params, v, *args, **kwargs)
    gradient = self._jt_op(params, residual, *args, **kwargs)


    delta_params = linear_solve.solve_cg(matvec, gradient)