Refactoring dataset_reduce module into loop_builder.

- Rename `build_dataset_reduce_fn` to `build_training_loop`
- Parameterize by a new enum `LoopImplementation`
- Update all previous usages

This refactoring will make it easier to introduce alternative looping strateies
in the future. Namely of interest is a `tf.foldl` based implementation that
works across "datasets-as-arrays" (tf.Tensor with the zeroth dimesion being the
number of examples) as introduced in https://arxiv.org/abs/2307.09619.

PiperOrigin-RevId: 656540932

Author: ZacharyGarrett

tensorflow_federated/python/learning/BUILD

@@ -46,14 +46,14 @@ py_library(
 )
 
 py_library(
-    name = "dataset_reduce",
-    srcs = ["dataset_reduce.py"],
+    name = "loop_builder",
+    srcs = ["loop_builder.py"],
 )
 
 py_cpu_gpu_test(
-    name = "dataset_reduce_test",
-    srcs = ["dataset_reduce_test.py"],
-    deps = [":dataset_reduce"],
+    name = "loop_builder_test",
+    srcs = ["loop_builder_test.py"],
+    deps = [":loop_builder"],
 )
 
 py_library(

tensorflow_federated/python/learning/algorithms/BUILD

@@ -69,7 +69,7 @@ py_cpu_gpu_test(
         ":fed_avg",
         "//tensorflow_federated/python/aggregators:factory_utils",
         "//tensorflow_federated/python/core/test:static_assert",
-        "//tensorflow_federated/python/learning:dataset_reduce",
+        "//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",
@@ -116,7 +116,7 @@ py_cpu_gpu_test(
     deps = [
         ":fed_avg_with_optimizer_schedule",
         "//tensorflow_federated/python/core/test:static_assert",
-        "//tensorflow_federated/python/learning:dataset_reduce",
+        "//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",
@@ -160,7 +160,7 @@ py_cpu_gpu_test(
         "//tensorflow_federated/python/aggregators:factory_utils",
         "//tensorflow_federated/python/core/templates:iterative_process",
         "//tensorflow_federated/python/core/test:static_assert",
-        "//tensorflow_federated/python/learning:dataset_reduce",
+        "//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",
@@ -293,7 +293,7 @@ py_library(
         "//tensorflow_federated/python/core/impl/types:computation_types",
         "//tensorflow_federated/python/core/impl/types:placements",
         "//tensorflow_federated/python/core/templates:measured_process",
-        "//tensorflow_federated/python/learning:dataset_reduce",
+        "//tensorflow_federated/python/learning:loop_builder",
         "//tensorflow_federated/python/learning:tensor_utils",
         "//tensorflow_federated/python/learning/metrics:aggregator",
         "//tensorflow_federated/python/learning/metrics:types",
@@ -317,7 +317,7 @@ py_cpu_gpu_test(
     deps = [
         ":fed_sgd",
         "//tensorflow_federated/python/core/test:static_assert",
-        "//tensorflow_federated/python/learning:dataset_reduce",
+        "//tensorflow_federated/python/learning:loop_builder",
         "//tensorflow_federated/python/learning:model_update_aggregator",
         "//tensorflow_federated/python/learning/metrics:aggregator",
         "//tensorflow_federated/python/learning/models:functional",
@@ -376,7 +376,7 @@ py_library(
         "//tensorflow_federated/python/core/impl/types:type_conversions",
         "//tensorflow_federated/python/core/templates:measured_process",
         "//tensorflow_federated/python/learning:client_weight_lib",
-        "//tensorflow_federated/python/learning:dataset_reduce",
+        "//tensorflow_federated/python/learning:loop_builder",
         "//tensorflow_federated/python/learning:tensor_utils",
         "//tensorflow_federated/python/learning/metrics:aggregator",
         "//tensorflow_federated/python/learning/metrics:types",
@@ -415,7 +415,7 @@ py_cpu_gpu_test(
         "//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:dataset_reduce",
+        "//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",
@@ -450,7 +450,7 @@ py_library(
         "//tensorflow_federated/python/core/impl/types:placements",
         "//tensorflow_federated/python/core/templates:aggregation_process",
         "//tensorflow_federated/python/core/templates:measured_process",
-        "//tensorflow_federated/python/learning:dataset_reduce",
+        "//tensorflow_federated/python/learning:loop_builder",
         "//tensorflow_federated/python/learning/metrics:sum_aggregation_factory",
         "//tensorflow_federated/python/learning/models:functional",
         "//tensorflow_federated/python/learning/models:model_weights",
@@ -520,7 +520,7 @@ py_cpu_gpu_test(
         ":personalization_eval",
         "//tensorflow_federated/python/core/backends/native:execution_contexts",
         "//tensorflow_federated/python/core/impl/types:computation_types",
-        "//tensorflow_federated/python/learning:dataset_reduce",
+        "//tensorflow_federated/python/learning:loop_builder",
         "//tensorflow_federated/python/learning/models:keras_utils",
         "//tensorflow_federated/python/learning/models:model_examples",
         "//tensorflow_federated/python/learning/models:model_weights",

tensorflow_federated/python/learning/algorithms/fed_avg_test.py

@@ -20,7 +20,7 @@
 
 from tensorflow_federated.python.aggregators import factory_utils
 from tensorflow_federated.python.core.test import static_assert
-from tensorflow_federated.python.learning import dataset_reduce
+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
@@ -60,9 +60,9 @@ def test_construction_calls_model_fn(self, optimizer_fn, aggregation_factory):
       ('simulation_tff_optimizer', True),
   )
   @mock.patch.object(
-      dataset_reduce,
+      loop_builder,
       '_dataset_reduce_fn',
-      wraps=dataset_reduce._dataset_reduce_fn,
+      wraps=loop_builder._dataset_reduce_fn,
   )
   def test_client_tf_dataset_reduce_fn(self, simulation, mock_method):
     fed_avg.build_weighted_fed_avg(

tensorflow_federated/python/learning/algorithms/fed_avg_with_optimizer_schedule_test.py

@@ -19,7 +19,7 @@
 import tensorflow as tf
 
 from tensorflow_federated.python.core.test import static_assert
-from tensorflow_federated.python.learning import dataset_reduce
+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
 from tensorflow_federated.python.learning.metrics import aggregator
@@ -68,9 +68,9 @@ def test_construction_of_functional_model(self):
       ('simulation', True),
   )
   @mock.patch.object(
-      dataset_reduce,
+      loop_builder,
       '_dataset_reduce_fn',
-      wraps=dataset_reduce._dataset_reduce_fn,
+      wraps=loop_builder._dataset_reduce_fn,
   )
   def test_client_tf_dataset_reduce_fn(self, use_simulation, mock_reduce):
     client_learning_rate_fn = lambda x: 0.5

tensorflow_federated/python/learning/algorithms/fed_eval.py

@@ -29,7 +29,7 @@
 from tensorflow_federated.python.core.impl.types import placements
 from tensorflow_federated.python.core.templates import aggregation_process
 from tensorflow_federated.python.core.templates import measured_process
-from tensorflow_federated.python.learning import dataset_reduce
+from tensorflow_federated.python.learning import loop_builder
 from tensorflow_federated.python.learning.metrics import sum_aggregation_factory
 from tensorflow_federated.python.learning.models import functional
 from tensorflow_federated.python.learning.models import model_weights as model_weights_lib
@@ -105,8 +105,10 @@ def reduce_fn(num_examples, batch):
       else:
         return num_examples + tf.cast(model_output.num_examples, tf.int64)
 
-    dataset_reduce_fn = dataset_reduce.build_dataset_reduce_fn(
-        use_experimental_simulation_loop
+    dataset_reduce_fn = loop_builder.build_training_loop(
+        loop_builder.LoopImplementation.DATASET_ITERATOR
+        if use_experimental_simulation_loop
+        else loop_builder.LoopImplementation.DATASET_REDUCE
     )
     num_examples = dataset_reduce_fn(
         reduce_fn, dataset, lambda: tf.zeros([], dtype=tf.int64)

tensorflow_federated/python/learning/algorithms/fed_prox_test.py

@@ -21,7 +21,7 @@
 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 dataset_reduce
+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
@@ -64,9 +64,9 @@ def test_construction_calls_model_fn(self, optimizer_fn, aggregation_factory):
       ('simulation_tff_optimizer', True),
   )
   @mock.patch.object(
-      dataset_reduce,
+      loop_builder,
       '_dataset_reduce_fn',
-      wraps=dataset_reduce._dataset_reduce_fn,
+      wraps=loop_builder._dataset_reduce_fn,
   )
   def test_client_tf_dataset_reduce_fn(self, simulation, mock_method):
     fed_prox.build_weighted_fed_prox(

tensorflow_federated/python/learning/algorithms/fed_sgd.py

@@ -37,7 +37,7 @@
 from tensorflow_federated.python.core.impl.types import computation_types
 from tensorflow_federated.python.core.impl.types import placements
 from tensorflow_federated.python.core.templates import measured_process
-from tensorflow_federated.python.learning import dataset_reduce
+from tensorflow_federated.python.learning import loop_builder
 from tensorflow_federated.python.learning import tensor_utils
 from tensorflow_federated.python.learning.metrics import aggregator as metric_aggregator
 from tensorflow_federated.python.learning.metrics import types
@@ -65,8 +65,10 @@ def _build_client_update(
   Returns:
     A `tf.function`.
   """
-  dataset_reduce_fn = dataset_reduce.build_dataset_reduce_fn(
-      use_experimental_simulation_loop
+  dataset_reduce_fn = loop_builder.build_training_loop(
+      loop_builder.LoopImplementation.DATASET_ITERATOR
+      if use_experimental_simulation_loop
+      else loop_builder.LoopImplementation.DATASET_REDUCE
   )
 
   @tf.function
@@ -215,8 +217,10 @@ def _build_functional_client_update(
   Returns:
     A `tf.function`.
   """
-  dataset_reduce_fn = dataset_reduce.build_dataset_reduce_fn(
-      use_experimental_simulation_loop
+  dataset_reduce_fn = loop_builder.build_training_loop(
+      loop_builder.LoopImplementation.DATASET_ITERATOR
+      if use_experimental_simulation_loop
+      else loop_builder.LoopImplementation.DATASET_REDUCE
   )
 
   @tf.function

tensorflow_federated/python/learning/algorithms/fed_sgd_test.py

@@ -20,7 +20,7 @@
 import tensorflow as tf
 
 from tensorflow_federated.python.core.test import static_assert
-from tensorflow_federated.python.learning import dataset_reduce
+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
 from tensorflow_federated.python.learning.metrics import aggregator
@@ -143,9 +143,9 @@ def test_non_finite_aggregation(self, bad_value):
       ('non-simulation', False), ('simulation', True)
   )
   @mock.patch.object(
-      dataset_reduce,
+      loop_builder,
       '_dataset_reduce_fn',
-      wraps=dataset_reduce._dataset_reduce_fn,
+      wraps=loop_builder._dataset_reduce_fn,
   )
   @tensorflow_test_utils.skip_test_for_multi_gpu
   def test_client_tf_dataset_reduce_fn(self, simulation, mock_method):
@@ -250,9 +250,9 @@ def test_non_finite_aggregation(self, bad_value):
       ('non-simulation', False), ('simulation', True)
   )
   @mock.patch.object(
-      dataset_reduce,
+      loop_builder,
       '_dataset_reduce_fn',
-      wraps=dataset_reduce._dataset_reduce_fn,
+      wraps=loop_builder._dataset_reduce_fn,
   )
   @tensorflow_test_utils.skip_test_for_multi_gpu
   def test_client_tf_dataset_reduce_fn(self, simulation, mock_method):

tensorflow_federated/python/learning/algorithms/mime.py

@@ -41,7 +41,7 @@
 from tensorflow_federated.python.core.impl.types import type_conversions
 from tensorflow_federated.python.core.templates import measured_process
 from tensorflow_federated.python.learning import client_weight_lib
-from tensorflow_federated.python.learning import dataset_reduce
+from tensorflow_federated.python.learning import loop_builder
 from tensorflow_federated.python.learning import tensor_utils
 from tensorflow_federated.python.learning.metrics import aggregator as metric_aggregator
 from tensorflow_federated.python.learning.metrics import types
@@ -80,8 +80,10 @@ def _build_client_update_fn_for_mime_lite(
   @tensorflow_computation.tf_computation
   def client_update_fn(global_optimizer_state, initial_weights, data):
     model = model_fn()
-    dataset_reduce_fn = dataset_reduce.build_dataset_reduce_fn(
-        use_experimental_simulation_loop
+    dataset_reduce_fn = loop_builder.build_training_loop(
+        loop_builder.LoopImplementation.DATASET_ITERATOR
+        if use_experimental_simulation_loop
+        else loop_builder.LoopImplementation.DATASET_REDUCE
     )
     weight_tensor_specs = type_conversions.type_to_tf_tensor_specs(
         model_weights_lib.weights_type_from_model(model)
@@ -320,8 +322,10 @@ def _build_functional_client_update_fn_for_mime_lite(
 
   @tensorflow_computation.tf_computation
   def client_update_fn(global_optimizer_state, incoming_weights, data):
-    dataset_reduce_fn = dataset_reduce.build_dataset_reduce_fn(
-        use_experimental_simulation_loop
+    dataset_reduce_fn = loop_builder.build_training_loop(
+        loop_builder.LoopImplementation.DATASET_ITERATOR
+        if use_experimental_simulation_loop
+        else loop_builder.LoopImplementation.DATASET_REDUCE
     )
     weight_tensor_specs = tf.nest.map_structure(
         lambda t: tf.TensorSpec(shape=t.shape, dtype=t.dtype),

tensorflow_federated/python/learning/algorithms/mime_test.py

@@ -31,7 +31,7 @@
 from tensorflow_federated.python.core.templates import measured_process
 from tensorflow_federated.python.core.test import static_assert
 from tensorflow_federated.python.learning import client_weight_lib
-from tensorflow_federated.python.learning import dataset_reduce
+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.algorithms import mime
@@ -191,9 +191,9 @@ class MimeLiteClientWorkExecutionTest(tf.test.TestCase, parameterized.TestCase):
       ('non-simulation', False), ('simulation', True)
   )
   @mock.patch.object(
-      dataset_reduce,
+      loop_builder,
       '_dataset_reduce_fn',
-      wraps=dataset_reduce._dataset_reduce_fn,
+      wraps=loop_builder._dataset_reduce_fn,
   )
   @tensorflow_test_utils.skip_test_for_multi_gpu
   def test_client_tf_dataset_reduce_fn(self, simulation, mock_method):
@@ -280,9 +280,9 @@ class MimeLiteFunctionalClientWorkExecutionTest(
       ('non-simulation', False), ('simulation', True)
   )
   @mock.patch.object(
-      dataset_reduce,
+      loop_builder,
       '_dataset_reduce_fn',
-      wraps=dataset_reduce._dataset_reduce_fn,
+      wraps=loop_builder._dataset_reduce_fn,
   )
   @tensorflow_test_utils.skip_test_for_gpu
   def test_client_tf_dataset_reduce_fn(self, simulation, mock_method):
@@ -390,9 +390,9 @@ def test_construction_calls_model_fn(self):
       ('simulation_tff_optimizer', True),
   )
   @mock.patch.object(
-      dataset_reduce,
+      loop_builder,
       '_dataset_reduce_fn',
-      wraps=dataset_reduce._dataset_reduce_fn,
+      wraps=loop_builder._dataset_reduce_fn,
   )
   def test_client_tf_dataset_reduce_fn(self, simulation, mock_method):
     mime.build_weighted_mime_lite(