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__' def dataclass(clz: type):
  def replace(self, **updates):
  def iterate_clz(x):
  def clz_from_iterable(meta, data): class FederatedAlgorithm: class ClientPreprocessor:
class FederatedData(abc.ABC):
class RepeatableIterator:
class SubsetFederatedData(FederatedData):
class FederatedDataBuilder(abc.ABC):
  def __init__(self,
               fns: Iterable[Callable[[ClientId, client_datasets.Examples],
                                      client_datasets.Examples]] = ()):
  def __call__(self, client_id: ClientId,
               examples: client_datasets.Examples) -> client_datasets.Examples:
  def append(
      self, fn: Callable[[ClientId, client_datasets.Examples],
                         client_datasets.Examples]
  ) -> 'ClientPreprocessor':
  def __str__(self) -> str:
  def __repr__(self) -> str:
  def slice(self,
            start: Optional[ClientId] = None,
            stop: Optional[ClientId] = None) -> 'FederatedData':
  def num_clients(self) -> int:
  def client_ids(self) -> Iterator[ClientId]:
  def client_sizes(self) -> Iterator[Tuple[ClientId, int]]:
  def client_size(self, client_id: ClientId) -> int:
  def clients(self) -> Iterator[Tuple[ClientId, client_datasets.ClientDataset]]:
  def shuffled_clients(
      self,
      buffer_size: int,
      seed: Optional[int] = None
  ) -> Iterator[Tuple[ClientId, client_datasets.ClientDataset]]:
  def get_clients(
      self, client_ids: Iterable[ClientId]
  ) -> Iterator[Tuple[ClientId, client_datasets.ClientDataset]]:
  def get_client(self, client_id: ClientId) -> client_datasets.ClientDataset:
  def preprocess_client(
      self, fn: Callable[[ClientId, client_datasets.Examples],
                         client_datasets.Examples]
  ) -> 'FederatedData':
  def preprocess_batch(
      self, fn: Callable[[client_datasets.Examples], client_datasets.Examples]
  ) -> 'FederatedData':
def shuffle_repeat_batch_federated_data(
    fd: FederatedData,
    batch_size: int,
    client_buffer_size: int,
    example_buffer_size: int,
    seed: Optional[int] = None) -> Iterator[client_datasets.Examples]:
def padded_batch_federated_data(fd: FederatedData,
                                hparams: Optional[
                                    client_datasets.PaddedBatchHParams] = None,
                                **kwargs) -> Iterator[client_datasets.Examples]:
