# Path: fedjax/core/dataclasses.py
# def dataclass(clz: type):
#   def replace(self, **updates):
#   def iterate_clz(x):
#   def clz_from_iterable(meta, data):


import collections
import itertools
import numpy as np
from typing import Any, Callable, Iterable, Iterator, Mapping, Optional
from fedjax.core import dataclasses
#
# 1.  They can be easily passed around as a single object, instead of a list of
#     arguments.
# 2.  The list of argument is only defined, documented, and assigned default
#     values once (in the hyperparams class).
#
# This setup is more convenient for library code where the caller specifies such
# a hyperparams object, but less so for users directly calling the actual
# batching function (e.g. compare `padded_batch(batch_size=3)` vs
# `padded_batch(PaddedBatchHParams(batch_size=3))`. Therefore, all our batching
# functions support two ways of invocation:
# -   They can be passed in a hyperparams object.
# -   They can also be passed in a list of keyword arguments, which will be used
#     to construct the hyperparams object.
#
# See ClientDataset.padded_batch() for an example.


@dataclasses.dataclass
class PaddedBatchHParams:
  """See :meth:`ClientDataset.padded_batch`.

  Attributes:
    batch_size: Desired batch size.
    num_batch_size_buckets: Number of batch size buckets for the final batch.
  """
  batch_size: int
  num_batch_size_buckets: int = 1



# Path: fedjax/core/dataclasses.py
# def dataclass(clz: type):
#   def replace(self, **updates):
#   def iterate_clz(x):
#   def clz_from_iterable(meta, data):


# Path: fedjax/core/client_datasets.py
import collections
import itertools
import numpy as np
from typing import Any, Callable, Iterable, Iterator, Mapping, Optional
from fedjax.core import dataclasses

#
# 1.  They can be easily passed around as a single object, instead of a list of
#     arguments.
# 2.  The list of argument is only defined, documented, and assigned default
#     values once (in the hyperparams class).
#
# This setup is more convenient for library code where the caller specifies such
# a hyperparams object, but less so for users directly calling the actual
# batching function (e.g. compare `padded_batch(batch_size=3)` vs
# `padded_batch(PaddedBatchHParams(batch_size=3))`. Therefore, all our batching
# functions support two ways of invocation:
# -   They can be passed in a hyperparams object.
# -   They can also be passed in a list of keyword arguments, which will be used
#     to construct the hyperparams object.
#
# See ClientDataset.padded_batch() for an example.


@dataclasses.dataclass
class PaddedBatchHParams:
  """See :meth:`ClientDataset.padded_batch`.

  Attributes:
    batch_size: Desired batch size.
    num_batch_size_buckets: Number of batch size buckets for the final batch.
  """
  batch_size: int
  num_batch_size_buckets: int = 1



@dataclasses.dataclass