# 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
    init_val = (opt_step, (args, kwargs))

    jit, unroll = self._get_loop_options()

    many_step = loop.while_loop(
        cond_fun=self._cond_fun, body_fun=self._body_fun,
        init_val=init_val, maxiter=self.maxiter - 1, jit=jit,
        unroll=unroll)[0]

    return tree_util.tree_map(
        functools.partial(_where, self.maxiter == 0), zero_step, many_step,
        is_leaf=lambda x: x is None)  # state attributes can sometimes be None

  def run(self,
          init_params: Any,
          *args,
          **kwargs) -> OptStep:
    """Runs the optimization loop.

    Args:
      init_params: pytree containing the initial parameters.
      *args: additional positional arguments to be passed to the update method.
      **kwargs: additional keyword arguments to be passed to the update method.
    Returns:
      (params, state)
    """
    run = self._run

    if getattr(self, "implicit_diff", True):
      reference_signature = getattr(self, "reference_signature", None)

# 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

    init_val = (opt_step, (args, kwargs))

    jit, unroll = self._get_loop_options()

    many_step = loop.while_loop(
        cond_fun=self._cond_fun, body_fun=self._body_fun,
        init_val=init_val, maxiter=self.maxiter - 1, jit=jit,
        unroll=unroll)[0]

    return tree_util.tree_map(
        functools.partial(_where, self.maxiter == 0), zero_step, many_step,
        is_leaf=lambda x: x is None)  # state attributes can sometimes be None

  def run(self,
          init_params: Any,
          *args,
          **kwargs) -> OptStep:
    """Runs the optimization loop.

    Args:
      init_params: pytree containing the initial parameters.
      *args: additional positional arguments to be passed to the update method.
      **kwargs: additional keyword arguments to be passed to the update method.
    Returns:
      (params, state)
    """
    run = self._run

    if getattr(self, "implicit_diff", True):
      reference_signature = getattr(self, "reference_signature", None)

      decorator = idf.custom_root(