def intersect_slice_ranges(
    current_start: Optional[ClientId], current_stop: Optional[ClientId],
    new_start: Optional[ClientId], new_stop: Optional[ClientId]
) -> Tuple[Optional[ClientId], Optional[ClientId]]:
  def __init__(self, base: Iterable[Any]):
  def __iter__(self) -> Iterator[Any]:
  def __next__(self) -> Any:
  def __init__(self,
               base: FederatedData,
               client_ids: Iterable[ClientId],
               validate=True):
  def slice(self,
            start: Optional[ClientId] = None,
            stop: Optional[ClientId] = None) -> FederatedData:
  def num_clients(self) -> int:
  def client_ids(self) -> Iterator[ClientId]:
  def client_sizes(self) -> Iterator[Tuple[ClientId, int]]:
  def client_size(self, client_id: ClientId) -> int:
  def clients(self) -> Iterator[Tuple[ClientId, client_datasets.ClientDataset]]:
  def shuffled_clients(
      self,
      buffer_size: int,
      seed: Optional[int] = None
  ) -> Iterator[Tuple[ClientId, client_datasets.ClientDataset]]:
  def get_clients(
      self, client_ids: Iterable[ClientId]
  ) -> Iterator[Tuple[ClientId, client_datasets.ClientDataset]]:
  def get_client(self, client_id: ClientId) -> client_datasets.ClientDataset:
  def preprocess_client(
      self, fn: Callable[[ClientId, client_datasets.Examples],
                         client_datasets.Examples]
  ) -> FederatedData:
  def preprocess_batch(
      self, fn: Callable[[client_datasets.Examples], client_datasets.Examples]
  ) -> FederatedData:
  def __enter__(self):
  def __exit__(self, exc_type, exc_value, exc_traceback):
  def add_many(self,
               client_ids_examples: Iterable[Tuple[bytes,
                                                   client_datasets.Examples]]): def for_each_client(client_init: ClientInit,
                    client_step: Union[ClientStep,
                                       Callable[[ClientStepState, BatchExample],
                                                ClientStepState]],
                    client_final: ClientFinal = lambda _, s: s,
                    with_step_result: bool = False):
  """Creates a function which maps over clients.

  For example, for_each_client could be used to define how to run client
  updates for each client in a federated training round. Another common use case
  of for_each_client is to run evaluation per client for a given set of model
  parameters.

  The underlying backend for for_each_client is customizable.
  For example, if multiple devies are available (e.g. TPU), a :func:`jax.pmap`
  based backend can be used to parallelize across devices.
  It's also possible to manually specify which backend to use (for debugging).

  The expected usage of for_each_client is as follows::

    # Map over clients and count how many points are greater than `limit` for
    # each client. Each client also has a different `start` that is specified
    # via client input.

    def client_init(shared_input, client_input):
      client_step_state = {
          'limit': shared_input['limit'],
          'count': client_input['start']
      }
      return client_step_state

    def client_step(client_step_state, batch):
      num = jnp.sum(batch['x'] > client_step_state['limit'])
      client_step_state = {
          'limit': client_step_state['limit'],
          'count': client_step_state['count'] + num
      }
      return client_step_state

    def client_final(shared_input, client_step_state):
      del shared_input  # Unused.
      return client_step_state['count']

    # Three clients with different data and starting counts.
    # clients = [(client_id, client_batches, client_input)]
    clients = [
        (b'cid0',
        [{'x': jnp.array([1, 2, 3, 4])}, {'x': jnp.array([1, 2, 3])}],
        {'start': jnp.array(2)}),
        (b'cid1',
        [{'x': jnp.array([1, 2])}, {'x': jnp.array([1, 2, 3, 4, 5])}],
        {'start': jnp.array(0)}),
        (b'cid2',
        [{'x': jnp.array([1])}],
        {'start': jnp.array(1)}),
    ]
    shared_input = {'limit': jnp.array(2)}

    func = fedjax.for_each_client.for_each_client(
        client_init, client_step, client_final)
    print(list(func(shared_input, clients)))
    # [(b'cid0', 5), (b'cid1', 3), (b'cid2', 1)]

  Here's the same example with per step results. ::

    # We'll also keep track of the `num` per step in our step results.

    def client_step_with_result(client_step_state, batch):
      num = jnp.sum(batch['x'] > client_step_state['limit'])
      client_step_state = {
          'limit': client_step_state['limit'],
          'count': client_step_state['count'] + num
      }
      client_step_result = {'num': num}
      return client_step_state, client_step_result

    func = fedjax.for_each_client.for_each_client(
        client_init, client_step_with_result, client_final, with_step_result=True)
    print(list(func(shared_input, clients)))
    # [
    #   (b'cid0', 5, [{'num': 2}, {'num': 1}]),
    #   (b'cid1', 3, [{'num': 0}, {'num': 3}]),
    #   (b'cid2', 1, [{'num': 0}]),
    # ]

  Args:
    client_init: Function that initializes the internal intermittent client step
      state from the share input and per client input. The shared input contains
      information like the global model parameters that are shared across all
      clients. The per client input is per client information. The initialized
      internal client step state is fed as intermittent input and output from
      `client_step` and `client_final`. This client step state usually contains
      the model parameters and optimizer state for each client that are updated
      at each `client_step`. This will be run once for each client.
    client_step: Function that takes the client step state and a batch of
      examples as input and outputs a (possibly updated) client step state.
      Optionally, per step results can also be returned as the second element if
      `with_step_result` is True. Per step results are usually diagnostics like
      gradient norm. This will be run for each batch for each client.
    client_final: Function that applies the final transformation on the internal
      client step state to the desired final client output. More meaningful
      transformations can be done here, like model update clipping. Defaults to
      just returning the client step state. This will be run once for each
      client.
    with_step_result: Indicates whether client_step returns a pair where the
      first element is considered the client output and the second element is
      the client step result.

  Returns:
    A for each client function that takes shared_input and the per client inputs
    as tuple (client_id, batched_client_data, client_rng) to map over and
    returns the outputs per client as specified in `client_final` along with
    optional per client per step results.
  """
  for_each_client_backend_ = get_for_each_client_backend()
  if with_step_result:
    return for_each_client_backend_(client_init, client_step, client_final)

  def client_step_with_result(client_step_state, batch):
    return client_step(client_step_state, batch), ()

  func = for_each_client_backend_(client_init, client_step_with_result,
                                  client_final)

  def run(shared_input, clients):
    for client_id, client_output, _ in func(shared_input, clients):
      yield client_id, client_output

  return run class Model:
class ModelEvaluator:
class AverageLossEvaluator:
  def __hash__(self) -> int:
  def __eq__(self, other: Any) -> bool:
