code: # Lint as: python3
  # Copyright 2019, The TensorFlow Federated Authors.
  #
  # 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.
  """Data loader for Stackoverflow."""
  from typing import List
  
  import numpy as np
  import tensorflow as tf
  import tensorflow_federated as tff
  
  EVAL_BATCH_SIZE = 100
  
  
  def create_vocab(vocab_size):
    """Creates vocab from `vocab_size` most common words in Stackoverflow."""
    vocab_dict = tff.simulation.datasets.stackoverflow.load_word_counts()
    return list(vocab_dict.keys())[:vocab_size]
  
  
  def split_input_target(chunk):
    """Generate input and target data.
  
    The task of language model is to predict the next word.
  
    Args:
      chunk: A Tensor of text data.
  
    Returns:
      A namedtuple of input and target data.
    """
    input_text = tf.map_fn(lambda x: x[:-1], chunk)
    target_text = tf.map_fn(lambda x: x[1:], chunk)
    return (input_text, target_text)
  
  
  def build_to_ids_fn(vocab, max_seq_len):
    """Constructs function mapping examples to sequences of token indices."""
    _, _, bos, eos = get_special_tokens(len(vocab))
  
    table_values = np.arange(len(vocab), dtype=np.int64)
    table = tf.lookup.StaticVocabularyTable(
        tf.lookup.KeyValueTensorInitializer(vocab, table_values),
        num_oov_buckets=1)
  
    def to_ids(example):
      sentence = tf.reshape(example['tokens'], shape=[1])
      words = tf.strings.split(sentence, sep=' ').values
      truncated_words = words[:max_seq_len]
      tokens = table.lookup(truncated_words) + 1
      tokens = tf.cond(
          tf.less(tf.size(tokens), max_seq_len),
          lambda: tf.concat([tokens, [eos]], 0), lambda: tokens)
  
      return tf.concat([[bos], tokens], 0)
  
    return to_ids
  
  
  def batch_and_split(dataset, max_seq_len, batch_size):
    return dataset.padded_batch(
        batch_size, padded_shapes=[max_seq_len + 1]).map(
            split_input_target, num_parallel_calls=tf.data.experimental.AUTOTUNE)
  
  
  def get_special_tokens(vocab_size):
    """Gets tokens dataset preprocessing code will add to Stackoverflow."""
    pad = 0
    oov = vocab_size + 1
    bos = vocab_size + 2
    eos = vocab_size + 3
    return pad, oov, bos, eos
  
  
  def create_train_dataset_preprocess_fn(vocab: List[str],
                                         client_batch_size: int,
                                         client_epochs_per_round: int,
                                         max_seq_len: int,
                                         max_training_elements_per_user: int,
                                         max_shuffle_buffer_size=10000):
    """Creates preprocessing functions for stackoverflow data.
  
    This function returns a function which takes a dataset and returns a dataset,
    generally for mapping over a set of unprocessed client datasets during
    training.
  
    Args:
      vocab: Vocabulary which defines the embedding.
      client_batch_size: Integer representing batch size to use on the clients.
      client_epochs_per_round: Number of epochs for which to repeat train client
        dataset.
      max_seq_len: Integer determining shape of padded batches. Sequences will be
        padded up to this length, and sentences longer than `max_seq_len` will be
        truncated to this length.
      max_training_elements_per_user: Integer controlling the maximum number of
        elements to take per user. If -1, takes all elements for each user.
      max_shuffle_buffer_size: Maximum shuffle buffer size.
  
    Returns:
      Two functions, the first `preprocess_train` and the second
      `preprocess_val_and_test`, as described above.
    """
    if client_batch_size <= 0:
      raise ValueError('client_batch_size must be a positive integer; you have '
                       'passed {}'.format(client_batch_size))
    elif client_epochs_per_round <= 0:
      raise ValueError('client_epochs_per_round must be a positive integer; you '
                       'have passed {}'.format(client_epochs_per_round))
    elif max_seq_len <= 0:
      raise ValueError('max_seq_len must be a positive integer; you have '
                       'passed {}'.format(max_seq_len))
    elif max_training_elements_per_user < -1:
      raise ValueError(
          'max_training_elements_per_user must be an integer at '
          'least -1; you have passed {}'.format(max_training_elements_per_user))
  
    if (max_training_elements_per_user == -1 or
        max_training_elements_per_user > max_shuffle_buffer_size):
      shuffle_buffer_size = max_shuffle_buffer_size
    else:
      shuffle_buffer_size = max_training_elements_per_user
  
    # TODO(b/155408842): need further investigation on why `tff.tf_compuation`
    # decorator causes b/153363900 for `to_ids`, and large memory consumption.
    def preprocess_train(dataset):
      to_ids = build_to_ids_fn(vocab, max_seq_len)
      dataset = dataset.take(max_training_elements_per_user)
      dataset = dataset.shuffle(shuffle_buffer_size)
      dataset = dataset.repeat(client_epochs_per_round)
      dataset = dataset.map(
          to_ids, num_parallel_calls=tf.data.experimental.AUTOTUNE)
      return batch_and_split(dataset, max_seq_len, client_batch_size)
  
    return preprocess_train
  
  
  def create_test_dataset_preprocess_fn(vocab: List[str], max_seq_len: int):
    """Creates preprocessing functions for stackoverflow data.
  
    This function returns a function which represents preprocessing logic
    for use on centralized validation and test datasets outside of TFF.
  
    Args:
      vocab: Vocabulary which defines the embedding.
      max_seq_len: Integer determining shape of padded batches. Sequences will be
        padded up to this length, and sentences longer than `max_seq_len` will be
        truncated to this length.
  
    Returns:
      `preprocess_val_and_test`, as described above.
    """
    if max_seq_len <= 0:
      raise ValueError('max_seq_len must be a positive integer; you have '
                       'passed {}'.format(max_seq_len))
  
    def preprocess_val_and_test(dataset):
      to_ids = build_to_ids_fn(vocab, max_seq_len)
      id_dataset = dataset.map(
          to_ids, num_parallel_calls=tf.data.experimental.AUTOTUNE)
      return batch_and_split(id_dataset, max_seq_len, EVAL_BATCH_SIZE)
  
    return preprocess_val_and_test
  
  
  def construct_word_level_datasets(vocab_size: int,
                                    client_batch_size: int,
                                    client_epochs_per_round: int,
                                    max_seq_len: int,
                                    max_training_elements_per_user: int,
                                    num_validation_examples: int,
                                    max_shuffle_buffer_size=10000):
    """Preprocessing for Stackoverflow data.
  
    Notice that this preprocessing function *ignores* the heldout Stackoverflow
    dataset for consistency with the other datasets in the proposed optimization
    paper, and returns a validation/test split of the Stackoverflow "test" data,
    containing more examples from users in the Stackoverflow train dataset.
  
    Args:
      vocab_size: Integer representing size of the vocab to use. Vocabulary will
        then be the `vocab_size` most frequent words in the Stackoverflow dataset.
      client_batch_size: Integer representing batch size to use on the clients.
      client_epochs_per_round: Number of epochs for which to repeat train client
        dataset.
      max_seq_len: Integer determining shape of padded batches. Sequences will be
        padded up to this length, and sentences longer than `max_seq_len` will be
        truncated to this length.
      max_training_elements_per_user: Integer controlling the maximum number of
        elements to take per user. If -1, takes all elements for each user.
      num_validation_examples: Number of examples from Stackoverflow test set to
        use for validation on each round.
      max_shuffle_buffer_size: Maximum shuffle buffer size.
  
    Returns:
      stackoverflow_train: An instance of `tff.simulation.ClientData`
        representing Stackoverflow data for training.
      stackoverflow_validation: A split of the Stackoverflow Test data as outlined
        in `tff.simulation.datasets.stackoverflow`, containing at most
        `num_validation_examples` examples.
      stackoverflow_test: A split of the same Stackoverflow Test data containing
      the examples not used in `stackoverflow_validation`.
    """
    if num_validation_examples < 1:
      raise ValueError(
          'num_validation_examples must be an integer at '
          'least 1; you have passed {}'.format(num_validation_examples))
    elif vocab_size <= 0:
      raise ValueError('vocab_size must be a positive integer; you have '
                       'passed {}'.format(vocab_size))
  
    (stackoverflow_train, _,
     stackoverflow_test) = tff.simulation.datasets.stackoverflow.load_data()
  
    vocab = create_vocab(vocab_size)
  
    raw_test_dataset = stackoverflow_test.create_tf_dataset_from_all_clients()
  
    preprocess_train = create_train_dataset_preprocess_fn(
        vocab, client_batch_size, client_epochs_per_round, max_seq_len,
        max_training_elements_per_user, max_shuffle_buffer_size)
  
    preprocess_val_and_test = create_test_dataset_preprocess_fn(
        vocab, max_seq_len)
  
    stackoverflow_train = stackoverflow_train.preprocess(preprocess_train)
    stackoverflow_val = preprocess_val_and_test(
        raw_test_dataset.take(num_validation_examples))
    stackoverflow_test = preprocess_val_and_test(
        raw_test_dataset.skip(num_validation_examples))
  
    return stackoverflow_train, stackoverflow_val, stackoverflow_test
  
  
  def get_centralized_train_dataset(vocab_size: int,
                                    batch_size: int,
                                    max_seq_len: int,
                                    shuffle_buffer_size: int = 10000):
    """Creates centralized approximately shuffled train dataset."""
  
    vocab = create_vocab(vocab_size)
    to_ids = build_to_ids_fn(vocab, max_seq_len)
    train, _, _ = tff.simulation.datasets.stackoverflow.load_data()
  
    train = train.create_tf_dataset_from_all_clients()
    train = train.shuffle(buffer_size=shuffle_buffer_size)
    return batch_and_split(
        train.map(to_ids, num_parallel_calls=tf.data.experimental.AUTOTUNE),
        max_seq_len, batch_size)
   # Lint as: python3
  # Copyright 2019, The TensorFlow Federated Authors.
  #
  # 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.
  
  from absl.testing import absltest
  import tensorflow as tf
  
  from tensorflow_federated.python.core.api import computation_types
  from tensorflow_federated.python.core.api import computations
  from tensorflow_federated.python.core.backends.mapreduce import canonical_form
  from tensorflow_federated.python.core.backends.mapreduce import test_utils
  
  tf.compat.v1.enable_v2_behavior()
  
  
  def _dummy_canonical_form_computations():
  
    @computations.tf_computation
    def initialize():
      return tf.constant(0)
  
    @computations.tf_computation(tf.int32)
    def prepare(server_state):
      del server_state  # Unused
      return tf.constant(1.0)
  
    @computations.tf_computation(
        computation_types.SequenceType(tf.float32), tf.float32)
    def work(client_data, client_input):
      del client_data  # Unused
      del client_input  # Unused
      return (True, []), []
  
    @computations.tf_computation
    def zero():
      return tf.constant(0), tf.constant(0)
  
    @computations.tf_computation((tf.int32, tf.int32), tf.bool)
    def accumulate(accumulator, client_update):
      del accumulator  # Unused
      del client_update  # Unused
      return tf.constant(1), tf.constant(1)
  
    @computations.tf_computation((tf.int32, tf.int32), (tf.int32, tf.int32))
    def merge(accumulator1, accumulator2):
      del accumulator1  # Unused
      del accumulator2  # Unused
      return tf.constant(1), tf.constant(1)
  
    @computations.tf_computation(tf.int32, tf.int32)
    def report(accumulator):
      del accumulator  # Unused
      return tf.constant(1.0)
  
    @computations.tf_computation
    def bitwidth():
      return []
  
    @computations.tf_computation(
        tf.int32, (tf.float32, computation_types.NamedTupleType([])))
    def update(server_state, global_update):
      del server_state  # Unused
      del global_update  # Unused
      return tf.constant(1), []
  
    return (initialize, prepare, work, zero, accumulate, merge, report, bitwidth,
            update)
  
  
  class CanonicalFormTest(absltest.TestCase):
  
    def test_init_does_not_raise_type_error(self):
      (initialize, prepare, work, zero, accumulate, merge, report, bitwidth,
       update) = _dummy_canonical_form_computations()
  
      try:
        canonical_form.CanonicalForm(initialize, prepare, work, zero, accumulate,
                                     merge, report, bitwidth, update)
      except TypeError:
        self.fail('Raised TypeError unexpectedly.')
  
    def test_init_does_not_raise_type_error_with_unknown_dimensions(self):
  
      @computations.tf_computation
      def initialize():
        return tf.constant(0)
  
      @computations.tf_computation(tf.int32)
      def prepare(server_state):
        del server_state  # Unused
        return tf.constant(1.0)
  
      @computations.tf_computation(
          computation_types.SequenceType(tf.float32), tf.float32)
      def work(client_data, client_input):
        del client_data  # Unused
        del client_input  # Unused
        return (True, []), []
  
      @computations.tf_computation
      def zero():
        return tf.constant([], dtype=tf.string)
  
      @computations.tf_computation(
          computation_types.TensorType(shape=[None], dtype=tf.string), tf.bool)
      def accumulate(accumulator, client_update):
        del accumulator  # Unused
        del client_update  # Unused
        return tf.constant(['abc'])
  
      @computations.tf_computation(
          computation_types.TensorType(shape=[None], dtype=tf.string),
          computation_types.TensorType(shape=[None], dtype=tf.string))
      def merge(accumulator1, accumulator2):
        del accumulator1  # Unused
        del accumulator2  # Unused
        return tf.constant(['abc'])
  
      @computations.tf_computation(
          computation_types.TensorType(shape=[None], dtype=tf.string))
      def report(accumulator):
        del accumulator  # Unused
        return tf.constant(1.0)
  
      @computations.tf_computation
      def bitwidth():
        return []
  
      @computations.tf_computation(
          tf.int32, (tf.float32, computation_types.NamedTupleType([])))
      def update(server_state, global_update):
        del server_state  # Unused
        del global_update  # Unused
        return tf.constant(1), []
  
      try:
        canonical_form.CanonicalForm(initialize, prepare, work, zero, accumulate,
                                     merge, report, bitwidth, update)
      except TypeError:
        self.fail('Raised TypeError unexpectedly.')
  
    def test_init_raises_type_error_with_bad_initialize_result_type(self):
      (_, prepare, work, zero, accumulate, merge, report, bitwidth,
       update) = _dummy_canonical_form_computations()
  
      @computations.tf_computation
      def initialize():
        return tf.constant(0.0)
  
      with self.assertRaises(TypeError):
        canonical_form.CanonicalForm(initialize, prepare, work, zero, accumulate,
                                     merge, report, bitwidth, update)
  
    def test_init_raises_type_error_with_bad_prepare_parameter_type(self):
      (initialize, _, work, zero, accumulate, merge, report, bitwidth,
       update) = _dummy_canonical_form_computations()
  
      @computations.tf_computation(tf.float32)
      def prepare(server_state):
        del server_state  # Unused
        return tf.constant(1.0)
  
      with self.assertRaises(TypeError):
        canonical_form.CanonicalForm(initialize, prepare, work, zero, accumulate,
                                     merge, report, bitwidth, update)
  
    def test_init_raises_type_error_with_bad_prepare_result_type(self):
      (initialize, _, work, zero, accumulate, merge, report, bitwidth,
       update) = _dummy_canonical_form_computations()
  
      @computations.tf_computation(tf.int32)
      def prepare(server_state):
        del server_state  # Unused
        return tf.constant(1)
  
      with self.assertRaises(TypeError):
        canonical_form.CanonicalForm(initialize, prepare, work, zero, accumulate,
                                     merge, report, bitwidth, update)
  
    def test_init_raises_type_error_with_bad_work_second_parameter_type(self):
      (initialize, prepare, _, zero, accumulate, merge, report, bitwidth,
       update) = _dummy_canonical_form_computations()
  
      @computations.tf_computation(
          computation_types.SequenceType(tf.float32), tf.int32)
      def work(client_data, client_input):
        del client_data  # Unused
        del client_input  # Unused
        return (True, []), []
  
      with self.assertRaises(TypeError):
        canonical_form.CanonicalForm(initialize, prepare, work, zero, accumulate,
                                     merge, report, bitwidth, update)
  
    def test_init_raises_type_error_with_bad_work_result_type(self):
      (initialize, prepare, _, zero, accumulate, merge, report, bitwidth,
       update) = _dummy_canonical_form_computations()
  
      @computations.tf_computation(
          computation_types.SequenceType(tf.float32), tf.float32)
      def work(client_data, client_input):
        del client_data  # Unused
        del client_input  # Unused
        return (tf.constant('abc'), []), []
  
      with self.assertRaises(TypeError):
        canonical_form.CanonicalForm(initialize, prepare, work, zero, accumulate,
                                     merge, report, bitwidth, update)
  
    def test_init_raises_type_error_with_bad_zero_result_type(self):
      (initialize, prepare, work, _, accumulate, merge, report, bitwidth,
       update) = _dummy_canonical_form_computations()
  
      @computations.tf_computation
      def zero():
        return tf.constant(0.0), tf.constant(0)
  
      with self.assertRaises(TypeError):
        canonical_form.CanonicalForm(initialize, prepare, work, zero, accumulate,
                                     merge, report, bitwidth, update)
  
    def test_init_raises_type_error_with_bad_accumulate_first_parameter_type(
        self):
      (initialize, prepare, work, zero, _, merge, report, bitwidth,
       update) = _dummy_canonical_form_computations()
  
      @computations.tf_computation((tf.float32, tf.int32), tf.bool)
      def accumulate(accumulator, client_update):
        del accumulator  # Unused
        del client_update  # Unused
        return tf.constant(1), tf.constant(1)
  
      with self.assertRaises(TypeError):
        canonical_form.CanonicalForm(initialize, prepare, work, zero, accumulate,
                                     merge, report, bitwidth, update)
  
    def test_init_raises_type_error_with_bad_accumulate_second_parameter_type(
        self):
      (initialize, prepare, work, zero, _, merge, report, bitwidth,
       update) = _dummy_canonical_form_computations()
  
      @computations.tf_computation((tf.float32, tf.float32), tf.string)
      def accumulate(accumulator, client_update):
        del accumulator  # Unused
        del client_update  # Unused
        return tf.constant(1), tf.constant(1)
  
      with self.assertRaises(TypeError):
        canonical_form.CanonicalForm(initialize, prepare, work, zero, accumulate,
                                     merge, report, bitwidth, update)
  
    def test_init_raises_type_error_with_bad_accumulate_result_type(self):
      (initialize, prepare, work, zero, _, merge, report, bitwidth,
       update) = _dummy_canonical_form_computations()
  
      @computations.tf_computation((tf.float32, tf.float32), tf.bool)
      def accumulate(accumulator, client_update):
        del accumulator  # Unused
        del client_update  # Unused
        return tf.constant(1.0), tf.constant(1)
  
      with self.assertRaises(TypeError):
        canonical_form.CanonicalForm(initialize, prepare, work, zero, accumulate,
                                     merge, report, bitwidth, update)
  
    def test_init_raises_type_error_with_bad_merge_first_parameter_type(self):
      (initialize, prepare, work, zero, accumulate, _, report, bitwidth,
       update) = _dummy_canonical_form_computations()
  
      @computations.tf_computation((tf.float32, tf.int32), (tf.int32, tf.int32))
      def merge(accumulator1, accumulator2):
        del accumulator1  # Unused
        del accumulator2  # Unused
        return tf.constant(1), tf.constant(1)
  
      with self.assertRaises(TypeError):
        canonical_form.CanonicalForm(initialize, prepare, work, zero, accumulate,
                                     merge, report, bitwidth, update)
  
    def test_init_raises_type_error_with_bad_merge_second_parameter_type(self):
      (initialize, prepare, work, zero, accumulate, _, report, bitwidth,
       update) = _dummy_canonical_form_computations()
  
      @computations.tf_computation((tf.int32, tf.int32), (tf.float32, tf.int32))
      def merge(accumulator1, accumulator2):
        del accumulator1  # Unused
        del accumulator2  # Unused
        return tf.constant(1), tf.constant(1)
  
      with self.assertRaises(TypeError):
        canonical_form.CanonicalForm(initialize, prepare, work, zero, accumulate,
                                     merge, report, bitwidth, update)
  
    def test_init_raises_type_error_with_bad_merge_result_type(self):
      (initialize, prepare, work, zero, accumulate, _, report, bitwidth,
       update) = _dummy_canonical_form_computations()
  
      @computations.tf_computation((tf.int32, tf.int32), (tf.int32, tf.int32))
      def merge(accumulator1, accumulator2):
        del accumulator1  # Unused
        del accumulator2  # Unused
        return tf.constant(1.0), tf.constant(1)
  
      with self.assertRaises(TypeError):
        canonical_form.CanonicalForm(initialize, prepare, work, zero, accumulate,
                                     merge, report, bitwidth, update)
  
    def test_init_raises_type_error_with_bad_report_parameter_type(self):
      (initialize, prepare, work, zero, accumulate, merge, _, bitwidth,
       update) = _dummy_canonical_form_computations()
  
      @computations.tf_computation(tf.float32, tf.int32)
      def report(accumulator):
        del accumulator  # Unused
        return tf.constant(1.0)
  
      with self.assertRaises(TypeError):
        canonical_form.CanonicalForm(initialize, prepare, work, zero, accumulate,
                                     merge, report, bitwidth, update)
  
    def test_init_raises_type_error_with_bad_report_result_type(self):
      (initialize, prepare, work, zero, accumulate, merge, _, bitwidth,
       update) = _dummy_canonical_form_computations()
  
      @computations.tf_computation(tf.int32, tf.int32)
      def report(accumulator):
        del accumulator  # Unused
        return tf.constant(1)
  
      with self.assertRaises(TypeError):
        canonical_form.CanonicalForm(initialize, prepare, work, zero, accumulate,
                                     merge, report, bitwidth, update)
  
    def test_init_raises_type_error_with_bad_update_first_parameter_type(self):
      (initialize, prepare, work, zero, accumulate, merge, report, bitwidth,
       _) = _dummy_canonical_form_computations()
  
      @computations.tf_computation(
          tf.float32, (tf.float32, computation_types.NamedTupleType([])))
      def update(server_state, global_update):
        del server_state  # Unused
        del global_update  # Unused
        return tf.constant(1), []
  
      with self.assertRaises(TypeError):
        canonical_form.CanonicalForm(initialize, prepare, work, zero, accumulate,
                                     merge, report, bitwidth, update)
  
    def test_init_raises_type_error_with_bad_update_second_parameter_type(self):
      (initialize, prepare, work, zero, accumulate, merge, report, bitwidth,
       _) = _dummy_canonical_form_computations()
  
      @computations.tf_computation(
          tf.int32, (tf.int32, computation_types.NamedTupleType([])))
      def update(server_state, global_update):
        del server_state  # Unused
        del global_update  # Unused
        return tf.constant(1), []
  
      with self.assertRaises(TypeError):
        canonical_form.CanonicalForm(initialize, prepare, work, zero, accumulate,
                                     merge, report, bitwidth, update)
  
    def test_init_raises_type_error_with_bad_update_result_type(self):
      (initialize, prepare, work, zero, accumulate, merge, report, bitwidth,
       _) = _dummy_canonical_form_computations()
  
      @computations.tf_computation(
          tf.int32, (tf.float32, computation_types.NamedTupleType([])))
      def update(server_state, global_update):
        del server_state  # Unused
        del global_update  # Unused
        return tf.constant(1.0), []
  
      with self.assertRaises(TypeError):
        canonical_form.CanonicalForm(initialize, prepare, work, zero, accumulate,
                                     merge, report, bitwidth, update)
  
    def test_summary(self):
      cf = test_utils.get_temperature_sensor_example()
  
      class CapturePrint(object):
  
        def __init__(self):
          self.summary = ''
  
        def __call__(self, msg):
          self.summary += msg + '\n'
  
      capture = CapturePrint()
      cf.summary(print_fn=capture)
      # pyformat: disable
      self.assertEqual(
          capture.summary,
          'initialize: ( -> <num_rounds=int32>)\n'
          'prepare   : (<num_rounds=int32> -> <max_temperature=float32>)\n'
          'work      : (<float32*,<max_temperature=float32>> -> <<<is_over=bool>,<>>,<num_readings=int32>>)\n'
          'zero      : ( -> <num_total=int32,num_over=int32>)\n'
          'accumulate: (<<num_total=int32,num_over=int32>,<is_over=bool>> -> <num_total=int32,num_over=int32>)\n'
          'merge     : (<<num_total=int32,num_over=int32>,<num_total=int32,num_over=int32>> -> <num_total=int32,num_over=int32>)\n'
          'report    : (<num_total=int32,num_over=int32> -> <ratio_over_threshold=float32>)\n'
          'bitwidth  : ( -> <>)\n'
          'update    : ( -> <num_rounds=int32>)\n'
      )
      # pyformat: enable
  
  
  if __name__ == '__main__':
    absltest.main()
   # Lint as: python3
  # Copyright 2020, The TensorFlow Federated Authors.
  #
  # 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.
  
  import string
  
  import tensorflow as tf
  import tensorflow_federated as tff
  
  from tensorflow_federated.python.research.triehh import triehh_tf
  from tensorflow_federated.python.research.triehh import triehh_tff
  
  
  class TriehhTffTest(tf.test.TestCase):
  
    def test_build_triehh_process_works_as_expeted(self):
      clients = 3
      num_sub_rounds = 4
      max_rounds = 6
      max_num_heavy_hitters = 3
      max_user_contribution = 100
      roots = (
          string.ascii_lowercase + string.digits + "'@#-;*:./" +
          triehh_tf.DEFAULT_TERMINATOR)
      possible_prefix_extensions = list(roots)
  
      iterative_process = triehh_tff.build_triehh_process(
          possible_prefix_extensions,
          num_sub_rounds,
          max_num_heavy_hitters,
          max_user_contribution,
          default_terminator=triehh_tf.DEFAULT_TERMINATOR)
  
      server_state = iterative_process.initialize()
      expected_discovered_prefixes = tf.constant([''], dtype=tf.string)
      expected_discovered_heavy_hitters = tf.constant([], dtype=tf.string)
      expected_accumulated_votes = tf.zeros(
          dtype=tf.int32,
          shape=[max_num_heavy_hitters,
                 len(possible_prefix_extensions)])
      expected_round_num = tf.constant(0, dtype=tf.int32)
  
      self.assertAllEqual(server_state.discovered_prefixes,
                          expected_discovered_prefixes)
      self.assertAllEqual(server_state.discovered_heavy_hitters,
                          expected_discovered_heavy_hitters)
      self.assertAllEqual(server_state.accumulated_votes,
                          expected_accumulated_votes)
      self.assertAllEqual(server_state.round_num, expected_round_num)
  
      def create_dataset_fn(client_id):
        del client_id
        return tf.data.Dataset.from_tensor_slices(['hello', 'hey', 'hi'])
  
      client_ids = list(range(100))
  
      client_data = tff.simulation.ClientData.from_clients_and_fn(
          client_ids=client_ids,
          create_tf_dataset_for_client_fn=create_dataset_fn)
  
      for round_num in range(max_rounds * num_sub_rounds):
        # TODO(b/152051528): Remove this once lookup table state is cleared in
        # eager executer.
        tff.framework.set_default_executor(tff.framework.local_executor_factory())
        sampled_clients = list(range(clients))
        sampled_datasets = [
            client_data.create_tf_dataset_for_client(client_id)
            for client_id in sampled_clients
        ]
        server_state, _ = iterative_process.next(server_state, sampled_datasets)
  
        if (round_num + 1) % num_sub_rounds == 0:
          if (max_num_heavy_hitters - len(server_state.discovered_heavy_hitters) <
              1) or (server_state.discovered_prefixes.size == 0):
            # Training is done.
            # All max_num_heavy_hitters have been discovered.
            break
  
      expected_discovered_heavy_hitters = tf.constant(['hi', 'hey', 'hello'],
                                                      dtype=tf.string)
  
      self.assertAllEqual(server_state.discovered_heavy_hitters,
                          expected_discovered_heavy_hitters)
  
  
  if __name__ == '__main__':
    tf.compat.v1.enable_v2_behavior()
    tf.test.main()
  
apis: tensorflow.size tensorflow.lookup.KeyValueTensorInitializer tensorflow.reshape tensorflow.map_fn tensorflow.strings.split tensorflow.concat
apis: tensorflow.constant tensorflow.compat.v1.enable_v2_behavior
apis: tensorflow.constant tensorflow.data.Dataset.from_tensor_slices tensorflow.test.main tensorflow.compat.v1.enable_v2_behavior