# Path: jaxopt/_src/projection.py
# def projection_simplex(x: jnp.ndarray, value: float = 1.0) -> jnp.ndarray:
#   r"""Projection onto a simplex:
# 
#   .. math::
# 
#     \underset{p}{\text{argmin}} ~ ||x - p||_2^2 \quad \textrm{subject to} \quad
#     p \ge 0, p^\top 1 = \text{value}
# 
#   By default, the projection is onto the probability simplex.
# 
#   Args:
#     x: vector to project, an array of shape (n,).
#     value: value p should sum to (default: 1.0).
#   Returns:
#     projected vector, an array with the same shape as ``x``.
#   """
#   if value is None:
#     value = 1.0
#   return value * _projection_unit_simplex(x / value)


import jax
import jax.numpy as jnp
from jax.nn import softplus
from jax.scipy.special import logsumexp
from jaxopt._src.projection import projection_simplex

  Args:
    label: ground-truth integer label, between 0 and n_classes - 1.
    logits: scores produced by the model, shape = (n_classes, ).
  Returns:
    loss value
  """
  n_classes = logits.shape[0]
  one_hot = jax.nn.one_hot(label, n_classes)
  # Logsumexp is the Fenchel conjugate of the Shannon negentropy on the simplex.
  # logsumexp = jnp.dot(proba, logits) - jnp.dot(proba, jnp.log(proba))
  return logsumexp(logits) - jnp.dot(logits, one_hot)


def multiclass_sparsemax_loss(label: int, scores: jnp.ndarray) -> float:
  """Multiclass sparsemax loss.

  Args:
    label: ground-truth integer label, between 0 and n_classes - 1.
    scores: scores produced by the model, shape = (n_classes, ).
  Returns:
    loss value

  References:
    From Softmax to Sparsemax: A Sparse Model of Attention and Multi-Label
    Classification. André F. T. Martins, Ramón Fernandez Astudillo.
    ICML 2016.
  """
  n_classes = scores.shape[0]
  one_hot = jax.nn.one_hot(label, n_classes)

# Path: jaxopt/_src/projection.py
# def projection_simplex(x: jnp.ndarray, value: float = 1.0) -> jnp.ndarray:
#   r"""Projection onto a simplex:
# 
#   .. math::
# 
#     \underset{p}{\text{argmin}} ~ ||x - p||_2^2 \quad \textrm{subject to} \quad
#     p \ge 0, p^\top 1 = \text{value}
# 
#   By default, the projection is onto the probability simplex.
# 
#   Args:
#     x: vector to project, an array of shape (n,).
#     value: value p should sum to (default: 1.0).
#   Returns:
#     projected vector, an array with the same shape as ``x``.
#   """
#   if value is None:
#     value = 1.0
#   return value * _projection_unit_simplex(x / value)


# Path: jaxopt/_src/loss.py
import jax
import jax.numpy as jnp
from jax.nn import softplus
from jax.scipy.special import logsumexp
from jaxopt._src.projection import projection_simplex


  Args:
    label: ground-truth integer label, between 0 and n_classes - 1.
    logits: scores produced by the model, shape = (n_classes, ).
  Returns:
    loss value
  """
  n_classes = logits.shape[0]
  one_hot = jax.nn.one_hot(label, n_classes)
  # Logsumexp is the Fenchel conjugate of the Shannon negentropy on the simplex.
  # logsumexp = jnp.dot(proba, logits) - jnp.dot(proba, jnp.log(proba))
  return logsumexp(logits) - jnp.dot(logits, one_hot)


def multiclass_sparsemax_loss(label: int, scores: jnp.ndarray) -> float:
  """Multiclass sparsemax loss.

  Args:
    label: ground-truth integer label, between 0 and n_classes - 1.
    scores: scores produced by the model, shape = (n_classes, ).
  Returns:
    loss value

  References:
    From Softmax to Sparsemax: A Sparse Model of Attention and Multi-Label
    Classification. André F. T. Martins, Ramón Fernandez Astudillo.
    ICML 2016.
  """
  n_classes = scores.shape[0]
  one_hot = jax.nn.one_hot(label, n_classes)

  proba = projection_simplex(scores)