# 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
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Unit tests for datasets.py."""


desired_batch_size = 23
batch_axes = [0, 0, 3, 2]
test_names = ['default', 'NHWC', 'HWCN', 'HWNC']
image_formats = [None, 'NHWC', 'HWCN', 'HWNC']
batch_size = 13
width = 11
num_channels = 3
input_shapes = [
    (batch_size, width, width, num_channels),
    (batch_size, width, width, num_channels),
    (width, width, num_channels, batch_size),
    (width, width, batch_size, num_channels),
]
test_parameters = zip(test_names, image_formats, batch_axes, input_shapes)


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)}

# 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

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Unit tests for datasets.py."""


desired_batch_size = 23
batch_axes = [0, 0, 3, 2]
test_names = ['default', 'NHWC', 'HWCN', 'HWNC']
image_formats = [None, 'NHWC', 'HWCN', 'HWNC']
batch_size = 13
width = 11
num_channels = 3
input_shapes = [
    (batch_size, width, width, num_channels),
    (batch_size, width, width, num_channels),
    (width, width, num_channels, batch_size),
    (width, width, batch_size, num_channels),
]
test_parameters = zip(test_names, image_formats, batch_axes, input_shapes)


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(