# Path: fedjax/datasets/cifar100.py
# SPLITS = ('train', 'test')
# _TFF_SQLITE_COMPRESSED_HEXDIGEST = '23d3916c9caa33395737ee560cc7cb77bbd05fc7b73647ee7be3a7e764172939'
# _TFF_SQLITE_COMPRESSED_NUM_BYTES = 153019920
# _FEDJAX_SQLITE_HEXDIGEST = {
#     'train': 'a4dc2f4ac4c9b6e7bd5234a1b568389384c00c903dcd34001c3cf50a4a81c713',
#     'test': '74416c6ee8f41b1086f0e0e3c5289ac6df5a641100df48e8b583411a55de891f'
# }
# _FEDJAX_SQLITE_NUM_BYTES = {'train': 140521472, 'test': 28135424}
# CIFAR100_PIXELS_MEAN = np.array([0.4914, 0.4822, 0.4465], dtype=np.float32)
# CIFAR100_PIXELS_INVERSE_STDDEV = (
#     1 / np.array([0.2023, 0.1994, 0.2010], dtype=np.float32))
# def cite():
# def _parse_tf_examples(vs: List[bytes]) -> client_datasets.Examples:
# def load_split(split: str,
#                mode: str = 'sqlite',
#                cache_dir: Optional[str] = None) -> federated_data.FederatedData:
# def load_data(
#     mode: str = 'sqlite',
#     cache_dir: Optional[str] = None
# ) -> Tuple[federated_data.FederatedData, federated_data.FederatedData]:
# def preprocess_client(
#     client_id: federated_data.ClientId,
#     examples: client_datasets.Examples) -> client_datasets.Examples:
# def preprocess_image(image: np.ndarray, is_train: bool) -> np.ndarray:
# def preprocess_batch(examples: client_datasets.Examples,
#                      is_train: bool) -> client_datasets.Examples:


from absl.testing import absltest
from fedjax.datasets import cifar100
import numpy as np
import numpy.testing as npt
            'label': label
        }), {
            'x': image,
            'y': label
        })

  def test_preprocess_image(self):
    # 4 images, each with one side being 1 after processing.
    image = np.zeros([4, 32, 32, 3], dtype=np.uint8)
    # [125, 123, 114] becomes nearly 0 after processing.
    image[:, :, :, :] = [125, 123, 114]
    # [177, 174, 165] becomes nearly 1 after processing.
    image[0, 0, :, :] = [177, 174, 165]
    image[1, -1, :, :] = [177, 174, 165]
    image[2, :, 0, :] = [177, 174, 165]
    image[3, :, -1, :] = [177, 174, 165]

    with self.subTest('is_train=False'):
      processed = cifar100.preprocess_image(image, is_train=False)
      self.assertEqual(processed.dtype, np.float32)
      ref = np.zeros([4, 32, 32, 3], dtype=np.float32)
      ref[0, 0, :, :] = 1
      ref[1, -1, :, :] = 1
      ref[2, :, 0, :] = 1
      ref[3, :, -1, :] = 1
      npt.assert_allclose(processed, ref, rtol=0.01, atol=0.01)

    with self.subTest('is_train=True'):
      # Expected crop start offsets: (2, 6); flip: yes.
      np.random.seed(4)

# Path: fedjax/datasets/cifar100.py
# SPLITS = ('train', 'test')
# _TFF_SQLITE_COMPRESSED_HEXDIGEST = '23d3916c9caa33395737ee560cc7cb77bbd05fc7b73647ee7be3a7e764172939'
# _TFF_SQLITE_COMPRESSED_NUM_BYTES = 153019920
# _FEDJAX_SQLITE_HEXDIGEST = {
#     'train': 'a4dc2f4ac4c9b6e7bd5234a1b568389384c00c903dcd34001c3cf50a4a81c713',
#     'test': '74416c6ee8f41b1086f0e0e3c5289ac6df5a641100df48e8b583411a55de891f'
# }
# _FEDJAX_SQLITE_NUM_BYTES = {'train': 140521472, 'test': 28135424}
# CIFAR100_PIXELS_MEAN = np.array([0.4914, 0.4822, 0.4465], dtype=np.float32)
# CIFAR100_PIXELS_INVERSE_STDDEV = (
#     1 / np.array([0.2023, 0.1994, 0.2010], dtype=np.float32))
# def cite():
# def _parse_tf_examples(vs: List[bytes]) -> client_datasets.Examples:
# def load_split(split: str,
#                mode: str = 'sqlite',
#                cache_dir: Optional[str] = None) -> federated_data.FederatedData:
# def load_data(
#     mode: str = 'sqlite',
#     cache_dir: Optional[str] = None
# ) -> Tuple[federated_data.FederatedData, federated_data.FederatedData]:
# def preprocess_client(
#     client_id: federated_data.ClientId,
#     examples: client_datasets.Examples) -> client_datasets.Examples:
# def preprocess_image(image: np.ndarray, is_train: bool) -> np.ndarray:
# def preprocess_batch(examples: client_datasets.Examples,
#                      is_train: bool) -> client_datasets.Examples:


# Path: fedjax/datasets/cifar100_test.py
from absl.testing import absltest
from fedjax.datasets import cifar100
import numpy as np
import numpy.testing as npt

            'label': label
        }), {
            'x': image,
            'y': label
        })

  def test_preprocess_image(self):
    # 4 images, each with one side being 1 after processing.
    image = np.zeros([4, 32, 32, 3], dtype=np.uint8)
    # [125, 123, 114] becomes nearly 0 after processing.
    image[:, :, :, :] = [125, 123, 114]
    # [177, 174, 165] becomes nearly 1 after processing.
    image[0, 0, :, :] = [177, 174, 165]
    image[1, -1, :, :] = [177, 174, 165]
    image[2, :, 0, :] = [177, 174, 165]
    image[3, :, -1, :] = [177, 174, 165]

    with self.subTest('is_train=False'):
      processed = cifar100.preprocess_image(image, is_train=False)
      self.assertEqual(processed.dtype, np.float32)
      ref = np.zeros([4, 32, 32, 3], dtype=np.float32)
      ref[0, 0, :, :] = 1
      ref[1, -1, :, :] = 1
      ref[2, :, 0, :] = 1
      ref[3, :, -1, :] = 1
      npt.assert_allclose(processed, ref, rtol=0.01, atol=0.01)

    with self.subTest('is_train=True'):
      # Expected crop start offsets: (2, 6); flip: yes.
      np.random.seed(4)

      processed = cifar100.preprocess_image(image, is_train=True)