# Path: jaxopt/objective.py

# 
# Path: jaxopt/projection.py

# 
# Path: jaxopt/prox.py

# 
# Path: jaxopt/_src/test_util.py
# def ridge_solver(X, y, lam):
# def ridge_solver_jac(X, y, lam, eps=1e-4):
# def lasso_skl(X, y, lam, tol=1e-5, fit_intercept=False):
# def lasso_skl_jac(X, y, lam, tol=1e-5, fit_intercept=False, eps=1e-4):
# def enet_skl(X, y, params_prox, tol=1e-5, fit_intercept=False):
# def enet_skl_jac(X, y, params_prox, tol=1e-5, fit_intercept=False, eps=1e-5):
# def multitask_lasso_skl(X, Y, lam, tol=1e-5):
# def logreg_skl(X, y, lam, tol=1e-5, fit_intercept=False,
#                penalty="l2", multiclass=True):
#   def _reshape_coef(coef):
# def logreg_skl_jac(X, y, lam, tol=1e-5, fit_intercept=False, penalty="l2",
#                    multiclass=True, eps=1e-5):
# def multiclass_linear_svm_skl(X, y, lam, tol=1e-5):
# def multiclass_linear_svm_skl_jac(X, y, lam, tol=1e-5, eps=1e-5):
# def lsq_linear_osp(X, y, bounds, tol=1e-10, max_iter=None):
# def lsq_linear_cube_osp(X, y, l, tol=1e-10, max_iter=None):
# def lsq_linear_cube_osp_jac(X, y, l, eps=1e-5, tol=1e-10, max_iter=None):
# def check_states_have_same_types(state1, state2):
# def _canonicalize_dtype(x64_enabled, dtype):
# def canonicalize_dtype(dtype):
# def _dtype(x):
# def default_tolerance():
# def tolerance(dtype, tol=None):
# def device_under_test():
# def _assert_numpy_allclose(a, b, atol=None, rtol=None, err_msg=''):
# def is_sequence(x):
#   def assertArraysEqual(self, x, y, *, check_dtypes=True, err_msg=''):
#   def assertArraysAllClose(self, x, y, *, check_dtypes=True, atol=None,
#                            rtol=None, err_msg=''):
#   def assertDtypesMatch(self, x, y, *, canonicalize_dtypes=True):
#   def assertAllClose(self, x, y, *, check_dtypes=True, atol=None, rtol=None,
#                      canonicalize_dtypes=True, err_msg=''):
#   XX = jnp.dot(X.T, X)
#   I = jnp.eye(X.shape[1])
# class JaxoptTestCase(parameterized.TestCase):


from absl.testing import absltest
from jaxopt import objective
from jaxopt import projection
from jaxopt import prox
from jaxopt._src import test_util
from sklearn import datasets
import jax
import jax.numpy as jnp
import jaxopt
import optax
    root_fun = jax.grad(fun)

    def fixed_point_fun(params, *args, **kwargs):
      return root_fun(params, *args, **kwargs) - params


    X, y = datasets.make_classification(n_samples=10, n_features=5, n_classes=3,
                                        n_informative=3, random_state=0)
    data = (X, y)
    params0 = jnp.zeros(X.shape[1])

    fpi = jaxopt.FixedPointIteration(fixed_point_fun=fixed_point_fun)

    # Gradient solvers and fixed point
    for solver in (jaxopt.GradientDescent(fun=fun),
                   jaxopt.PolyakSGD(fun=fun),
                   jaxopt.OptaxSolver(opt=optax.adam(1e-3), fun=fun),
                   fpi,
                   jaxopt.AndersonAcceleration(fixed_point_fun=fixed_point_fun),
                   jaxopt.AndersonWrapper(fpi),
                   jaxopt.ArmijoSGD(fun=fun),
                   jaxopt.LBFGS(fun=fun)):

      hash(solver)  # Checks that hash works.
      update = jax.jit(solver.update)
      state0 = solver.init_state(params0)
      params, state = update(params0, state0, data=data)
      test_util.check_states_have_same_types(state0, state)

    # Proximal gradient solvers.

# Path: jaxopt/objective.py

# 
# Path: jaxopt/projection.py

