# 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'))):


from clu import metrics
from init2winit.model_lib import losses
import flax
import jax.numpy as jnp
import numpy as np
   logits: [batch, ..., num_classes] float array.
   targets: one hot vector of shape [batch, ..., num_classes].
   weights: None or array of shape [batch x ...] (rank of one_hot_targets -1).

  Returns:
    Binary vector indicated which examples are misclassified ([batch, ...]).
  """
  if logits.ndim != targets.ndim:
    raise ValueError('Incorrect shapes. Got shape %s logits and %s targets' %
                     (str(logits.shape), str(targets.shape)))
  preds = jnp.argmax(logits, axis=-1)
  pred_targets = jnp.argmax(targets, axis=-1)
  loss = jnp.not_equal(preds, pred_targets)
  if weights is not None:
    if weights.ndim != targets.ndim - 1:
      raise ValueError(
          'Incorrect shapes. Got shape %s weights and %s one_hot_targets' %
          (str(weights.shape), str(targets.shape)))
    loss = loss * weights
  return loss


# All metrics used here must take three arguments named `logits`, `targets`,
# `weights`. We don't use CLU's `mask` argument, the metric gets
# that information from `weights`. The total weight for calculating the average
# is `weights.sum()`.
_METRICS = {
    'binary_autoencoder_metrics':
        metrics.Collection.create(
            sigmoid_binary_cross_entropy=weighted_average_metric(

# 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/metrics.py
from clu import metrics
from init2winit.model_lib import losses
import flax
import jax.numpy as jnp
import numpy as np

   logits: [batch, ..., num_classes] float array.
   targets: one hot vector of shape [batch, ..., num_classes].
   weights: None or array of shape [batch x ...] (rank of one_hot_targets -1).

  Returns:
    Binary vector indicated which examples are misclassified ([batch, ...]).
  """
  if logits.ndim != targets.ndim:
    raise ValueError('Incorrect shapes. Got shape %s logits and %s targets' %
                     (str(logits.shape), str(targets.shape)))
  preds = jnp.argmax(logits, axis=-1)
  pred_targets = jnp.argmax(targets, axis=-1)
  loss = jnp.not_equal(preds, pred_targets)
  if weights is not None:
    if weights.ndim != targets.ndim - 1:
      raise ValueError(
          'Incorrect shapes. Got shape %s weights and %s one_hot_targets' %
          (str(weights.shape), str(targets.shape)))
    loss = loss * weights
  return loss


# All metrics used here must take three arguments named `logits`, `targets`,
# `weights`. We don't use CLU's `mask` argument, the metric gets
# that information from `weights`. The total weight for calculating the average
# is `weights.sum()`.
_METRICS = {
    'binary_autoencoder_metrics':
        metrics.Collection.create(
            sigmoid_binary_cross_entropy=weighted_average_metric(

                losses.unnormalized_sigmoid_binary_cross_entropy),