# Path: fedjax/core/client_datasets.py
# class BatchPreprocessor:
# class PaddedBatchHParams:
# class ShuffleRepeatBatchHParams:
# class BatchHParams:
# class ClientDataset:
# class PaddedBatchView:
# class ShuffleRepeatBatchView:
# class BatchView:
#   def __init__(self, fns: Iterable[Callable[[Examples], Examples]] = ()):
#   def __call__(self, examples: Examples) -> Examples:
#   def append(self, fn: Callable[[Examples], Examples]) -> 'BatchPreprocessor':
#   def __str__(self) -> str:
#   def __repr__(self) -> str:
#   def __init__(self,
#                raw_examples: Examples,
#                preprocessor: BatchPreprocessor = NoOpBatchPreprocessor):
#   def __len__(self) -> int:
#   def __getitem__(self, index: slice) -> 'ClientDataset':
#   def all_examples(self) -> Examples:
#   def padded_batch(self,
#                    hparams: Optional[PaddedBatchHParams] = None,
#                    **kwargs) -> Iterable[Examples]:
#   def shuffle_repeat_batch(self,
#                            hparams: Optional[ShuffleRepeatBatchHParams] = None,
#                            **kwargs) -> Iterable[Examples]:
#   def batch(self,
#             hparams: Optional[BatchHParams] = None,
#             **kwargs) -> Iterable[Examples]:
#   def __init__(self, client_dataset: ClientDataset,
#                hparams: PaddedBatchHParams):
#   def __iter__(self) -> Iterator[Examples]:
# def _pick_final_batch_size(data_size: int, batch_size: int,
#                            num_batch_size_buckets: int) -> int:
#   def __init__(self, client_dataset: ClientDataset,
#                hparams: ShuffleRepeatBatchHParams):
#   def __iter__(self) -> Iterator[Examples]:
#   def __init__(self, client_dataset: ClientDataset, hparams: BatchHParams):
#   def __iter__(self) -> Iterator[Examples]:
# def padded_batch_client_datasets(datasets: Iterable[ClientDataset],
#                                  hparams: Optional[PaddedBatchHParams] = None,
#                                  **kwargs) -> Iterator[Examples]:
# def buffered_shuffle(source: Iterable[Any], buffer_size: int,
#                      rng: np.random.RandomState) -> Iterator[Any]:
# def buffered_shuffle_batch_client_datasets(
#     datasets: Iterable[ClientDataset], batch_size: int, buffer_size: int,
#     rng: np.random.RandomState) -> Iterator[Examples]:
#   def gen_items():
# def assert_consistent_rows(examples: Examples):
# def num_examples(examples: Examples, validate: bool = True) -> int:
# def slice_examples(examples: Examples, index: slice) -> Examples:
# def concat_examples(many_examples: Iterable[Examples]) -> Examples:
# def attach_mask(examples: Examples, mask: np.ndarray) -> Examples:
# def pad_examples(examples: Examples, size: int) -> Examples:
# EXAMPLE_MASK_KEY = '__mask__'