# 
# Path: jaxopt/prox.py

# 
# Path: jaxopt/_src/test_util.py
# def ridge_solver(X, y, lam):
# def ridge_solver_jac(X, y, lam, eps=1e-4):
# def lasso_skl(X, y, lam, tol=1e-5, fit_intercept=False):
# def lasso_skl_jac(X, y, lam, tol=1e-5, fit_intercept=False, eps=1e-4):
# def enet_skl(X, y, params_prox, tol=1e-5, fit_intercept=False):
# def enet_skl_jac(X, y, params_prox, tol=1e-5, fit_intercept=False, eps=1e-5):
# def multitask_lasso_skl(X, Y, lam, tol=1e-5):
# def logreg_skl(X, y, lam, tol=1e-5, fit_intercept=False,
#                penalty="l2", multiclass=True):
#   def _reshape_coef(coef):
# def logreg_skl_jac(X, y, lam, tol=1e-5, fit_intercept=False, penalty="l2",
#                    multiclass=True, eps=1e-5):
# def multiclass_linear_svm_skl(X, y, lam, tol=1e-5):
# def multiclass_linear_svm_skl_jac(X, y, lam, tol=1e-5, eps=1e-5):
# def lsq_linear_osp(X, y, bounds, tol=1e-10, max_iter=None):
# def lsq_linear_cube_osp(X, y, l, tol=1e-10, max_iter=None):
# def lsq_linear_cube_osp_jac(X, y, l, eps=1e-5, tol=1e-10, max_iter=None):
# def check_states_have_same_types(state1, state2):
# def _canonicalize_dtype(x64_enabled, dtype):
# def canonicalize_dtype(dtype):
# def _dtype(x):
# def default_tolerance():
# def tolerance(dtype, tol=None):
# def device_under_test():
# def _assert_numpy_allclose(a, b, atol=None, rtol=None, err_msg=''):
# def is_sequence(x):
#   def assertArraysEqual(self, x, y, *, check_dtypes=True, err_msg=''):
#   def assertArraysAllClose(self, x, y, *, check_dtypes=True, atol=None,
#                            rtol=None, err_msg=''):
#   def assertDtypesMatch(self, x, y, *, canonicalize_dtypes=True):
#   def assertAllClose(self, x, y, *, check_dtypes=True, atol=None, rtol=None,
#                      canonicalize_dtypes=True, err_msg=''):
#   XX = jnp.dot(X.T, X)
#   I = jnp.eye(X.shape[1])
# class JaxoptTestCase(parameterized.TestCase):


# Path: tests/common_test.py
from absl.testing import absltest
from jaxopt import objective
from jaxopt import projection
from jaxopt import prox
from jaxopt._src import test_util
from sklearn import datasets
import jax
import jax.numpy as jnp
import jaxopt
import optax

    root_fun = jax.grad(fun)

    def fixed_point_fun(params, *args, **kwargs):
      return root_fun(params, *args, **kwargs) - params


    X, y = datasets.make_classification(n_samples=10, n_features=5, n_classes=3,
                                        n_informative=3, random_state=0)
    data = (X, y)
    params0 = jnp.zeros(X.shape[1])

    fpi = jaxopt.FixedPointIteration(fixed_point_fun=fixed_point_fun)

    # Gradient solvers and fixed point
    for solver in (jaxopt.GradientDescent(fun=fun),
                   jaxopt.PolyakSGD(fun=fun),
                   jaxopt.OptaxSolver(opt=optax.adam(1e-3), fun=fun),
                   fpi,
                   jaxopt.AndersonAcceleration(fixed_point_fun=fixed_point_fun),
                   jaxopt.AndersonWrapper(fpi),
                   jaxopt.ArmijoSGD(fun=fun),
                   jaxopt.LBFGS(fun=fun)):

      hash(solver)  # Checks that hash works.
      update = jax.jit(solver.update)
      state0 = solver.init_state(params0)
      params, state = update(params0, state0, data=data)
      test_util.check_states_have_same_types(state0, state)

    # Proximal gradient solvers.

    for solver in (jaxopt.ProximalGradient(fun=fun, prox=prox.prox_lasso),