# 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
import numpy as np
from absl.testing import absltest
from flax import linen as nn
from init2winit.model_lib import normalization
  model_init_fn = jax.jit(
      functools.partial(flax_module.init, use_running_average=False))
  xs = np.zeros(input_shape)
  init_dict = model_init_fn({'params': rng}, xs)
  params = init_dict['params']
  batch_stats = init_dict['batch_stats']
  return params, batch_stats


class NormalizationTest(absltest.TestCase):
  """Tests for virtual batch norm.

  We test gradient accumulation behavior in
  optimizer_lib/test_gradient_accumulator.py.
  """

  def test_batch_norm(self):
    """Test virtual BN recovers BN when the virtual size equals batch size."""
    rng = jax.random.PRNGKey(0)
    batch_size = 10
    feature_size = 7
    input_shape = (batch_size, 3, 3, feature_size)
    half_input_shape = (batch_size // 2, 3, 3, feature_size)
    twos = 2.0 * jnp.ones(half_input_shape)
    nines = 9.0 * jnp.ones(half_input_shape)
    x = jnp.concatenate((twos, nines))

    bn_flax_module = nn.BatchNorm(momentum=0.9)
    bn_params, bn_state = _init(bn_flax_module, rng, input_shape)


# 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/test_normalization.py
import functools
import jax
import jax.numpy as jnp
import numpy as np
from absl.testing import absltest
from flax import linen as nn
from init2winit.model_lib import normalization

  model_init_fn = jax.jit(
      functools.partial(flax_module.init, use_running_average=False))
  xs = np.zeros(input_shape)
  init_dict = model_init_fn({'params': rng}, xs)
  params = init_dict['params']
  batch_stats = init_dict['batch_stats']
  return params, batch_stats


class NormalizationTest(absltest.TestCase):
  """Tests for virtual batch norm.

  We test gradient accumulation behavior in
  optimizer_lib/test_gradient_accumulator.py.
  """

  def test_batch_norm(self):
    """Test virtual BN recovers BN when the virtual size equals batch size."""
    rng = jax.random.PRNGKey(0)
    batch_size = 10
    feature_size = 7
    input_shape = (batch_size, 3, 3, feature_size)
    half_input_shape = (batch_size // 2, 3, 3, feature_size)
    twos = 2.0 * jnp.ones(half_input_shape)
    nines = 9.0 * jnp.ones(half_input_shape)
    x = jnp.concatenate((twos, nines))

    bn_flax_module = nn.BatchNorm(momentum=0.9)
    bn_params, bn_state = _init(bn_flax_module, rng, input_shape)


    vbn_flax_module = normalization.VirtualBatchNorm(