# Path: jaxopt/loss.py

# 
# Path: jaxopt/projection.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 jax.nn import softmax
from jax.scipy.special import expit as sigmoid
from jaxopt import loss
from jaxopt import projection
from jaxopt._src import test_util
import jax
import jax.numpy as jnp

  def test_binary_logistic_loss(self):
    self._test_binary_loss_function(loss.binary_logistic_loss, sigmoid)

  def _test_multiclass_loss_function(self, loss_fun, inv_link_fun):
    # Check that loss is zero when all weights goes to the correct label.
    loss_val = loss_fun(0, jnp.array([1e5, 0, 0]))
    self.assertEqual(loss_val, 0)

    # Check that gradient has the expected form.
    logits = jnp.array([1.2, 0.3, 2.3])
    grad = jax.grad(loss_fun, argnums=1)(0, logits)
    self.assertArraysAllClose(grad, inv_link_fun(logits) - jnp.array([1, 0, 0]))

    # Check that vmap works as expected.
    logits = jnp.array([[1.3, 2.5],
                        [1.7, -0.3],
                        [0.2, 1.2]])
    labels = jnp.array([0, 0, 1])
    losses = jax.vmap(loss_fun)(labels, logits)
    losses2 = jnp.array([loss_fun(labels[0], logits[0]),
                         loss_fun(labels[1], logits[1]),
                         loss_fun(labels[2], logits[2])])
    self.assertArraysAllClose(losses, losses2)

  def test_multiclass_logistic_loss(self):
    self._test_multiclass_loss_function(loss.multiclass_logistic_loss, softmax)

  def test_multiclass_sparsemax_loss(self):
    self._test_multiclass_loss_function(loss.multiclass_sparsemax_loss,

# Path: jaxopt/loss.py

# 
# Path: jaxopt/projection.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/loss_test.py
from absl.testing import absltest
from jax.nn import softmax
from jax.scipy.special import expit as sigmoid
from jaxopt import loss
from jaxopt import projection
from jaxopt._src import test_util
import jax
import jax.numpy as jnp


  def test_binary_logistic_loss(self):
    self._test_binary_loss_function(loss.binary_logistic_loss, sigmoid)

  def _test_multiclass_loss_function(self, loss_fun, inv_link_fun):
    # Check that loss is zero when all weights goes to the correct label.
    loss_val = loss_fun(0, jnp.array([1e5, 0, 0]))
    self.assertEqual(loss_val, 0)

    # Check that gradient has the expected form.
    logits = jnp.array([1.2, 0.3, 2.3])
    grad = jax.grad(loss_fun, argnums=1)(0, logits)
    self.assertArraysAllClose(grad, inv_link_fun(logits) - jnp.array([1, 0, 0]))

    # Check that vmap works as expected.
    logits = jnp.array([[1.3, 2.5],
                        [1.7, -0.3],
                        [0.2, 1.2]])
    labels = jnp.array([0, 0, 1])
    losses = jax.vmap(loss_fun)(labels, logits)
    losses2 = jnp.array([loss_fun(labels[0], logits[0]),
                         loss_fun(labels[1], logits[1]),
                         loss_fun(labels[2], logits[2])])
    self.assertArraysAllClose(losses, losses2)

  def test_multiclass_logistic_loss(self):
    self._test_multiclass_loss_function(loss.multiclass_logistic_loss, softmax)

  def test_multiclass_sparsemax_loss(self):
    self._test_multiclass_loss_function(loss.multiclass_sparsemax_loss,

                                        projection.projection_simplex)