# 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/tree_util.py
# def tree_add_scalar_mul(tree_x, scalar, tree_y):
#   """Compute tree_x + scalar * tree_y."""
#   return tree_multimap(lambda x, y: x + scalar * y, tree_x, tree_y)
# 
# Path: jaxopt/_src/tree_util.py
# def tree_l2_norm(tree_x, squared=False):
#   """Compute the l2 norm ||tree_x||."""
#   squared_tree = tree_map(jnp.square, tree_x)
#   sqnorm = tree_sum(squared_tree)
#   if squared:
#     return sqnorm
#   else:
#     return jnp.sqrt(sqnorm)
# 
# Path: jaxopt/_src/tree_util.py
# def broadcast_pytrees(*trees):
# def tree_scalar_mul(scalar, tree_x):
# def tree_add_scalar_mul(tree_x, scalar, tree_y):
# def _vdot_safe(a, b):
# def tree_vdot(tree_x, tree_y):
# def tree_dot(tree_x, tree_y):
# def tree_sum(tree_x):
# def tree_l2_norm(tree_x, squared=False):
# def tree_zeros_like(tree_x):
# def tree_ones_like(tree_x):
# def tree_average(trees, weights):
# def tree_gram(a):
# def tree_inf_norm(tree_x):
# def tree_where(cond, a, b):
# def tree_negative(tree):
# def tree_reciproqual(tree):
# def tree_mean(tree):


import inspect
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 dataclasses import dataclass
from jaxopt._src import base
from jaxopt._src.tree_util import tree_add_scalar_mul
from jaxopt._src.tree_util import tree_l2_norm
from jaxopt._src.tree_util import tree_sub
  """
  fun: Callable
  projection_grad: Optional[Callable]
  stepsize: Union[float, Callable]
  maxiter: int = 500
  tol: float = 1e-2
  verbose: int = 0
  implicit_diff: bool = True
  implicit_diff_solve: Optional[Callable] = None
  has_aux: bool = False
  jit: base.AutoOrBoolean = "auto"
  unroll: base.AutoOrBoolean = "auto"

  @staticmethod
  def make_projection_grad(projection: Callable,
                           mapping_fun: Callable) -> Callable:
    """Instantiates `projection_grad` argument from projection and mirror map.

    Args:
      projection: projection operator of the form
        ``projection(x, hyperparams_proj)``, typically
        ``argmin_z D_{gen_fun}(z, mapping_fun^{-1}(y))``.
      mapping_fun: the mirror map, typically of the form
        ``mapping_fun = grad(gen_fun)``, where `gen_fun` is the generating
        function of the Bregman divergence.
    Returns:
      A function `projection_grad(x, g, stepsize, hyperparams_proj)`
      representing the mirror descent update for iterate x and gradient g.
    """
    def projection_grad(x, x_fun_grad, stepsize, hyperparams_proj):

# 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/tree_util.py
# def tree_add_scalar_mul(tree_x, scalar, tree_y):
#   """Compute tree_x + scalar * tree_y."""
#   return tree_multimap(lambda x, y: x + scalar * y, tree_x, tree_y)
# 
# Path: jaxopt/_src/tree_util.py
# def tree_l2_norm(tree_x, squared=False):
#   """Compute the l2 norm ||tree_x||."""
#   squared_tree = tree_map(jnp.square, tree_x)
#   sqnorm = tree_sum(squared_tree)
#   if squared:
#     return sqnorm
#   else:
#     return jnp.sqrt(sqnorm)
# 
# Path: jaxopt/_src/tree_util.py
# def broadcast_pytrees(*trees):
# def tree_scalar_mul(scalar, tree_x):
# def tree_add_scalar_mul(tree_x, scalar, tree_y):
# def _vdot_safe(a, b):
# def tree_vdot(tree_x, tree_y):
# def tree_dot(tree_x, tree_y):
# def tree_sum(tree_x):
# def tree_l2_norm(tree_x, squared=False):
# def tree_zeros_like(tree_x):
# def tree_ones_like(tree_x):
# def tree_average(trees, weights):
# def tree_gram(a):
# def tree_inf_norm(tree_x):
# def tree_where(cond, a, b):
# def tree_negative(tree):
# def tree_reciproqual(tree):
# def tree_mean(tree):


# Path: jaxopt/_src/mirror_descent.py
import inspect
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 dataclasses import dataclass
from jaxopt._src import base
from jaxopt._src.tree_util import tree_add_scalar_mul
from jaxopt._src.tree_util import tree_l2_norm
from jaxopt._src.tree_util import tree_sub

  """
  fun: Callable
  projection_grad: Optional[Callable]
  stepsize: Union[float, Callable]
  maxiter: int = 500
  tol: float = 1e-2
  verbose: int = 0
  implicit_diff: bool = True
  implicit_diff_solve: Optional[Callable] = None
  has_aux: bool = False
  jit: base.AutoOrBoolean = "auto"
  unroll: base.AutoOrBoolean = "auto"

  @staticmethod
  def make_projection_grad(projection: Callable,
                           mapping_fun: Callable) -> Callable:
    """Instantiates `projection_grad` argument from projection and mirror map.

    Args:
      projection: projection operator of the form
        ``projection(x, hyperparams_proj)``, typically
        ``argmin_z D_{gen_fun}(z, mapping_fun^{-1}(y))``.
      mapping_fun: the mirror map, typically of the form
        ``mapping_fun = grad(gen_fun)``, where `gen_fun` is the generating
        function of the Bregman divergence.
    Returns:
      A function `projection_grad(x, g, stepsize, hyperparams_proj)`
      representing the mirror descent update for iterate x and gradient g.
    """
    def projection_grad(x, x_fun_grad, stepsize, hyperparams_proj):

      update = tree_add_scalar_mul(mapping_fun(x), -stepsize, x_fun_grad)