Automated rollback of commit 078120c0

PiperOrigin-RevId: 491782781

tensorflow_federated/python/aggregators/BUILD

@@ -275,7 +275,8 @@ py_library(
     srcs = ["elias_gamma_encode.py"],
     srcs_version = "PY3",
     deps = [
-        ":factory",
+        "//tensorflow_federated/python/aggregators:concat",
+        "//tensorflow_federated/python/aggregators:factory",
         "//tensorflow_federated/python/common_libs:py_typecheck",
         "//tensorflow_federated/python/common_libs:structure",
         "//tensorflow_federated/python/core/impl/federated_context:federated_computation",
@@ -283,7 +284,6 @@ py_library(
         "//tensorflow_federated/python/core/impl/tensorflow_context:tensorflow_computation",
         "//tensorflow_federated/python/core/impl/types:computation_types",
         "//tensorflow_federated/python/core/impl/types:placements",
-        "//tensorflow_federated/python/core/impl/types:type_conversions",
         "//tensorflow_federated/python/core/templates:aggregation_process",
         "//tensorflow_federated/python/core/templates:measured_process",
     ],
@@ -297,7 +297,7 @@ py_test(
     srcs_version = "PY3",
     deps = [
         ":elias_gamma_encode",
-        ":mean",
+        "//tensorflow_federated/python/aggregators:mean",
         "//tensorflow_federated/python/core/backends/native:execution_contexts",
         "//tensorflow_federated/python/core/impl/types:computation_types",
         "//tensorflow_federated/python/core/impl/types:type_test_utils",

tensorflow_federated/python/aggregators/elias_gamma_encode.py

@@ -19,6 +19,7 @@ from typing import Optional
 import tensorflow as tf
 import tensorflow_compression as tfc
 
+from tensorflow_federated.python.aggregators import concat
 from tensorflow_federated.python.aggregators import factory
 from tensorflow_federated.python.common_libs import py_typecheck
 from tensorflow_federated.python.common_libs import structure
@@ -27,15 +28,16 @@ from tensorflow_federated.python.core.impl.federated_context import intrinsics
 from tensorflow_federated.python.core.impl.tensorflow_context import tensorflow_computation
 from tensorflow_federated.python.core.impl.types import computation_types
 from tensorflow_federated.python.core.impl.types import placements
-from tensorflow_federated.python.core.impl.types import type_conversions
 from tensorflow_federated.python.core.templates import aggregation_process
 from tensorflow_federated.python.core.templates import measured_process
 
 
 @tensorflow_computation.tf_computation
-def _get_bits(value):
-  """Return size (in bits) of an encoded tensor."""
-  return 8. * tf.cast(tf.strings.length(value, unit="BYTE"), dtype=tf.float64)
+def _get_bitrate(value, num_elements):
+  """Return size (in bits) of an encoded value."""
+  bitstring_length = 8. * tf.cast(
+      tf.strings.length(value, unit="BYTE"), dtype=tf.float64)
+  return tf.math.divide_no_nan(bitstring_length, num_elements)
 
 
 def _is_int32_or_structure_of_int32s(type_spec: computation_types.Type) -> bool:
@@ -114,36 +116,28 @@ class EliasGammaEncodedSumFactory(factory.UnweightedAggregationFactory):
       raise ValueError("Expect value_type to be an int32 tensor or a structure "
                        "containing only other structures of int32 tensors, "
                        f"found {value_type}.")
+    concat_fn, unconcat_fn = concat.create_concat_fns(value_type)
+    concat_value_type = concat_fn.type_signature.result
 
     if self._bitrate_mean_factory is not None:
       bitrate_mean_process = self._bitrate_mean_factory.create(
           computation_types.to_type(tf.float64))
 
-    @tensorflow_computation.tf_computation(value_type)
-    def encode(value):
-      return tf.nest.map_structure(
-          lambda x: tfc.run_length_gamma_encode(data=x), value)
-
     def sum_encoded_value(value):
 
       @tensorflow_computation.tf_computation
       def get_accumulator():
-        return type_conversions.structure_from_tensor_type_tree(
-            lambda x: tf.zeros(shape=x.shape, dtype=tf.int32), value_type)
+        return tf.zeros(shape=concat_value_type.shape, dtype=tf.int32)
 
       @tensorflow_computation.tf_computation
       def decode_accumulate_values(accumulator, encoded_value):
-        shapes = type_conversions.structure_from_tensor_type_tree(
-            lambda x: x.shape, value_type)
-        return tf.nest.map_structure(
-            lambda a, x, y: a + tfc.run_length_gamma_decode(code=x, shape=y),
-            accumulator, encoded_value, shapes)
+        decoded_value = tfc.run_length_gamma_decode(
+            code=encoded_value, shape=concat_value_type.shape)
+        return accumulator + decoded_value
 
       @tensorflow_computation.tf_computation
       def merge_decoded_values(decoded_value_1, decoded_value_2):
-        return tf.nest.map_structure(
-            tensorflow_computation.tf_computation(lambda x, y: x + y),
-            decoded_value_1, decoded_value_2)
+        return decoded_value_1 + decoded_value_2
 
       @tensorflow_computation.tf_computation
       def report_decoded_summation(summed_decoded_values):
@@ -164,35 +158,27 @@ class EliasGammaEncodedSumFactory(factory.UnweightedAggregationFactory):
         init_fn.type_signature.result, computation_types.at_clients(value_type))
     def next_fn(state, value):
       measurements = ()
-      encoded_value = intrinsics.federated_map(encode, value)
+      concat_value = intrinsics.federated_map(concat_fn, value)
+      encoded_value = intrinsics.federated_map(
+          tensorflow_computation.tf_computation(
+              lambda x: tfc.run_length_gamma_encode(data=x)), concat_value)
       if self._bitrate_mean_factory is not None:
 
         @tensorflow_computation.tf_computation
         def get_num_elements():
-          num_elements = type_conversions.structure_from_tensor_type_tree(
-              lambda x: tf.constant(x.shape.num_elements(), tf.float64),
-              value_type)
-          return tf.math.add_n(tf.nest.flatten(num_elements))
+          return tf.constant(concat_value_type.shape.num_elements(), tf.float64)
 
         num_elements = intrinsics.federated_eval(get_num_elements,
                                                  placements.CLIENTS)
-
-        @tensorflow_computation.tf_computation
-        def struct_get_bits(x):
-          return tf.math.add_n([_get_bits(t) for t in tf.nest.flatten(x)])
-
-        total_bits = intrinsics.federated_map(struct_get_bits, encoded_value)
-
-        bitrates = intrinsics.federated_map(
-            tensorflow_computation.tf_computation(
-                lambda x, y: tf.math.divide_no_nan(x=x, y=y, name="divide")),
-            (total_bits, num_elements))
+        bitrates = intrinsics.federated_map(_get_bitrate,
+                                            (encoded_value, num_elements))
         avg_bitrate = bitrate_mean_process.next(
             bitrate_mean_process.initialize(), bitrates).result
         measurements = intrinsics.federated_zip(
             collections.OrderedDict(elias_gamma_code_avg_bitrate=avg_bitrate))
       decoded_value = sum_encoded_value(encoded_value)
+      unconcat_value = intrinsics.federated_map(unconcat_fn, decoded_value)
       return measured_process.MeasuredProcessOutput(
-          state=state, result=decoded_value, measurements=measurements)
+          state=state, result=unconcat_value, measurements=measurements)
 
     return aggregation_process.AggregationProcess(init_fn, next_fn)

tensorflow_federated/python/aggregators/elias_gamma_encode_test.py

@@ -63,15 +63,9 @@ _test_expected_result_struct_int32_tensors = collections.OrderedDict({
     'layer2': [-6, 2, 0, 0],
 })
 
-
-# Avg num bits accross clients to represent each element
-_test_avg_bitrate_int32_tensor_rank_1 = 16.0 / 4.0
-_test_avg_bitrate_int32_tensor_rank_2 = 32.0 / 8.0
-_test_avg_bitrate_struct_int32_tensors = 32.0 / 8.0
-
-# Avg num bits across clients to represent each tensor
-_test_avg_bitstring_length_int32_tensor_rank_1 = 16.0
-_test_avg_bitstring_length_int32_tensor_rank_2 = 32.0
+_test_avg_bitrate_int32_tensor_rank_1 = 16. / 4.
+_test_avg_bitrate_int32_tensor_rank_2 = 32. / 8.
+_test_avg_bitrate_struct_int32_tensors = 32. / 8.
 
 
 class EncodeUtilTest(tf.test.TestCase, parameterized.TestCase):
@@ -174,13 +168,16 @@ class EncodeExecutionTest(tf.test.TestCase, parameterized.TestCase):
 
   @parameterized.named_parameters(
       ('int32_tensor_rank_1', _test_client_values_int32_tensor_rank_1,
-       _test_avg_bitstring_length_int32_tensor_rank_1),
+       _test_avg_bitrate_int32_tensor_rank_1),
       ('int32_tensor_rank_2', _test_client_values_int32_tensor_rank_2,
-       _test_avg_bitstring_length_int32_tensor_rank_2))
-  def test_bitstring_impl(self, client_values, avg_total_bits):
+       _test_avg_bitrate_int32_tensor_rank_2))
+  def test_bitstring_impl(self, client_values, avg_bitrate):
+    num_elements = tf.size(client_values[0], out_type=tf.float64)
     bitstrings = [tfc.run_length_gamma_encode(x) for x in client_values]
