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

# 
# Path: jaxopt/tree_util.py

# 
# Path: jaxopt/tree_util.py

# 
# Path: jaxopt/tree_util.py

# 
# Path: jaxopt/tree_util.py

# 
# Path: jaxopt/tree_util.py



from typing import Any
from typing import Callable
from typing import NamedTuple
from typing import Optional
from jaxopt._src import base
from jaxopt.tree_util import tree_add
from jaxopt.tree_util import tree_add_scalar_mul
from jaxopt.tree_util import tree_l2_norm
from jaxopt.tree_util import tree_scalar_mul
from jaxopt.tree_util import tree_sub
from jaxopt.tree_util import tree_zeros_like
import dataclasses
import jax
import jax.numpy as jnp
                          value=jnp.asarray(jnp.inf),
                          stepsize=jnp.asarray(1.0),
                          aux=None,
                          velocity=velocity)

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

    Args:
      params: pytree containing the parameters.
      state: named tuple containing the solver state.
      *args: additional positional arguments to be passed to ``fun``.
      **kwargs: additional keyword arguments to be passed to ``fun``.
    Returns:
      (params, state)
    """
    if self.pre_update:
      params, state = self.pre_update(params, state, *args, **kwargs)

    (value, aux), grad = self._value_and_grad_fun(params, *args, **kwargs)

    grad_sqnorm = tree_l2_norm(grad, squared=True)
    stepsize = jax.lax.min(value / (grad_sqnorm + self.delta),
                           self.max_stepsize)

    if self.momentum == 0:

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

# 
# Path: jaxopt/tree_util.py

# 
# Path: jaxopt/tree_util.py

# 
# Path: jaxopt/tree_util.py

# 
# Path: jaxopt/tree_util.py

# 
# Path: jaxopt/tree_util.py



# Path: jaxopt/_src/polyak_sgd.py
from typing import Any
from typing import Callable
from typing import NamedTuple
from typing import Optional
from jaxopt._src import base
from jaxopt.tree_util import tree_add
from jaxopt.tree_util import tree_add_scalar_mul
from jaxopt.tree_util import tree_l2_norm
from jaxopt.tree_util import tree_scalar_mul
from jaxopt.tree_util import tree_sub
from jaxopt.tree_util import tree_zeros_like
import dataclasses
import jax
import jax.numpy as jnp

                          value=jnp.asarray(jnp.inf),
                          stepsize=jnp.asarray(1.0),
                          aux=None,
                          velocity=velocity)

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

    Args:
      params: pytree containing the parameters.
      state: named tuple containing the solver state.
      *args: additional positional arguments to be passed to ``fun``.
      **kwargs: additional keyword arguments to be passed to ``fun``.
    Returns:
      (params, state)
    """
    if self.pre_update:
      params, state = self.pre_update(params, state, *args, **kwargs)

    (value, aux), grad = self._value_and_grad_fun(params, *args, **kwargs)

    grad_sqnorm = tree_l2_norm(grad, squared=True)
    stepsize = jax.lax.min(value / (grad_sqnorm + self.delta),
                           self.max_stepsize)

    if self.momentum == 0:

      new_params = tree_add_scalar_mul(params, -stepsize, grad)