# Path: jaxopt/tree_util.py

# 
# Path: jaxopt/tree_util.py

# 
# Path: jaxopt/tree_util.py

# 
# 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/loop.py
# def _while_loop_scan(cond_fun, body_fun, init_val, max_iter):
#   def _iter(val):
#   def _fun(tup, it):
# def _while_loop_python(cond_fun, body_fun, init_val, maxiter):
# def _while_loop_lax(cond_fun, body_fun, init_val, maxiter):
#   def _cond_fun(_val):
#   def _body_fun(_val):
# def while_loop(cond_fun, body_fun, init_val, maxiter, unroll=False, jit=False):


from dataclasses import dataclass
from functools import partial
from typing import Any
from typing import Callable
from typing import NamedTuple
from typing import Optional
from jaxopt.tree_util import tree_add_scalar_mul, tree_l2_norm
from jaxopt.tree_util import tree_scalar_mul, tree_zeros_like
from jaxopt.tree_util import tree_add, tree_sub
from jaxopt._src import base
from jaxopt._src import loop
import jax
import jax.lax as lax
import jax.numpy as jnp
  def update_stepsize(t):
    """Multiply stepsize per factor, return new params and new value."""
    stepsize, factor = t
    stepsize = stepsize * factor
    stepsize = jnp.minimum(stepsize, max_stepsize)
    next_params = tree_add_scalar_mul(params, -stepsize, grad)
    f_next, _ = fun_with_aux(next_params, *args, **kwargs)
    return stepsize, next_params, f_next

  def body_fun(t):
    stepsize, next_params, f_next, _ = t

    violated = wolfe_cond_violated(stepsize, coef, f_cur, f_next, grad_sqnorm)
    stepsize, next_params, f_next = lax.cond(
      violated, update_stepsize,
      lambda _: (stepsize, next_params, f_next),
      operand=(stepsize, decrease_factor))

    if goldstein:
      goldstein_violated = curvature_cond_violated(stepsize, coef, f_cur,
                                                   f_next, grad_sqnorm)
      stepsize, next_params, f_next = lax.cond(
        goldstein_violated, update_stepsize,
        lambda _: (stepsize, next_params, f_next),
        operand=(stepsize, increase_factor))
      violated = jnp.logical_or(violated, goldstein_violated)

    return stepsize, next_params, f_next, violated

  init_val = stepsize, next_params, f_next, jnp.array(True)

# Path: jaxopt/tree_util.py

# 
# Path: jaxopt/tree_util.py

# 
# Path: jaxopt/tree_util.py

# 
# 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/loop.py
# def _while_loop_scan(cond_fun, body_fun, init_val, max_iter):
#   def _iter(val):
#   def _fun(tup, it):
# def _while_loop_python(cond_fun, body_fun, init_val, maxiter):
# def _while_loop_lax(cond_fun, body_fun, init_val, maxiter):
#   def _cond_fun(_val):
#   def _body_fun(_val):
# def while_loop(cond_fun, body_fun, init_val, maxiter, unroll=False, jit=False):


# Path: jaxopt/_src/armijo_sgd.py
from dataclasses import dataclass
from functools import partial
from typing import Any
from typing import Callable
from typing import NamedTuple
from typing import Optional
from jaxopt.tree_util import tree_add_scalar_mul, tree_l2_norm
from jaxopt.tree_util import tree_scalar_mul, tree_zeros_like
from jaxopt.tree_util import tree_add, tree_sub
from jaxopt._src import base
from jaxopt._src import loop
import jax
import jax.lax as lax
import jax.numpy as jnp

  def update_stepsize(t):
    """Multiply stepsize per factor, return new params and new value."""
    stepsize, factor = t
    stepsize = stepsize * factor
    stepsize = jnp.minimum(stepsize, max_stepsize)
    next_params = tree_add_scalar_mul(params, -stepsize, grad)
    f_next, _ = fun_with_aux(next_params, *args, **kwargs)
    return stepsize, next_params, f_next

  def body_fun(t):
    stepsize, next_params, f_next, _ = t

    violated = wolfe_cond_violated(stepsize, coef, f_cur, f_next, grad_sqnorm)
    stepsize, next_params, f_next = lax.cond(
      violated, update_stepsize,
      lambda _: (stepsize, next_params, f_next),
      operand=(stepsize, decrease_factor))

    if goldstein:
      goldstein_violated = curvature_cond_violated(stepsize, coef, f_cur,
                                                   f_next, grad_sqnorm)
      stepsize, next_params, f_next = lax.cond(
        goldstein_violated, update_stepsize,
        lambda _: (stepsize, next_params, f_next),
        operand=(stepsize, increase_factor))
      violated = jnp.logical_or(violated, goldstein_violated)

    return stepsize, next_params, f_next, violated

  init_val = stepsize, next_params, f_next, jnp.array(True)

  ret = loop.while_loop(cond_fun=lambda t: t[-1],  # check boolean violated