-    total_bits = [elias_gamma_encode._get_bits(x) for x in bitstrings]
-    self.assertEqual(np.mean(total_bits), avg_total_bits)
+    bitrates = [
+        elias_gamma_encode._get_bitrate(x, num_elements) for x in bitstrings
+    ]
+    self.assertEqual(np.mean(bitrates), avg_bitrate)
 
 
 if __name__ == '__main__':

tensorflow_federated/python/aggregators/stochastic_discretization.py

@@ -187,9 +187,9 @@ def _discretize_struct(struct, step_size):
     seed = tf.cast(
         tf.stack([tf.timestamp() * 1e6,
                   tf.timestamp() * 1e6]), dtype=tf.int64)
-    scaled_x = tf.divide(x, tf.cast(step_size, x.dtype))
-    prob_x = scaled_x - tf.floor(scaled_x)
-    random_x = tf.random.stateless_uniform(x.shape, seed=seed, dtype=x.dtype)
+    scaled_x = tf.divide(tf.cast(x, tf.float32), step_size)
+    prob_x = scaled_x - tf.cast(tf.floor(scaled_x), tf.float32)
+    random_x = tf.random.stateless_uniform(x.shape, seed=seed, dtype=tf.float32)
     discretized_x = tf.where(
         tf.less_equal(random_x, prob_x), tf.math.ceil(scaled_x),
         tf.math.floor(scaled_x))
