# Path: jaxopt/implicit_diff.py

# 
# 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/loop.py

# 
# Path: jaxopt/tree_util.py



import abc
import dataclasses
import functools
import itertools
import jax
import jax.numpy as jnp
from typing import Any
from typing import Callable
from typing import NamedTuple
from typing import Optional
from typing import Union
from typing import Tuple
from jaxopt import implicit_diff as idf
from jaxopt._src import linear_solve
from jaxopt import loop
from jaxopt import tree_util

  def _run(self,
           init_params: Any,
           *args,
           **kwargs) -> OptStep:
    state = self.init_state(init_params, *args, **kwargs)

    # We unroll the very first iteration. This allows `init_val` and `body_fun`
    # below to have the same output type, which is a requirement of
    # lax.while_loop and lax.scan.
    #
    # TODO(frostig,mblondel): if we could check concreteness of self.maxiter,
    # and we knew that it is concrete here, then we could optimize away the
    # redundant first step, e.g.:
    #
    #   maxiter = get_maybe_concrete(self.maxiter)  # concrete value or None
    #   if maxiter == 0:
    #     return OptStep(params=init_params, state=state)
    #
    # In the general case below, we prefer to use `jnp.where` instead
    # of a `lax.cond` for now in order to avoid staging the initial
    # update and the run loop. They might not be staging compatible.

    zero_step = self._make_zero_step(init_params, state)

    opt_step = self.update(init_params, state, *args, **kwargs)
    init_val = (opt_step, (args, kwargs))

    jit, unroll = self._get_loop_options()


# Path: jaxopt/implicit_diff.py

# 
# 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/loop.py

# 
# Path: jaxopt/tree_util.py



# Path: jaxopt/_src/base.py
import abc
import dataclasses
import functools
import itertools
import jax
import jax.numpy as jnp
from typing import Any
from typing import Callable
from typing import NamedTuple
from typing import Optional
from typing import Union
from typing import Tuple
from jaxopt import implicit_diff as idf
from jaxopt._src import linear_solve
from jaxopt import loop
from jaxopt import tree_util


  def _run(self,
           init_params: Any,
           *args,
           **kwargs) -> OptStep:
    state = self.init_state(init_params, *args, **kwargs)

    # We unroll the very first iteration. This allows `init_val` and `body_fun`
    # below to have the same output type, which is a requirement of
    # lax.while_loop and lax.scan.
    #
    # TODO(frostig,mblondel): if we could check concreteness of self.maxiter,
    # and we knew that it is concrete here, then we could optimize away the
    # redundant first step, e.g.:
    #
    #   maxiter = get_maybe_concrete(self.maxiter)  # concrete value or None
    #   if maxiter == 0:
    #     return OptStep(params=init_params, state=state)
    #
    # In the general case below, we prefer to use `jnp.where` instead
    # of a `lax.cond` for now in order to avoid staging the initial
    # update and the run loop. They might not be staging compatible.

    zero_step = self._make_zero_step(init_params, state)

    opt_step = self.update(init_params, state, *args, **kwargs)
    init_val = (opt_step, (args, kwargs))

    jit, unroll = self._get_loop_options()


    many_step = loop.while_loop(