# 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/loss.py
# def huber_loss(target: float, pred: float, delta: float = 1.0) -> float:
# def binary_logistic_loss(label: int, logit: float) -> float:
# def multiclass_logistic_loss(label: int, logits: jnp.ndarray) -> float:
# def multiclass_sparsemax_loss(label: int, scores: jnp.ndarray) -> float:


from typing import Tuple
from jaxopt._src import base
from jaxopt._src import loss
import jax
import jax.numpy as jnp

def ridge_regression(
  params: jnp.ndarray,
  l2reg: float,
  data: Tuple[jnp.ndarray, jnp.ndarray]) -> float:
  r"""
  Ridge regression, i.e L2-regularized least squares.

  .. math::

    \frac{1}{2n} ||XW - y||_2^2 +
    0.5 \cdot \text{l2reg} \cdot ||W||_2^2

  Args:
    W: parameters.
    l2reg: strenght of regularization.
    data: a tuple ``(X, y)`` where ``X`` is a matrix of shape ``(n_samples,
      n_features)`` and ``y`` is a vector of shape ``(n_samples,)``.
  Returns:
    objective value.

  Example::

    value = ridge_regression(W, l2reg, (X, y))
  """
  least_squares = LeastSquares()(params, data)
  ridge = 0.5 * l2reg * jnp.sum(params ** 2)
  return least_squares + ridge



# 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/loss.py
# def huber_loss(target: float, pred: float, delta: float = 1.0) -> float:
# def binary_logistic_loss(label: int, logit: float) -> float:
# def multiclass_logistic_loss(label: int, logits: jnp.ndarray) -> float:
# def multiclass_sparsemax_loss(label: int, scores: jnp.ndarray) -> float:


# Path: jaxopt/_src/objective.py
from typing import Tuple
from jaxopt._src import base
from jaxopt._src import loss
import jax
import jax.numpy as jnp


def ridge_regression(
  params: jnp.ndarray,
  l2reg: float,
  data: Tuple[jnp.ndarray, jnp.ndarray]) -> float:
  r"""
  Ridge regression, i.e L2-regularized least squares.

  .. math::

    \frac{1}{2n} ||XW - y||_2^2 +
    0.5 \cdot \text{l2reg} \cdot ||W||_2^2

  Args:
    W: parameters.
    l2reg: strenght of regularization.
    data: a tuple ``(X, y)`` where ``X`` is a matrix of shape ``(n_samples,
      n_features)`` and ``y`` is a vector of shape ``(n_samples,)``.
  Returns:
    objective value.

  Example::

    value = ridge_regression(W, l2reg, (X, y))
  """
  least_squares = LeastSquares()(params, data)
  ridge = 0.5 * l2reg * jnp.sum(params ** 2)
  return least_squares + ridge



_logloss_vmap = jax.vmap(loss.multiclass_logistic_loss)