@@ -202,7 +202,7 @@ def _undiscretize_struct(struct, step_size, tf_dtype_struct):
   """Unscales the discretized structure and casts back to original dtypes."""
 
   def undiscretize_tensor(x, original_dtype):
-    unscaled_x = tf.cast(x, original_dtype) * tf.cast(step_size, original_dtype)
-    return unscaled_x
+    unscaled_x = tf.cast(x, tf.float32) * step_size
+    return tf.cast(unscaled_x, original_dtype)
 
   return tf.nest.map_structure(undiscretize_tensor, struct, tf_dtype_struct)

tensorflow_federated/python/aggregators/stochastic_discretization_test.py

@@ -35,16 +35,16 @@ _test_struct_type_int = [tf.int32, (tf.int32, (2,)), (tf.int32, (3, 3))]
 _test_struct_type_float = [tf.float32, (tf.float32, (2,)), (tf.float32, (3, 3))]
 
 _test_nested_struct_type_float = collections.OrderedDict(
-    a=[tf.float16, [(tf.float32, (2, 2, 1))]], b=(tf.float16, (3, 3)))
+    a=[tf.float32, [(tf.float32, (2, 2, 1))]], b=(tf.float32, (3, 3)))
 
 
 def _make_test_nested_struct_value(value):
   return collections.OrderedDict(
       a=[
-          tf.constant(value, dtype=tf.float16),
+          tf.constant(value, dtype=tf.float32),
           [tf.constant(value, dtype=tf.float32, shape=[2, 2, 1])]
       ],
-      b=tf.constant(value, dtype=tf.float16, shape=(3, 3)))
+      b=tf.constant(value, dtype=tf.float32, shape=(3, 3)))
 
 
 def _named_test_cases_product(*args):