def create_model_from_haiku(
    transformed_forward_pass: hk.Transformed,
    sample_batch: BatchExample,
    train_loss: Callable[[BatchExample, BatchTrainOutput], jnp.ndarray],
    eval_metrics: Optional[Mapping[str, metrics.Metric]] = None,
    train_kwargs: Optional[Mapping[str, Any]] = None,
    eval_kwargs: Optional[Mapping[str, Any]] = None) -> Model:
  def init(rng):
  def apply_for_train(params, batch, rng=None):
  def apply_for_eval(params, batch):
def create_model_from_stax(
    stax_init: Callable[..., Params],
    stax_apply: Callable[..., jnp.ndarray],
    sample_shape: Tuple[int, ...],
    train_loss: Callable[[BatchExample, BatchTrainOutput], jnp.ndarray],
    eval_metrics: Optional[Mapping[str, metrics.Metric]] = None,
    train_kwargs: Optional[Mapping[str, Any]] = None,
    eval_kwargs: Optional[Mapping[str, Any]] = None,
    input_key: str = 'x') -> Model:
  def init(rng):
  def apply_for_train(params, batch, rng=None):
  def apply_for_eval(params, batch):
def _evaluate_model_step(model: Model, params: Params, batch: BatchExample,
                         stat: metrics.Stat) -> Dict[str, metrics.Stat]:
def evaluate_model(model: Model, params: Params,
                   batches: Iterable[BatchExample]) -> Dict[str, jnp.ndarray]:
  def __init__(self, model: Model):
    def client_init(shared_input, client_input):
    def client_step(state, batch):
    def client_final(shared_input, state):
  def evaluate_global_params(
      self, params: Params, clients: Iterable[Tuple[federated_data.ClientId,
                                                    Iterable[BatchExample]]]
  ) -> Iterator[Tuple[federated_data.ClientId, Dict[str, jnp.ndarray]]]:
  def evaluate_per_client_params(
      self, clients: Iterable[Tuple[federated_data.ClientId,
                                    Iterable[BatchExample], Params]]
  ) -> Iterator[Tuple[federated_data.ClientId, Dict[str, jnp.ndarray]]]:
def model_per_example_loss(
    model: Model) -> Callable[[Params, BatchExample, PRNGKey], jnp.ndarray]:
  def per_example_loss(params, batch_example, rng):
def _evaluate_average_loss_step(per_example_loss, params, batch, rng,
                                accum_loss, num_examples):
def _finalize_average_loss(regularizer, params, accum_loss, num_examples):
def evaluate_average_loss(
    params: Params,
    batches: Iterable[BatchExample],
    rng: PRNGKey,
    per_example_loss: Callable[[Params, BatchExample, PRNGKey], jnp.ndarray],
    regularizer: Optional[Callable[[Params],
                                   jnp.ndarray]] = None) -> jnp.ndarray:
  def __init__(self,
               per_example_loss: Callable[[Params, BatchExample, PRNGKey],
                                          jnp.ndarray],
               regularizer: Optional[Callable[[Params], jnp.ndarray]] = None):
    def client_init(shared_input, client_input):
    def client_step(state, batch):
    def client_final(shared_input, state):
  def evaluate_global_params(
      self, params: Params, clients: Iterable[Tuple[federated_data.ClientId,
                                                    Iterable[BatchExample],
                                                    PRNGKey]]
  ) -> Iterator[Tuple[federated_data.ClientId, jnp.ndarray]]:
  def evaluate_per_client_params(
      self, clients: Iterable[Tuple[federated_data.ClientId,
                                    Iterable[BatchExample], PRNGKey, Params]]
  ) -> Iterator[Tuple[federated_data.ClientId, jnp.ndarray]]:
def grad(
    per_example_loss: Callable[[Params, BatchExample, PRNGKey], jnp.ndarray],
    regularizer: Optional[Callable[[Params], jnp.ndarray]] = None
) -> Callable[[Params, BatchExample, PRNGKey], Params]:
  def scalar_loss(params, batch_example, rng):
def model_grad(
    model: Model,
    regularizer: Optional[Callable[[Params], jnp.ndarray]] = None
) -> Callable[[Params, BatchExample, PRNGKey], Params]: class Optimizer:
def create_optimizer_from_optax(opt: optax.GradientTransformation) -> Optimizer:
  def apply(grads, opt_state, params):
def ignore_grads_haiku(optimizer: Optimizer,
                       non_trainable_names: List[Tuple[str, str]]) -> Optimizer:
  def non_trainable_to_none(module_name, name, value):
  def init(params):
  def apply(grads, opt_state, params):
def adagrad(learning_rate: ScalarOrSchedule,
            initial_accumulator_value: float = 0.1,
            eps: float = 1e-6) -> Optimizer:
def adam(learning_rate: ScalarOrSchedule,
         b1: float = 0.9,
         b2: float = 0.999,
         eps: float = 1e-8,
         eps_root: float = 0.0) -> Optimizer:
