# 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
    ...  # possibly modify the default hps.
    model = model(hps, dataset.meta_data)
    model_init_fn = jax.jit(functools.partial(flax_module.init, train=False))
    init_dict = model_init_fn(
        {'params': params_rng, 'dropout': dropout_rng},
        fake_input_batch)
    # Trainable model parameters.
    params = init_dict['params']
    batch_stats = init_dict.get('batch_stats', {})

    # this is how to call the model fprop.
    logits, vars = flax_module.apply(
        {'params': params, 'batch_stats': batch_stats},
        batch,
        mutable=['batch_stats'],
        rngs={'dropout': dropout_rng},
        train=True))
    new_batch_stats = vars['batch_stats']
  ```

  Note for models without batch norm, `flax_module.apply` will only return a
  single value (logits). See the Flax docs for more info:
  https://flax.readthedocs.io/en/latest/flax.linen.html#flax.linen.Module.apply.
  """

  def __init__(self, hps, dataset_meta_data, loss_name, metrics_name):
    self.hps = hps
    self.dataset_meta_data = dataset_meta_data
    self.loss_fn = losses.get_loss_fn(loss_name)
    self.output_activation_fn = losses.get_output_activation_fn(loss_name)

# 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

    ...  # possibly modify the default hps.
    model = model(hps, dataset.meta_data)
    model_init_fn = jax.jit(functools.partial(flax_module.init, train=False))
    init_dict = model_init_fn(
        {'params': params_rng, 'dropout': dropout_rng},
        fake_input_batch)
    # Trainable model parameters.
    params = init_dict['params']
    batch_stats = init_dict.get('batch_stats', {})

    # this is how to call the model fprop.
    logits, vars = flax_module.apply(
        {'params': params, 'batch_stats': batch_stats},
        batch,
        mutable=['batch_stats'],
        rngs={'dropout': dropout_rng},
        train=True))
    new_batch_stats = vars['batch_stats']
  ```

  Note for models without batch norm, `flax_module.apply` will only return a
  single value (logits). See the Flax docs for more info:
  https://flax.readthedocs.io/en/latest/flax.linen.html#flax.linen.Module.apply.
  """

  def __init__(self, hps, dataset_meta_data, loss_name, metrics_name):
    self.hps = hps
    self.dataset_meta_data = dataset_meta_data
    self.loss_fn = losses.get_loss_fn(loss_name)
    self.output_activation_fn = losses.get_output_activation_fn(loss_name)

    self.metrics_bundle = metrics.get_metrics(metrics_name)