Remove unused aggregators from tff.learning.

PiperOrigin-RevId: 687336045

RELEASE.md

@@ -23,6 +23,10 @@ and this project adheres to
 ### Removed
 
 *   `tff.types.tensorflow_to_type`, this function is no longer used.
+*   `tff.learning.dp_aggregator` removed. Prefer using the class methods on
+    `tff.aggregators.DifferentiallyPrivateFactory`.
+*   `tff.learning.ddp_secure_aggregator` and `tff.learning.secure_aggregator`
+    removed.
 
 ## Release 0.88.0
 

tensorflow_federated/python/learning/BUILD

@@ -61,9 +61,6 @@ py_library(
     name = "model_update_aggregator",
     srcs = ["model_update_aggregator.py"],
     deps = [
-        "//tensorflow_federated/python/aggregators:differential_privacy",
-        "//tensorflow_federated/python/aggregators:distributed_dp",
-        "//tensorflow_federated/python/aggregators:encoded",
         "//tensorflow_federated/python/aggregators:factory",
         "//tensorflow_federated/python/aggregators:mean",
         "//tensorflow_federated/python/aggregators:quantile_estimation",
@@ -88,7 +85,6 @@ py_test(
         "//tensorflow_federated/python/core/impl/types:type_analysis",
         "//tensorflow_federated/python/core/templates:aggregation_process",
         "//tensorflow_federated/python/core/templates:iterative_process",
-        "//tensorflow_federated/python/core/test:static_assert",
     ],
 )
 

tensorflow_federated/python/learning/__init__.py

@@ -46,8 +46,4 @@
 from tensorflow_federated.python.learning.debug_measurements import add_debug_measurements
 from tensorflow_federated.python.learning.debug_measurements import add_debug_measurements_with_mixed_dtype
 from tensorflow_federated.python.learning.loop_builder import LoopImplementation
-from tensorflow_federated.python.learning.model_update_aggregator import compression_aggregator
-from tensorflow_federated.python.learning.model_update_aggregator import ddp_secure_aggregator
-from tensorflow_federated.python.learning.model_update_aggregator import dp_aggregator
 from tensorflow_federated.python.learning.model_update_aggregator import robust_aggregator
-from tensorflow_federated.python.learning.model_update_aggregator import secure_aggregator

tensorflow_federated/python/learning/algorithms/BUILD

@@ -69,12 +69,9 @@ py_cpu_gpu_test(
     deps = [
         ":fed_avg",
         "//tensorflow_federated/python/aggregators:factory_utils",
-        "//tensorflow_federated/python/core/test:static_assert",
         "//tensorflow_federated/python/learning:loop_builder",
         "//tensorflow_federated/python/learning:model_update_aggregator",
-        "//tensorflow_federated/python/learning/metrics:aggregator",
         "//tensorflow_federated/python/learning/models:model_examples",
-        "//tensorflow_federated/python/learning/models:test_models",
         "//tensorflow_federated/python/learning/optimizers:sgdm",
     ],
 )
@@ -118,7 +115,6 @@ py_cpu_gpu_test(
     shard_count = 10,
     deps = [
         ":fed_avg_with_optimizer_schedule",
-        "//tensorflow_federated/python/core/test:static_assert",
         "//tensorflow_federated/python/learning:loop_builder",
         "//tensorflow_federated/python/learning:model_update_aggregator",
         "//tensorflow_federated/python/learning/metrics:aggregator",
@@ -164,10 +160,8 @@ py_cpu_gpu_test(
         ":fed_prox",
         "//tensorflow_federated/python/aggregators:factory_utils",
         "//tensorflow_federated/python/core/templates:iterative_process",
-        "//tensorflow_federated/python/core/test:static_assert",
         "//tensorflow_federated/python/learning:loop_builder",
         "//tensorflow_federated/python/learning:model_update_aggregator",
-        "//tensorflow_federated/python/learning/metrics:aggregator",
         "//tensorflow_federated/python/learning/models:model_examples",
         "//tensorflow_federated/python/learning/models:model_weights",
         "//tensorflow_federated/python/learning/models:test_models",
@@ -330,7 +324,6 @@ py_cpu_gpu_test(
     deps = [
         ":fed_sgd",
         "//tensorflow_federated/python/core/environments/tensorflow_backend:tensorflow_test_utils",
-        "//tensorflow_federated/python/core/test:static_assert",
         "//tensorflow_federated/python/learning:loop_builder",
         "//tensorflow_federated/python/learning:model_update_aggregator",
         "//tensorflow_federated/python/learning/metrics:aggregator",
@@ -429,11 +422,9 @@ py_cpu_gpu_test(
         "//tensorflow_federated/python/core/impl/types:placements",
         "//tensorflow_federated/python/core/templates:iterative_process",
         "//tensorflow_federated/python/core/templates:measured_process",
-        "//tensorflow_federated/python/core/test:static_assert",
         "//tensorflow_federated/python/learning:client_weight_lib",
         "//tensorflow_federated/python/learning:loop_builder",
         "//tensorflow_federated/python/learning:model_update_aggregator",
-        "//tensorflow_federated/python/learning/metrics:aggregator",
         "//tensorflow_federated/python/learning/metrics:counters",
         "//tensorflow_federated/python/learning/models:functional",
         "//tensorflow_federated/python/learning/models:keras_utils",

tensorflow_federated/python/learning/algorithms/fed_avg_test.py

@@ -18,38 +18,25 @@
 from absl.testing import parameterized
 
 from tensorflow_federated.python.aggregators import factory_utils
-from tensorflow_federated.python.core.test import static_assert
 from tensorflow_federated.python.learning import loop_builder
 from tensorflow_federated.python.learning import model_update_aggregator
 from tensorflow_federated.python.learning.algorithms import fed_avg
-from tensorflow_federated.python.learning.metrics import aggregator
 from tensorflow_federated.python.learning.models import model_examples
-from tensorflow_federated.python.learning.models import test_models
 from tensorflow_federated.python.learning.optimizers import sgdm
 
 
 class FedAvgTest(parameterized.TestCase):
   """Tests construction of the FedAvg training process."""
 
-  @parameterized.product(
-      optimizer_fn=[
-          sgdm.build_sgdm(learning_rate=0.1),
-      ],
-      aggregation_factory=[
-          model_update_aggregator.robust_aggregator,
-          model_update_aggregator.compression_aggregator,
-          model_update_aggregator.secure_aggregator,
-      ],
-  )
-  def test_construction_calls_model_fn(self, optimizer_fn, aggregation_factory):
+  def test_construction_calls_model_fn(self):
     # Assert that the process building does not call `model_fn` too many times.
     # `model_fn` can potentially be expensive (loading weights, processing, etc
     # ).
     mock_model_fn = mock.Mock(side_effect=model_examples.LinearRegression)
     fed_avg.build_weighted_fed_avg(
         model_fn=mock_model_fn,
-        client_optimizer_fn=optimizer_fn,
-        model_aggregator=aggregation_factory(),
+        client_optimizer_fn=sgdm.build_sgdm(learning_rate=0.1),
+        model_aggregator=model_update_aggregator.robust_aggregator(),
     )
     self.assertEqual(mock_model_fn.call_count, 3)
 
@@ -125,34 +112,6 @@ def test_unweighted_fed_avg_raises_on_weighted_aggregator(self):
           model_aggregator=model_aggregator,
       )
 
-  def test_weighted_fed_avg_with_only_secure_aggregation(self):
-    model_fn = model_examples.LinearRegression
-    learning_process = fed_avg.build_weighted_fed_avg(
-        model_fn,
-        client_optimizer_fn=sgdm.build_sgdm(),
-        model_aggregator=model_update_aggregator.secure_aggregator(
-            weighted=True
-        ),
-        metrics_aggregator=aggregator.secure_sum_then_finalize,
-    )
-    static_assert.assert_not_contains_unsecure_aggregation(
-        learning_process.next
-    )
-
-  def test_unweighted_fed_avg_with_only_secure_aggregation(self):
-    model_fn = model_examples.LinearRegression
-    learning_process = fed_avg.build_unweighted_fed_avg(
-        model_fn,
-        client_optimizer_fn=sgdm.build_sgdm(),
-        model_aggregator=model_update_aggregator.secure_aggregator(
-            weighted=False
-        ),
-        metrics_aggregator=aggregator.secure_sum_then_finalize,
-    )
-    static_assert.assert_not_contains_unsecure_aggregation(
-        learning_process.next
-    )
-
 
 class FunctionalFedAvgTest(parameterized.TestCase):
   """Tests construction of the FedAvg training process."""
@@ -167,25 +126,6 @@ def test_raises_on_non_callable_or_functional_model(self, constructor):
           model_fn=0, client_optimizer_fn=sgdm.build_sgdm(learning_rate=0.1)
       )
 
-  @parameterized.named_parameters(
-      ('weighted', fed_avg.build_weighted_fed_avg),
-      ('unweighted', fed_avg.build_unweighted_fed_avg),
-  )
-  def test_weighted_fed_avg_with_only_secure_aggregation(self, constructor):
-    model = test_models.build_functional_linear_regression()
-    learning_process = constructor(
-        model_fn=model,
-        client_optimizer_fn=sgdm.build_sgdm(learning_rate=0.1),
-        server_optimizer_fn=sgdm.build_sgdm(),
-        model_aggregator=model_update_aggregator.secure_aggregator(
-            weighted=constructor is fed_avg.build_weighted_fed_avg
-        ),
-        metrics_aggregator=aggregator.secure_sum_then_finalize,
-    )
-    static_assert.assert_not_contains_unsecure_aggregation(
-        learning_process.next
-    )
-
 
 if __name__ == '__main__':
   absltest.main()

tensorflow_federated/python/learning/algorithms/fed_avg_with_optimizer_schedule_test.py

@@ -18,7 +18,6 @@
 from absl.testing import parameterized
 import tensorflow as tf
 
-from tensorflow_federated.python.core.test import static_assert
 from tensorflow_federated.python.learning import loop_builder
 from tensorflow_federated.python.learning import model_update_aggregator
 from tensorflow_federated.python.learning.algorithms import fed_avg_with_optimizer_schedule
@@ -30,17 +29,7 @@
 
 class ClientScheduledFedAvgTest(parameterized.TestCase):
 
-  @parameterized.product(
-      optimizer_fn=[
-          lambda x: sgdm.build_sgdm(learning_rate=x),
-      ],
-      aggregation_factory=[
-          model_update_aggregator.robust_aggregator,
-          model_update_aggregator.compression_aggregator,
-          model_update_aggregator.secure_aggregator,
-      ],
-  )
-  def test_construction_calls_model_fn(self, optimizer_fn, aggregation_factory):
+  def test_construction_calls_model_fn(self):
     # Assert that the process building does not call `model_fn` too many times.
     # `model_fn` can potentially be expensive (loading weights, processing, etc
     # ).
@@ -49,8 +38,8 @@ def test_construction_calls_model_fn(self, optimizer_fn, aggregation_factory):
     fed_avg_with_optimizer_schedule.build_weighted_fed_avg_with_optimizer_schedule(
         model_fn=mock_model_fn,
         client_learning_rate_fn=learning_rate_fn,
-        client_optimizer_fn=optimizer_fn,
-        model_aggregator=aggregation_factory(),
+        client_optimizer_fn=lambda lr: sgdm.build_sgdm(learning_rate=lr),
+        model_aggregator=model_update_aggregator.robust_aggregator(),
     )
     self.assertEqual(mock_model_fn.call_count, 3)
 
@@ -143,21 +132,6 @@ def test_raises_on_non_callable_model_fn(self):
           client_optimizer_fn=lambda _: sgdm.build_sgdm(),
       )
 
-  def test_construction_with_only_secure_aggregation(self):
-    model_fn = model_examples.LinearRegression
-    learning_process = fed_avg_with_optimizer_schedule.build_weighted_fed_avg_with_optimizer_schedule(
-        model_fn,
-        client_learning_rate_fn=lambda x: 0.5,
-        client_optimizer_fn=lambda x: sgdm.build_sgdm(),
-        model_aggregator=model_update_aggregator.secure_aggregator(
-            weighted=True
-        ),
-        metrics_aggregator=aggregator.secure_sum_then_finalize,
-    )
-    static_assert.assert_not_contains_unsecure_aggregation(
-        learning_process.next
-    )
-
   def test_measurements_include_client_learning_rate(self):
     client_work = fed_avg_with_optimizer_schedule.build_scheduled_client_work(
         model_fn=model_examples.LinearRegression,

tensorflow_federated/python/learning/algorithms/fed_prox_test.py

@@ -19,11 +19,9 @@
 
 from tensorflow_federated.python.aggregators import factory_utils
 from tensorflow_federated.python.core.templates import iterative_process
-from tensorflow_federated.python.core.test import static_assert
 from tensorflow_federated.python.learning import loop_builder
 from tensorflow_federated.python.learning import model_update_aggregator
 from tensorflow_federated.python.learning.algorithms import fed_prox
-from tensorflow_federated.python.learning.metrics import aggregator
 from tensorflow_federated.python.learning.models import model_examples
 from tensorflow_federated.python.learning.models import model_weights
 from tensorflow_federated.python.learning.models import test_models
@@ -34,26 +32,16 @@
 class FedProxConstructionTest(parameterized.TestCase):
   """Tests construction of the FedProx training process."""
 
-  @parameterized.product(
-      optimizer_fn=[
-          sgdm.build_sgdm(learning_rate=0.1),
-      ],
-      aggregation_factory=[
-          model_update_aggregator.robust_aggregator,
-          model_update_aggregator.compression_aggregator,
-          model_update_aggregator.secure_aggregator,
-      ],
-  )
-  def test_construction_calls_model_fn(self, optimizer_fn, aggregation_factory):
+  def test_construction_calls_model_fn(self):
     # Assert that the process building does not call `model_fn` too many times.
     # `model_fn` can potentially be expensive (loading weights, processing, etc
     # ).
     mock_model_fn = mock.Mock(side_effect=model_examples.LinearRegression)
     fed_prox.build_weighted_fed_prox(
         model_fn=mock_model_fn,
         proximal_strength=1.0,
-        client_optimizer_fn=optimizer_fn,
-        model_aggregator=aggregation_factory(),
+        client_optimizer_fn=sgdm.build_sgdm(learning_rate=0.1),
+        model_aggregator=model_update_aggregator.robust_aggregator(),
     )
     self.assertEqual(mock_model_fn.call_count, 3)
 
@@ -160,36 +148,6 @@ def test_unweighted_fed_avg_raises_on_weighted_aggregator(self):
           model_aggregator=model_aggregator,
       )
 
-  def test_weighted_fed_prox_with_only_secure_aggregation(self):
-    model_fn = model_examples.LinearRegression
-    learning_process = fed_prox.build_weighted_fed_prox(
-        model_fn,
-        proximal_strength=1.0,
-        client_optimizer_fn=sgdm.build_sgdm(),
-        model_aggregator=model_update_aggregator.secure_aggregator(
-            weighted=True
-        ),
-        metrics_aggregator=aggregator.secure_sum_then_finalize,
-    )
-    static_assert.assert_not_contains_unsecure_aggregation(
-        learning_process.next
-    )
-
-  def test_unweighted_fed_prox_with_only_secure_aggregation(self):
-    model_fn = model_examples.LinearRegression
-    learning_process = fed_prox.build_unweighted_fed_prox(
-        model_fn,
-        proximal_strength=1.0,
-        client_optimizer_fn=sgdm.build_sgdm(),
-        model_aggregator=model_update_aggregator.secure_aggregator(
-            weighted=False
-        ),
-        metrics_aggregator=aggregator.secure_sum_then_finalize,
-    )
-    static_assert.assert_not_contains_unsecure_aggregation(
-        learning_process.next
-    )
-
 
 if __name__ == '__main__':
   absltest.main()

tensorflow_federated/python/learning/algorithms/fed_sgd_test.py

@@ -20,7 +20,6 @@
 import tensorflow as tf
 
 from tensorflow_federated.python.core.environments.tensorflow_backend import tensorflow_test_utils
-from tensorflow_federated.python.core.test import static_assert
 from tensorflow_federated.python.learning import loop_builder
 from tensorflow_federated.python.learning import model_update_aggregator
 from tensorflow_federated.python.learning.algorithms import fed_sgd
@@ -160,11 +159,6 @@ def test_client_tf_dataset_reduce_fn(self, loop_implementation, mock_method):
 
   @parameterized.named_parameters(
       ('robust_aggregator', model_update_aggregator.robust_aggregator),
-      (
-          'compression_aggregator',
-          model_update_aggregator.compression_aggregator,
-      ),
-      ('secure_aggreagtor', model_update_aggregator.secure_aggregator),
   )
   def test_construction_calls_model_fn(self, aggregation_factory):
     # Assert that the process building does not call `model_fn` too many times.
@@ -177,17 +171,6 @@ def test_construction_calls_model_fn(self, aggregation_factory):
     # TODO: b/186451541 - reduce the number of calls to model_fn.
     self.assertEqual(mock_model_fn.call_count, 3)
 
-  def test_no_unsecure_aggregation_with_secure_aggregator(self):
-    model_fn = model_examples.LinearRegression
-    learning_process = fed_sgd.build_fed_sgd(
-        model_fn,
-        model_aggregator=model_update_aggregator.secure_aggregator(),
-        metrics_aggregator=aggregator.secure_sum_then_finalize,
-    )
-    static_assert.assert_not_contains_unsecure_aggregation(
-        learning_process.next
-    )
-
 
 class FunctionalFederatedSgdTest(tf.test.TestCase, parameterized.TestCase):
 
@@ -276,17 +259,6 @@ def test_build_functional_fed_sgd_succeeds(self):
     model = _build_functional_model()
     fed_sgd.build_fed_sgd(model_fn=model)
 
-  def test_no_unsecure_aggregation_with_secure_aggregator(self):
-    model = _build_functional_model()
-    learning_process = fed_sgd.build_fed_sgd(
-        model,
-        model_aggregator=model_update_aggregator.secure_aggregator(),
-        metrics_aggregator=aggregator.secure_sum_then_finalize,
-    )
-    static_assert.assert_not_contains_unsecure_aggregation(
-        learning_process.next
-    )
-
 
 if __name__ == '__main__':
   tf.test.main()