@@ -141,14 +141,14 @@ class StochasticDiscretizationExecutionTest(tf.test.TestCase,
   def test_discretize_impl(self, value_type, client_values, expected_sum):
     factory = stochastic_discretization.StochasticDiscretizationFactory(
         inner_agg_factory=_measurement_aggregator,
-        step_size=0.125,
+        step_size=0.1,
         distortion_aggregation_factory=mean.UnweightedMeanFactory())
     value_type = computation_types.to_type(value_type)
     process = factory.create(value_type)
     state = process.initialize()
 
     expected_result = expected_sum
-    expected_quantized_result = tf.nest.map_structure(lambda x: x * 8,
+    expected_quantized_result = tf.nest.map_structure(lambda x: x * 10,
                                                       expected_sum)
     expected_measurements = collections.OrderedDict(
         stochastic_discretization=expected_quantized_result, distortion=0.)
@@ -160,21 +160,19 @@ class StochasticDiscretizationExecutionTest(tf.test.TestCase,
       result = output.result
       self.assertAllClose(result, expected_result)
 
-  @parameterized.named_parameters(
-      ('float16', tf.float16), ('float32', tf.float32), ('float64', tf.float64))
+  @parameterized.named_parameters(('int32', tf.int32), ('int64', tf.int64),
+                                  ('float64', tf.float64))
   def test_output_dtype(self, dtype):
     """Checks the tensor type gets casted during preprocessing."""
-    x = tf.range(8)
-    x = tf.cast(x, dtype=dtype)
+    x = tf.range(8, dtype=dtype)
     encoded_x = stochastic_discretization._discretize_struct(x, step_size=0.1)
     self.assertEqual(encoded_x.dtype, stochastic_discretization.OUTPUT_TF_TYPE)
 
-  @parameterized.named_parameters(
-      ('float16', tf.float16), ('float32', tf.float32), ('float64', tf.float64))
+  @parameterized.named_parameters(('int32', tf.int32), ('int64', tf.int64),
+                                  ('float64', tf.float64))
   def test_revert_to_input_dtype(self, dtype):
     """Checks that postprocessing restores the original dtype."""
-    x = tf.range(8)
-    x = tf.cast(x, dtype=dtype)
+    x = tf.range(8, dtype=dtype)
     encoded_x = stochastic_discretization._discretize_struct(x, step_size=1)
     decoded_x = stochastic_discretization._undiscretize_struct(
         encoded_x, step_size=1, tf_dtype_struct=dtype)
@@ -234,15 +232,15 @@ class ScalingTest(tf.test.TestCase, parameterized.TestCase):
                                   ('step_size_3', 10**-5))
   def test_scaling(self, step_size):
     # Integers to prevent rounding.
-    x = tf.cast(
-        tf.random.stateless_uniform([100], (1, 1), -100, 100, dtype=tf.int32),
-        tf.float32)
-    discretized_x = stochastic_discretization._discretize_struct(x, step_size)
+    x = tf.random.stateless_uniform([100], (1, 1), -100, 100, dtype=tf.int32)
+    discretized_x = stochastic_discretization._discretize_struct(
+        x, tf.cast(step_size, tf.float32))
     reverted_x = stochastic_discretization._undiscretize_struct(
-        discretized_x, step_size, tf_dtype_struct=tf.float32)
+        discretized_x, step_size, tf_dtype_struct=tf.int32)
     x, discretized_x, reverted_x = self.evaluate([x, discretized_x, reverted_x])
-    self.assertAllEqual(tf.round(tf.divide(x, step_size)),
-                        discretized_x)  # Scaling up.
+    self.assertAllEqual(
+        tf.round(tf.divide(tf.cast(x, tf.float32), step_size)),
+        discretized_x)  # Scaling up.
     self.assertAllEqual(x, reverted_x)  # Scaling down.