import itertools
import numpy as np
import numpy.testing as npt
from absl.testing import absltest
from fedjax.core import client_datasets

  def test_slice_examples(self):
    examples = {'a': np.arange(10), 'b': np.arange(20).reshape([10, 2])}

    with self.subTest('Slicing like [3:]'):
      sliced = client_datasets.slice_examples(examples, slice(3))
      npt.assert_equal(sliced, {'a': [0, 1, 2], 'b': [[0, 1], [2, 3], [4, 5]]})

    with self.subTest('Slicing like [3:6]'):
      sliced = client_datasets.slice_examples(examples, slice(3, 6))
      npt.assert_equal(sliced, {
          'a': [3, 4, 5],
          'b': [[6, 7], [8, 9], [10, 11]]
      })

    with self.subTest('Slicing like [3:6:2]'):
      sliced = client_datasets.slice_examples(examples, slice(3, 6, 2))
      npt.assert_equal(sliced, {'a': [3, 5], 'b': [[6, 7], [10, 11]]})

  def test_attach_mask(self):
    examples = {'a': np.arange(5), 'b': np.arange(10).reshape([5, 2])}
    mask = np.array([True, False, True, False, False])
    with self.subTest('mask key already used'):
      with self.assertRaises(ValueError):
        client_datasets.attach_mask(
            {
                **examples, client_datasets.EXAMPLE_MASK_KEY: None
            }, mask)
    with self.subTest('success'):
      npt.assert_equal(

# Path: fedjax/core/client_datasets.py
# class BatchPreprocessor:
# class PaddedBatchHParams:
# class ShuffleRepeatBatchHParams:
# class BatchHParams:
# class ClientDataset:
# class PaddedBatchView:
# class ShuffleRepeatBatchView:
# class BatchView:
#   def __init__(self, fns: Iterable[Callable[[Examples], Examples]] = ()):
#   def __call__(self, examples: Examples) -> Examples:
#   def append(self, fn: Callable[[Examples], Examples]) -> 'BatchPreprocessor':
#   def __str__(self) -> str:
#   def __repr__(self) -> str:
#   def __init__(self,
#                raw_examples: Examples,
#                preprocessor: BatchPreprocessor = NoOpBatchPreprocessor):
#   def __len__(self) -> int:
#   def __getitem__(self, index: slice) -> 'ClientDataset':
#   def all_examples(self) -> Examples:
#   def padded_batch(self,
#                    hparams: Optional[PaddedBatchHParams] = None,
#                    **kwargs) -> Iterable[Examples]:
#   def shuffle_repeat_batch(self,
#                            hparams: Optional[ShuffleRepeatBatchHParams] = None,
#                            **kwargs) -> Iterable[Examples]:
#   def batch(self,
#             hparams: Optional[BatchHParams] = None,
#             **kwargs) -> Iterable[Examples]:
#   def __init__(self, client_dataset: ClientDataset,
#                hparams: PaddedBatchHParams):
#   def __iter__(self) -> Iterator[Examples]:
# def _pick_final_batch_size(data_size: int, batch_size: int,
#                            num_batch_size_buckets: int) -> int:
#   def __init__(self, client_dataset: ClientDataset,
#                hparams: ShuffleRepeatBatchHParams):
#   def __iter__(self) -> Iterator[Examples]:
#   def __init__(self, client_dataset: ClientDataset, hparams: BatchHParams):
#   def __iter__(self) -> Iterator[Examples]:
# def padded_batch_client_datasets(datasets: Iterable[ClientDataset],
#                                  hparams: Optional[PaddedBatchHParams] = None,
#                                  **kwargs) -> Iterator[Examples]:
# def buffered_shuffle(source: Iterable[Any], buffer_size: int,
#                      rng: np.random.RandomState) -> Iterator[Any]:
# def buffered_shuffle_batch_client_datasets(
#     datasets: Iterable[ClientDataset], batch_size: int, buffer_size: int,
#     rng: np.random.RandomState) -> Iterator[Examples]:
#   def gen_items():
# def assert_consistent_rows(examples: Examples):
# def num_examples(examples: Examples, validate: bool = True) -> int:
# def slice_examples(examples: Examples, index: slice) -> Examples:
# def concat_examples(many_examples: Iterable[Examples]) -> Examples:
# def attach_mask(examples: Examples, mask: np.ndarray) -> Examples:
# def pad_examples(examples: Examples, size: int) -> Examples:
# EXAMPLE_MASK_KEY = '__mask__'


# Path: fedjax/core/client_datasets_test.py
import itertools
import numpy as np
import numpy.testing as npt
from absl.testing import absltest
from fedjax.core import client_datasets


  def test_slice_examples(self):
    examples = {'a': np.arange(10), 'b': np.arange(20).reshape([10, 2])}

    with self.subTest('Slicing like [3:]'):
      sliced = client_datasets.slice_examples(examples, slice(3))
      npt.assert_equal(sliced, {'a': [0, 1, 2], 'b': [[0, 1], [2, 3], [4, 5]]})

    with self.subTest('Slicing like [3:6]'):
      sliced = client_datasets.slice_examples(examples, slice(3, 6))
      npt.assert_equal(sliced, {
          'a': [3, 4, 5],
          'b': [[6, 7], [8, 9], [10, 11]]
      })

    with self.subTest('Slicing like [3:6:2]'):
      sliced = client_datasets.slice_examples(examples, slice(3, 6, 2))
      npt.assert_equal(sliced, {'a': [3, 5], 'b': [[6, 7], [10, 11]]})

  def test_attach_mask(self):
    examples = {'a': np.arange(5), 'b': np.arange(10).reshape([5, 2])}
    mask = np.array([True, False, True, False, False])
    with self.subTest('mask key already used'):
      with self.assertRaises(ValueError):
        client_datasets.attach_mask(
            {
                **examples, client_datasets.EXAMPLE_MASK_KEY: None
            }, mask)
    with self.subTest('success'):
      npt.assert_equal(

          client_datasets.attach_mask(examples, mask), {