# 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 abc
import numpy as np
from typing import Any, Callable, Iterable, Iterator, Optional, Tuple
from fedjax.core import client_datasets
    example_buffer_size: int,
    seed: Optional[int] = None) -> Iterator[client_datasets.Examples]:
  """Shuffle-repeat-batch all client datasets in a federated dataset for training a centralized baseline.

  Shuffling is done using two levels of buffered shuffling, first at the client
  level, then at the example level.

  This produces an infinite stream of batches. itertools.islice() can be used to
  cap the number of batches, if so desired.

  Args:
    fd: Federated dataset.
    batch_size: Desired batch size.
    client_buffer_size: Buffer size for client level shuffling.
    example_buffer_size: Buffer size for example level shuffling.
    seed: Optional RNG seed.

  Yields:
    Batches of preprocessed examples.
  """
  rng = np.random.RandomState(seed)
  datasets = (client_dataset for _, client_dataset in fd.shuffled_clients(
      client_buffer_size, rng.randint(1 << 32)))
  yield from client_datasets.buffered_shuffle_batch_client_datasets(
      datasets, batch_size=batch_size, buffer_size=example_buffer_size, rng=rng)


def padded_batch_federated_data(fd: FederatedData,
                                hparams: Optional[
                                    client_datasets.PaddedBatchHParams] = None,

# 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/federated_data.py
import abc
import numpy as np
from typing import Any, Callable, Iterable, Iterator, Optional, Tuple
from fedjax.core import client_datasets

    example_buffer_size: int,
    seed: Optional[int] = None) -> Iterator[client_datasets.Examples]:
  """Shuffle-repeat-batch all client datasets in a federated dataset for training a centralized baseline.

  Shuffling is done using two levels of buffered shuffling, first at the client
  level, then at the example level.

  This produces an infinite stream of batches. itertools.islice() can be used to
  cap the number of batches, if so desired.

  Args:
    fd: Federated dataset.
    batch_size: Desired batch size.
    client_buffer_size: Buffer size for client level shuffling.
    example_buffer_size: Buffer size for example level shuffling.
    seed: Optional RNG seed.

  Yields:
    Batches of preprocessed examples.
  """
  rng = np.random.RandomState(seed)
  datasets = (client_dataset for _, client_dataset in fd.shuffled_clients(
      client_buffer_size, rng.randint(1 << 32)))
  yield from client_datasets.buffered_shuffle_batch_client_datasets(
      datasets, batch_size=batch_size, buffer_size=example_buffer_size, rng=rng)


def padded_batch_federated_data(fd: FederatedData,
                                hparams: Optional[
                                    client_datasets.PaddedBatchHParams] = None,

                                **kwargs) -> Iterator[client_datasets.Examples]: