Amend visibility of some execution context symbols.
@@ -20,13 +20,20 @@ from tensorflow_federated.python.core.impl.executors import executor_stacks
from tensorflow_federated.python.core.impl.wrappers import computation_wrapper_instances
-def create_test_execution_context(default_num_clients=0, clients_per_thread=1):
- """Creates an execution context that executes computations locally."""
- factory = executor_stacks.local_executor_factory(
- default_num_clients=default_num_clients,
- clients_per_thread=clients_per_thread)
+def replace_secure_intrinsics_with_bodies(comp):
+ """Replace `secure_...` intrinsics with insecure TensorFlow equivalents.
+
+ Designed for use in tests, this function replaces
+ `tff.federated_secure_{sum, sum_bitwidth, modular_sum}` usages with equivalent
+ TensorFlow computations. The resulting computation can then be run on TFF
+ runtimes which do not implement secure computation.
+
+ Args:
+ comp: The computation to transform.
- def compiler(comp):
+ Returns:
+ `comp` with secure intrinsics replaced with insecure TensorFlow equivalents.
+ """
# Compile secure_sum and secure_sum_bitwidth intrinsics to insecure
# TensorFlow computations for testing purposes.
replaced_intrinsic_bodies, _ = intrinsic_reductions.replace_secure_intrinsics_with_insecure_bodies(
@@ -34,8 +41,15 @@ def create_test_execution_context(default_num_clients=0, clients_per_thread=1):
return computation_wrapper_instances.building_block_to_computation(
replaced_intrinsic_bodies)
+
+def create_test_execution_context(default_num_clients=0, clients_per_thread=1):
+ """Creates an execution context that executes computations locally."""
+ factory = executor_stacks.local_executor_factory(
+ default_num_clients=default_num_clients,
+ clients_per_thread=clients_per_thread)
+
return sync_execution_context.ExecutionContext(
- executor_fn=factory, compiler_fn=compiler)
+ executor_fn=factory, compiler_fn=replace_secure_intrinsics_with_bodies)
def set_test_execution_context(default_num_clients=0, clients_per_thread=1):
