# Path: init2winit/model_lib/losses.py
# def unnormalized_sigmoid_binary_cross_entropy(logits, targets, weights=None):
# def sigmoid_binary_cross_entropy(logits, targets, weights=None):
# def unnormalized_sigmoid_mean_squared_error(logits, targets, weights=None):
# def sigmoid_mean_squared_error(logits, targets, weights=None):
# def weighted_unnormalized_cross_entropy(logits, targets, weights=None):
# def weighted_cross_entropy(logits, targets, weights=None):
# def ctc_loss(logits, logit_paddings, labels, label_paddings, blank_id=0):
# def get_loss_fn(loss_name):
# def get_output_activation_fn(loss_name):
#   def compute(self):
# _ALL_LOSS_FUNCTIONS = {
#     'sigmoid_mean_squared_error': (sigmoid_mean_squared_error, jax.nn.sigmoid),
#     'sigmoid_binary_cross_entropy':
#         (sigmoid_binary_cross_entropy, jax.nn.sigmoid),
#     'cross_entropy': (weighted_cross_entropy, jax.nn.softmax),
#     'ctc': (ctc_loss, jax.nn.log_softmax)
# }
# class MeanAveragePrecision(
#     metrics.CollectingMetric.from_outputs(('logits', 'targets', 'weights'))):


import types
import jax
import numpy as np
from absl.testing import absltest
from absl.testing import parameterized
from init2winit.model_lib import losses
    'logits':
        np.array([[4, 7], [-2, 5], [8, 6], [-10, -4], [3, -5]]).astype(float),
    'targets':
        np.array([[1, 0], [0, 1], [1, 0], [1, 0], [0, 1]]),
    'weights':
        np.array([2, 0, 0, 6, 1]),
}]

CLASSIFICATION_KEYS = [
    (loss_name, loss_name) for loss_name in CLASSIFICATION_LOSSES
]
RECONSTRUCTION_KEYS = [
    (loss_name, loss_name) for loss_name in RECONSTRUCTION_LOSSES
]


class LossesTest(parameterized.TestCase):
  """Tests for losses.py."""

  def test_loss_fn_registry(self):
    for loss_name in losses._ALL_LOSS_FUNCTIONS:  # pylint: disable=protected-access
      loss_fn = losses.get_loss_fn(loss_name)
      self.assertIsInstance(loss_fn, types.FunctionType)
    with self.assertRaises(ValueError):
      losses.get_loss_fn('__test__loss__name__')

  def test_output_activation_fn_registry(self):
    activation_fn = losses.get_output_activation_fn('cross_entropy')
    self.assertEqual(activation_fn.__name__, 'softmax')
    with self.assertRaises(ValueError):

# Path: init2winit/model_lib/losses.py
# def unnormalized_sigmoid_binary_cross_entropy(logits, targets, weights=None):
# def sigmoid_binary_cross_entropy(logits, targets, weights=None):
# def unnormalized_sigmoid_mean_squared_error(logits, targets, weights=None):
# def sigmoid_mean_squared_error(logits, targets, weights=None):
# def weighted_unnormalized_cross_entropy(logits, targets, weights=None):
# def weighted_cross_entropy(logits, targets, weights=None):
# def ctc_loss(logits, logit_paddings, labels, label_paddings, blank_id=0):
# def get_loss_fn(loss_name):
# def get_output_activation_fn(loss_name):
#   def compute(self):
# _ALL_LOSS_FUNCTIONS = {
#     'sigmoid_mean_squared_error': (sigmoid_mean_squared_error, jax.nn.sigmoid),
#     'sigmoid_binary_cross_entropy':
#         (sigmoid_binary_cross_entropy, jax.nn.sigmoid),
#     'cross_entropy': (weighted_cross_entropy, jax.nn.softmax),
#     'ctc': (ctc_loss, jax.nn.log_softmax)
# }
# class MeanAveragePrecision(
#     metrics.CollectingMetric.from_outputs(('logits', 'targets', 'weights'))):


# Path: init2winit/model_lib/test_losses.py
import types
import jax
import numpy as np
from absl.testing import absltest
from absl.testing import parameterized
from init2winit.model_lib import losses

    'logits':
        np.array([[4, 7], [-2, 5], [8, 6], [-10, -4], [3, -5]]).astype(float),
    'targets':
        np.array([[1, 0], [0, 1], [1, 0], [1, 0], [0, 1]]),
    'weights':
        np.array([2, 0, 0, 6, 1]),
}]

CLASSIFICATION_KEYS = [
    (loss_name, loss_name) for loss_name in CLASSIFICATION_LOSSES
]
RECONSTRUCTION_KEYS = [
    (loss_name, loss_name) for loss_name in RECONSTRUCTION_LOSSES
]


class LossesTest(parameterized.TestCase):
  """Tests for losses.py."""

  def test_loss_fn_registry(self):
    for loss_name in losses._ALL_LOSS_FUNCTIONS:  # pylint: disable=protected-access
      loss_fn = losses.get_loss_fn(loss_name)
      self.assertIsInstance(loss_fn, types.FunctionType)
    with self.assertRaises(ValueError):
      losses.get_loss_fn('__test__loss__name__')

  def test_output_activation_fn_registry(self):
    activation_fn = losses.get_output_activation_fn('cross_entropy')
    self.assertEqual(activation_fn.__name__, 'softmax')
    with self.assertRaises(ValueError):

      losses.get_output_activation_fn('__test__loss__name__')