# Path: optax/_src/loss.py
# def l2_loss(
#     predictions: chex.Array,
#     targets: Optional[chex.Array] = None,
# ) -> chex.Array:
# def huber_loss(
#     predictions: chex.Array,
#     targets: Optional[chex.Array] = None,
#     delta: float = 1.) -> chex.Array:
# def smooth_labels(
#     labels: chex.Array,
#     alpha: float,
# ) -> jnp.ndarray:
# def sigmoid_binary_cross_entropy(logits, labels):
# def softmax_cross_entropy(
#     logits: chex.Array,
#     labels: chex.Array,
# ) -> chex.Array:
# def cosine_similarity(
#     predictions: chex.Array,
#     targets: chex.Array,
#     epsilon: float = 0.,
# ) -> chex.Array:
# def cosine_distance(
#     predictions: chex.Array,
#     targets: chex.Array,
#     epsilon: float = 0.,
# ) -> chex.Array:
# def log_cosh(
#     predictions: chex.Array,
#     targets: Optional[chex.Array] = None,
# ) -> chex.Array:
# def ctc_loss(logits: chex.Array,
#              logitpaddings: chex.Array,
#              labels: chex.Array,
#              labelpaddings: chex.Array,
#              blank_id: int = 0,
#              logepsilon: float = -1e5) -> chex.Array:
#   def update_phi_score(phi, added_score):
#   def loop_body(prev, x):


from absl.testing import absltest
from absl.testing import parameterized
from optax._src import loss
import chex
import jax
import jax.numpy as jnp
import numpy as np
    self.ys = np.array([[10., 1., -2.], [1., 4., 0.2]], dtype=np.float32)
    self.ts = np.array([[0., 1.2, 0.2], [1., -0.3, 0.]], dtype=np.float32)
    # distance computed expected output from `scipy 1.20`.
    self.exp = np.array([0.9358251989, 1.0464068465], dtype=np.float32)

  @chex.all_variants()
  def test_scalar_distance(self):
    """Tests for a full batch."""
    np.testing.assert_allclose(
        self.variant(loss.cosine_distance)(self.ys[0], self.ts[0]),
        self.exp[0], atol=1e-4)

  @chex.all_variants()
  def test_scalar_similarity(self):
    """Tests for a full batch."""
    np.testing.assert_allclose(
        self.variant(loss.cosine_similarity)(self.ys[0], self.ts[0]),
        1. - self.exp[0], atol=1e-4)

  @chex.all_variants()
  def test_batched_distance(self):
    """Tests for a full batch."""
    np.testing.assert_allclose(
        self.variant(loss.cosine_distance)(self.ys, self.ts),
        self.exp, atol=1e-4)

  @chex.all_variants()
  def test_batched_similarity(self):
    """Tests for a full batch."""
    np.testing.assert_allclose(

# Path: optax/_src/loss.py
# def l2_loss(
#     predictions: chex.Array,
#     targets: Optional[chex.Array] = None,
# ) -> chex.Array:
# def huber_loss(
#     predictions: chex.Array,
#     targets: Optional[chex.Array] = None,
#     delta: float = 1.) -> chex.Array:
# def smooth_labels(
#     labels: chex.Array,
#     alpha: float,
# ) -> jnp.ndarray:
# def sigmoid_binary_cross_entropy(logits, labels):
# def softmax_cross_entropy(
#     logits: chex.Array,
#     labels: chex.Array,
# ) -> chex.Array:
# def cosine_similarity(
#     predictions: chex.Array,
#     targets: chex.Array,
#     epsilon: float = 0.,
# ) -> chex.Array:
# def cosine_distance(
#     predictions: chex.Array,
#     targets: chex.Array,
#     epsilon: float = 0.,
# ) -> chex.Array:
# def log_cosh(
#     predictions: chex.Array,
#     targets: Optional[chex.Array] = None,
# ) -> chex.Array:
# def ctc_loss(logits: chex.Array,
#              logitpaddings: chex.Array,
#              labels: chex.Array,
#              labelpaddings: chex.Array,
#              blank_id: int = 0,
#              logepsilon: float = -1e5) -> chex.Array:
#   def update_phi_score(phi, added_score):
#   def loop_body(prev, x):


# Path: optax/_src/loss_test.py
from absl.testing import absltest
from absl.testing import parameterized
from optax._src import loss
import chex
import jax
import jax.numpy as jnp
import numpy as np

    self.ys = np.array([[10., 1., -2.], [1., 4., 0.2]], dtype=np.float32)
    self.ts = np.array([[0., 1.2, 0.2], [1., -0.3, 0.]], dtype=np.float32)
    # distance computed expected output from `scipy 1.20`.
    self.exp = np.array([0.9358251989, 1.0464068465], dtype=np.float32)

  @chex.all_variants()
  def test_scalar_distance(self):
    """Tests for a full batch."""
    np.testing.assert_allclose(
        self.variant(loss.cosine_distance)(self.ys[0], self.ts[0]),
        self.exp[0], atol=1e-4)

  @chex.all_variants()
  def test_scalar_similarity(self):
    """Tests for a full batch."""
    np.testing.assert_allclose(
        self.variant(loss.cosine_similarity)(self.ys[0], self.ts[0]),
        1. - self.exp[0], atol=1e-4)

  @chex.all_variants()
  def test_batched_distance(self):
    """Tests for a full batch."""
    np.testing.assert_allclose(
        self.variant(loss.cosine_distance)(self.ys, self.ts),
        self.exp, atol=1e-4)

  @chex.all_variants()
  def test_batched_similarity(self):
    """Tests for a full batch."""
    np.testing.assert_allclose(

        self.variant(loss.cosine_similarity)(self.ys, self.ts),