# Path: init2winit/model_lib/normalization.py
# def _absolute_dims(rank, dims):
# def _get_batch_axis(
#     data_format,
#     x,
#     virtual_batch_size,
#     use_running_average,
#     axis_index_groups):
#   def __call__(self, x, use_running_average: Optional[bool] = None):
# class VirtualBatchNorm(nn.Module):


import functools
import jax
import jax.numpy as jnp
from typing import Any, Callable, Iterable
from absl import logging
from flax import linen as nn
from flax import optim
from init2winit.model_lib import normalization
from jax import lax
from jax.nn import initializers

  Args:
    normalizer: One of ['batch_norm', 'virtual_batch_norm', 'layer_norm',
      'none'].
    train: Boolean indiciating if we are running in train or inference mode
      for batch norm.
    batch_size: only used for virtual batch norm, the batch size.
    virtual_batch_size: only used for virtual batch norm, the virtual batch
      size.
    total_batch_size: only used for virtual batch norm when using gradient
      accumulation, the total batch size used to calculate gradients with.
    dtype: data type used for normalizer inputs and outputs.
    data_format: only used for virtual batch norm, used to determine the batch
      axis.

  Returns:
    A function that when called will create the normalizer module/function.

  Raises:
    ValueError if normalizer not recognized.
  """
  if normalizer == 'batch_norm':
    return functools.partial(
        nn.BatchNorm,
        use_running_average=not train,
        momentum=0.9,
        epsilon=1e-5,
        dtype=dtype)
  elif normalizer == 'virtual_batch_norm':
    return functools.partial(

# Path: init2winit/model_lib/normalization.py
# def _absolute_dims(rank, dims):
# def _get_batch_axis(
#     data_format,
#     x,
#     virtual_batch_size,
#     use_running_average,
#     axis_index_groups):
#   def __call__(self, x, use_running_average: Optional[bool] = None):
# class VirtualBatchNorm(nn.Module):


# Path: init2winit/model_lib/model_utils.py
import functools
import jax
import jax.numpy as jnp
from typing import Any, Callable, Iterable
from absl import logging
from flax import linen as nn
from flax import optim
from init2winit.model_lib import normalization
from jax import lax
from jax.nn import initializers


  Args:
    normalizer: One of ['batch_norm', 'virtual_batch_norm', 'layer_norm',
      'none'].
    train: Boolean indiciating if we are running in train or inference mode
      for batch norm.
    batch_size: only used for virtual batch norm, the batch size.
    virtual_batch_size: only used for virtual batch norm, the virtual batch
      size.
    total_batch_size: only used for virtual batch norm when using gradient
      accumulation, the total batch size used to calculate gradients with.
    dtype: data type used for normalizer inputs and outputs.
    data_format: only used for virtual batch norm, used to determine the batch
      axis.

  Returns:
    A function that when called will create the normalizer module/function.

  Raises:
    ValueError if normalizer not recognized.
  """
  if normalizer == 'batch_norm':
    return functools.partial(
        nn.BatchNorm,
        use_running_average=not train,
        momentum=0.9,
        epsilon=1e-5,
        dtype=dtype)
  elif normalizer == 'virtual_batch_norm':
    return functools.partial(

        normalization.VirtualBatchNorm,