def rmsprop(learning_rate: ScalarOrSchedule,
            decay: float = 0.9,
            eps: float = 1e-8,
            initial_scale: float = 0.,
            centered: bool = False,
            momentum: Optional[float] = None,
            nesterov: bool = False) -> Optimizer:
def sgd(learning_rate: ScalarOrSchedule,
        momentum: Optional[float] = None,
        nesterov: bool = False) -> Optimizer: def tree_weight(pytree: PyTree, weight: float) -> PyTree:
def tree_inverse_weight(pytree: PyTree, weight: float) -> PyTree:
def tree_zeros_like(pytree: PyTree) -> PyTree:
def tree_add(left: PyTree, right: PyTree) -> PyTree:
def tree_sum(pytrees: Iterable[PyTree]) -> PyTree:
def tree_mean(pytrees_and_weights: Iterable[Tuple[PyTree, float]]) -> PyTree:
def tree_size(pytree: PyTree) -> int:
def tree_l2_squared(pytree: PyTree) -> float:
def tree_l2_norm(pytree: PyTree) -> float:
def tree_clip_by_global_norm(pytree: PyTree, max_norm: float) -> PyTree:
from typing import Any, Callable, Mapping, Optional, Sequence, Tuple
from fedjax.core import client_datasets
from fedjax.core import dataclasses
from fedjax.core import federated_algorithm
from fedjax.core import federated_data
from fedjax.core import for_each_client
from fedjax.core import models
from fedjax.core import optimizers
from fedjax.core import tree_util
from fedjax.core.typing import BatchExample
from fedjax.core.typing import Params
from fedjax.core.typing import PRNGKey
import jax
import jax.numpy as jnp
# Copyright 2021 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# 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.
"""Mime implementation.

Mime: Mimicking Centralized Stochastic Algorithms in Federated Learning
    Sai Praneeth Karimireddy, Martin Jaggi, Satyen Kale, Mehryar Mohri,
    Sashank J. Reddi, Sebastian U. Stich, Ananda Theertha Suresh
    https://arxiv.org/abs/2008.03606
"""




Grads = Params


def create_grads_for_each_client(grad_fn):
  """Builds for_each_client for gradient computation."""

  def client_init(server_params, client_rng):
    client_step_state = {
        'params': server_params,
        'rng': client_rng,
        'num_sum': 0.,
        'grads_sum': jax.tree_util.tree_map(jnp.zeros_like, server_params)
    }
    return client_step_state

  def client_step(client_step_state, batch):
    rng, use_rng = jax.random.split(client_step_state['rng'])
    grads = grad_fn(client_step_state['params'], batch, use_rng)
    num = jnp.sum(batch[client_datasets.EXAMPLE_MASK_KEY])
    grads_sum = tree_util.tree_add(
        tree_util.tree_weight(grads, num), client_step_state['grads_sum'])
    next_client_step_state = {
        'params': client_step_state['params'],
        'rng': rng,
        'num_sum': client_step_state['num_sum'] + num,
        'grads_sum': grads_sum
    }
    return next_client_step_state

  def client_final(server_params, client_step_state):
    del server_params
    client_output = (client_step_state['grads_sum'],
                     client_step_state['num_sum'])
    return client_output

  return for_each_client.for_each_client(client_init, client_step, client_final)


def create_train_for_each_client(grad_fn, base_optimizer):
  """Builds for_each_client for client training."""

  def client_init(shared_input, client_rng):
    client_step_state = {
        'params': shared_input['params'],
        'opt_state': shared_input['opt_state'],
        'rng': client_rng,
        'init_params': shared_input['params'],
        'control_variate': shared_input['control_variate'],
    }
    return client_step_state

  def client_step(client_step_state, batch):
    rng, use_rng = jax.random.split(client_step_state['rng'])
    client_control_variate = grad_fn(client_step_state['init_params'], batch,
                                     use_rng)
    grads = grad_fn(client_step_state['params'], batch, use_rng)
    adjusted_grads = jax.tree_util.tree_multimap(
        lambda g, cc, c: g - cc + c, grads, client_control_variate,
        client_step_state['control_variate'])
    _, params = base_optimizer.apply(adjusted_grads,
                                     client_step_state['opt_state'],
                                     client_step_state['params'])
    next_client_step_state = {
        'params': params,
        'opt_state': client_step_state['opt_state'],
        'rng': rng,
        'init_params': client_step_state['init_params'],
        'control_variate': client_step_state['control_variate'],
    }
    return next_client_step_state

  def client_final(shared_input, client_step_state):
    delta_params = jax.tree_util.tree_multimap(lambda a, b: a - b,
                                               shared_input['params'],
                                               client_step_state['params'])
    return delta_params


  return for_each_client.for_each_client(client_init, client_step, client_final)