# 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


class Solver(abc.ABC):
  """Base class for solvers.

  A solver should implement a `run` method:

    `params, state = run(init_params, *args, **kwargs)`

  In addition, we will assume that the solver defines a notion of optimality:
    - `pytree = optimality_fun(params, *args, **kwargs)

  Note: subclasses of `Solver` which are also Python dataclasses must disable
  the generation of an `__eq__` method (`eq=False`) for compatibility with
  Python 3.7. Not doing so results in the generation of a `__hash__` method
  that clashes with JAX transformations when applied directly to methods of
  the class. This restriction will be lifted in the future when support for
  Python 3.7 is dropped and / or another solution is found.
  """

  @abc.abstractmethod
  def run(self,
          init_params: Any,
          *args,
          **kwargs) -> OptStep:
    pass

  def l2_optimality_error(self, params, *args, **kwargs):
    """Computes the L2 optimality error."""
    optimality = self.optimality_fun(params, *args, **kwargs)

# 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



class Solver(abc.ABC):
  """Base class for solvers.

  A solver should implement a `run` method:

    `params, state = run(init_params, *args, **kwargs)`

  In addition, we will assume that the solver defines a notion of optimality:
    - `pytree = optimality_fun(params, *args, **kwargs)

  Note: subclasses of `Solver` which are also Python dataclasses must disable
  the generation of an `__eq__` method (`eq=False`) for compatibility with
  Python 3.7. Not doing so results in the generation of a `__hash__` method
  that clashes with JAX transformations when applied directly to methods of
  the class. This restriction will be lifted in the future when support for
  Python 3.7 is dropped and / or another solution is found.
  """

  @abc.abstractmethod
  def run(self,
          init_params: Any,
          *args,
          **kwargs) -> OptStep:
    pass

  def l2_optimality_error(self, params, *args, **kwargs):
    """Computes the L2 optimality error."""
    optimality = self.optimality_fun(params, *args, **kwargs)

    return tree_util.tree_l2_norm(optimality)