# Path: init2winit/dataset_lib/data_utils.py
# def iterator_as_numpy(iterator):
# def image_iterator(data,
#                    rescale,
#                    output_shape,
#                    is_one_hot,
#                    autoencoder,
#                    shuffle_rng=None,
#                    augment_fn=None,
#                    include_example_keys=False):
# def maybe_pad_batch(batch,
#                     desired_batch_size,
#                     data_format=None,
#                     mask_key=None):
#   def zero_pad(ar, pad_axis):
# def shard(batch, n_devices=None):
#   def _shard_array(array):
# def tf_to_numpy(tfds_data):


from absl.testing import absltest
from absl.testing import parameterized
from init2winit.dataset_lib import data_utils
import numpy as np

class DataUtilsTest(parameterized.TestCase):
  """Unit tests for datasets.py."""

  @parameterized.named_parameters(*test_parameters)
  def test_padding(self, image_format, batch_axis, input_shape):
    """Test that the shape is the expected padded shape."""
    batch = {'inputs': np.ones(input_shape)}
    padded_batch = data_utils.maybe_pad_batch(
        batch, desired_batch_size, image_format)
    expected_shapes = list(input_shape)
    expected_shapes[batch_axis] = desired_batch_size
    self.assertEqual(padded_batch['inputs'].shape, tuple(expected_shapes))
    self.assertEqual(padded_batch['weights'].shape, (desired_batch_size,))

  def test_padding_seq2seq(self):
    """Test padding for sequence-to-sequence models."""
    input_len_max = 25
    input_len_true = 22  # true input_seq_length for each example in batch.
    target_len_max = 25
    target_len_true = 21  # true target_seq_length for each example in batch.

    inputs_shape = (batch_size, input_len_max)
    targets_shape = (batch_size, target_len_max)
    batch = {'inputs': np.ones(inputs_shape), 'targets': np.ones(targets_shape)}
    batch['inputs'][:, input_len_true:] = 0  # zero-pad extra inputs tokens
    batch['targets'][:, target_len_true:] = 0  # zero-pad extra targets tokens
    expected_inputs_shape = (desired_batch_size, input_len_max)
    expected_targets_shape = (desired_batch_size, target_len_max)
    expected_weights_shape = (desired_batch_size, target_len_max)

# Path: init2winit/dataset_lib/data_utils.py
# def iterator_as_numpy(iterator):
# def image_iterator(data,
#                    rescale,
#                    output_shape,
#                    is_one_hot,
#                    autoencoder,
#                    shuffle_rng=None,
#                    augment_fn=None,
#                    include_example_keys=False):
# def maybe_pad_batch(batch,
#                     desired_batch_size,
#                     data_format=None,
#                     mask_key=None):
#   def zero_pad(ar, pad_axis):
# def shard(batch, n_devices=None):
#   def _shard_array(array):
# def tf_to_numpy(tfds_data):


# Path: init2winit/dataset_lib/test_data_utils.py
from absl.testing import absltest
from absl.testing import parameterized
from init2winit.dataset_lib import data_utils
import numpy as np


class DataUtilsTest(parameterized.TestCase):
  """Unit tests for datasets.py."""

  @parameterized.named_parameters(*test_parameters)
  def test_padding(self, image_format, batch_axis, input_shape):
    """Test that the shape is the expected padded shape."""
    batch = {'inputs': np.ones(input_shape)}
    padded_batch = data_utils.maybe_pad_batch(
        batch, desired_batch_size, image_format)
    expected_shapes = list(input_shape)
    expected_shapes[batch_axis] = desired_batch_size
    self.assertEqual(padded_batch['inputs'].shape, tuple(expected_shapes))
    self.assertEqual(padded_batch['weights'].shape, (desired_batch_size,))

  def test_padding_seq2seq(self):
    """Test padding for sequence-to-sequence models."""
    input_len_max = 25
    input_len_true = 22  # true input_seq_length for each example in batch.
    target_len_max = 25
    target_len_true = 21  # true target_seq_length for each example in batch.

    inputs_shape = (batch_size, input_len_max)
    targets_shape = (batch_size, target_len_max)
    batch = {'inputs': np.ones(inputs_shape), 'targets': np.ones(targets_shape)}
    batch['inputs'][:, input_len_true:] = 0  # zero-pad extra inputs tokens
    batch['targets'][:, target_len_true:] = 0  # zero-pad extra targets tokens
    expected_inputs_shape = (desired_batch_size, input_len_max)
    expected_targets_shape = (desired_batch_size, target_len_max)
    expected_weights_shape = (desired_batch_size, target_len_max)

    padded_batch = data_utils.maybe_pad_batch(