# 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
# class NumExamples(metrics.Metric):
#   class _Metric(metrics.Metric):
#   def from_model_output(cls, logits, targets, weights, **_):
#   def merge(self, other):
#   def compute(self):
# def weighted_average_metric(fun):
#     def from_model_output(cls, logits, targets, weights, **_):
#     def merge(self, other):
#     def compute(self):
# def weighted_misclassifications(logits, targets, weights=None):
# def get_metrics(metrics_name):
# _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)
# }
# 
# Path: init2winit/model_lib/model_utils.py
# ACTIVATIONS = {
#     'relu': jax.nn.relu,
#     'id': lambda x: x,
#     'tanh': jnp.tanh,
#     'sigmoid': jax.nn.sigmoid,
#     'silu': jax.nn.silu,
#     'leaky_relu': jax.nn.leaky_relu,
# }
# INITIALIZERS = {
#     'delta_orthogonal': initializers.delta_orthogonal,
#     'orthogonal': initializers.orthogonal,
#     'lecun_normal': lecun_normal,
#     'xavier_uniform': initializers.xavier_uniform,
# }
# class ScalarMultiply(nn.Module):
#   def __call__(self, x):
# def get_normalizer(normalizer,
#                    train,
#                    batch_size=None,
#                    virtual_batch_size=None,
#                    total_batch_size=None,
#                    dtype=jnp.float32,
#                    data_format='NHWC'):
#     def identity_wrapper(*args, **kwargs):
#       def identity(x, *args, **kwargs):
# def apply_label_smoothing(one_hot_targets, label_smoothing):
# def sync_batchnorm_stats(state):
# def cross_device_avg(pytree):
# def l2_regularization(params, l2_decay_rank_threshold):
# def flatten_dict(nested_dict, sep='/'):
# def rescale_layers(params, layer_rescale_factors):


from init2winit.model_lib import losses
from init2winit.model_lib import metrics
from init2winit.model_lib import model_utils
from jax.nn import one_hot
import jax.numpy as jnp
                    apply_output_activation_fn=False):
    """Returns predictions from all the model outputs on the given batch."""
    if apply_output_activation_fn:
      return _predict_batch(self.flax_module, batch_stats, batch,
                            self.output_activation_fn)
    else:
      return _predict_batch(self.flax_module, params, batch_stats, batch)

  def training_cost(self, params, batch, batch_stats=None, dropout_rng=None):
    """Return loss with an optional L2 penalty on the weights."""
    all_variables = {'params': params}
    apply_kwargs = {'train': True}

    if batch_stats is not None:
      all_variables['batch_stats'] = batch_stats
      apply_kwargs['mutable'] = ['batch_stats']

    if dropout_rng is not None:
      apply_kwargs['rngs'] = {'dropout': dropout_rng}

    # For more information on flax.linen.Module.apply, see the docs at
    # https://flax.readthedocs.io/en/latest/flax.linen.html#flax.linen.Module.apply.
    logits, new_batch_stats = self.flax_module.apply(
        all_variables, batch['inputs'], **apply_kwargs)
    weights = batch.get('weights')
    targets = batch['targets']
    if self.dataset_meta_data['apply_one_hot_in_loss']:
      targets = one_hot(targets, logits.shape[-1])
    # Optionally apply label smoothing.
    if self.hps.get('label_smoothing') is not None:

# 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
# class NumExamples(metrics.Metric):
#   class _Metric(metrics.Metric):
#   def from_model_output(cls, logits, targets, weights, **_):
#   def merge(self, other):
#   def compute(self):
# def weighted_average_metric(fun):
#     def from_model_output(cls, logits, targets, weights, **_):
#     def merge(self, other):
#     def compute(self):
# def weighted_misclassifications(logits, targets, weights=None):
# def get_metrics(metrics_name):
# _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)
# }
# 
# Path: init2winit/model_lib/model_utils.py
# ACTIVATIONS = {
#     'relu': jax.nn.relu,
#     'id': lambda x: x,
#     'tanh': jnp.tanh,
#     'sigmoid': jax.nn.sigmoid,
#     'silu': jax.nn.silu,
#     'leaky_relu': jax.nn.leaky_relu,
# }
# INITIALIZERS = {
#     'delta_orthogonal': initializers.delta_orthogonal,
#     'orthogonal': initializers.orthogonal,
#     'lecun_normal': lecun_normal,
#     'xavier_uniform': initializers.xavier_uniform,
# }
# class ScalarMultiply(nn.Module):
#   def __call__(self, x):
# def get_normalizer(normalizer,
#                    train,
#                    batch_size=None,
#                    virtual_batch_size=None,
#                    total_batch_size=None,
#                    dtype=jnp.float32,
#                    data_format='NHWC'):
#     def identity_wrapper(*args, **kwargs):
#       def identity(x, *args, **kwargs):
# def apply_label_smoothing(one_hot_targets, label_smoothing):
# def sync_batchnorm_stats(state):
# def cross_device_avg(pytree):
# def l2_regularization(params, l2_decay_rank_threshold):
# def flatten_dict(nested_dict, sep='/'):
# def rescale_layers(params, layer_rescale_factors):


# Path: init2winit/model_lib/base_model.py
from init2winit.model_lib import losses
from init2winit.model_lib import metrics
from init2winit.model_lib import model_utils
from jax.nn import one_hot
import jax.numpy as jnp

                    apply_output_activation_fn=False):
    """Returns predictions from all the model outputs on the given batch."""
    if apply_output_activation_fn:
      return _predict_batch(self.flax_module, batch_stats, batch,
                            self.output_activation_fn)
    else:
      return _predict_batch(self.flax_module, params, batch_stats, batch)

  def training_cost(self, params, batch, batch_stats=None, dropout_rng=None):
    """Return loss with an optional L2 penalty on the weights."""
    all_variables = {'params': params}
    apply_kwargs = {'train': True}

    if batch_stats is not None:
      all_variables['batch_stats'] = batch_stats
      apply_kwargs['mutable'] = ['batch_stats']

    if dropout_rng is not None:
      apply_kwargs['rngs'] = {'dropout': dropout_rng}

    # For more information on flax.linen.Module.apply, see the docs at
    # https://flax.readthedocs.io/en/latest/flax.linen.html#flax.linen.Module.apply.
    logits, new_batch_stats = self.flax_module.apply(
        all_variables, batch['inputs'], **apply_kwargs)
    weights = batch.get('weights')
    targets = batch['targets']
    if self.dataset_meta_data['apply_one_hot_in_loss']:
      targets = one_hot(targets, logits.shape[-1])
    # Optionally apply label smoothing.
    if self.hps.get('label_smoothing') is not None:

      targets = model_utils.apply_label_smoothing(