# 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
      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),
            sigmoid_mean_squared_error=weighted_average_metric(
                losses.unnormalized_sigmoid_mean_squared_error),
            num_examples=NumExamples),
    'classification_metrics':
        metrics.Collection.create(
            error_rate=weighted_average_metric(weighted_misclassifications),
            ce_loss=weighted_average_metric(
                losses.weighted_unnormalized_cross_entropy),
            num_examples=NumExamples),
    'binary_classification_metrics':
        metrics.Collection.create(
            ce_loss=weighted_average_metric(
                losses.sigmoid_binary_cross_entropy),
            num_examples=NumExamples,

# 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

      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),
            sigmoid_mean_squared_error=weighted_average_metric(
                losses.unnormalized_sigmoid_mean_squared_error),
            num_examples=NumExamples),
    'classification_metrics':
        metrics.Collection.create(
            error_rate=weighted_average_metric(weighted_misclassifications),
            ce_loss=weighted_average_metric(
                losses.weighted_unnormalized_cross_entropy),
            num_examples=NumExamples),
    'binary_classification_metrics':
        metrics.Collection.create(
            ce_loss=weighted_average_metric(
                losses.sigmoid_binary_cross_entropy),
            num_examples=NumExamples,

            average_precision=losses.MeanAveragePrecision)