Add integration test that exercises a tensorflow dataset.

PiperOrigin-RevId: 773442984

Author: TensorFlow Federated Team

tensorflow_federated/python/core/impl/executor_stacks/cpp_executor_factory.py

@@ -122,6 +122,12 @@ def local_cpp_executor_factory(
     ] = None,
 ) -> federated_language.framework.ExecutorFactory:
   """Local ExecutorFactory backed by C++ Executor bindings."""
+
+  print(
+      "dalyk in cpp_executor_factory local_cpp_executor_factory, meaning we're"
+      ' using pybind'
+  )
+
   _check_num_clients_is_valid(default_num_clients)
 
   def _executor_fn(

tensorflow_federated/python/core/impl/executor_stacks/executor_factory.py

@@ -60,6 +60,11 @@ def local_cpp_executor_factory(
     RuntimeError: If an internal C++ worker binary can not be found.
   """
 
+  print(
+      "dalyk in executor_factory local_cpp_executor_factory, meaning we're"
+      ' using grpc'
+  )
+
   # This path is specified relative to this file because the relative location
   # of the worker binary will remain the same when this function is executed
   # from the Python package and from a Bazel test.

tensorflow_federated/python/tests/backend_test.py

@@ -173,6 +173,41 @@ def foo():
     result = foo()
     self.assertEqual(result, 10)
 
+  @test_contexts.with_contexts(*test_contexts.get_all_contexts())
+  def test_dataset(self):
+
+    client_data_type = federated_language.FederatedType(
+        federated_language.TensorType(np.int32, [3]),
+        federated_language.CLIENTS,
+    )
+    server_state_type = federated_language.FederatedType(
+        np.int32, federated_language.SERVER
+    )
+
+    def reduce_fn(state, batch):
+      return state + batch
+
+    @tff.tensorflow.computation(client_data_type.member)
+    def client_transform(data):
+      dataset = tf.data.Dataset.from_tensor_slices(data)
+      return dataset.reduce(0, reduce_fn)
+
+    @federated_language.federated_computation(
+        [client_data_type, server_state_type]
+    )
+    def my_comp(client_data, server_data):
+      transformed_client_data = federated_language.federated_map(
+          client_transform, client_data
+      )
+      return (
+          federated_language.federated_sum(transformed_client_data),
+          server_data,
+      )
+
+    result_1, result_2 = my_comp([[1, 2, 3], [4, 5, 6]], 30)
+    self.assertEqual(result_1, 21)
+    self.assertEqual(result_2, 30)
+
   @test_contexts.with_contexts(*test_contexts.get_all_contexts())
   def test_empty_tuple(self):