def test_federated_secure_modular_sum(self, value_dtype, modulus_type):
uri = intrinsic_defs.FEDERATED_SECURE_MODULAR_SUM.uri
comp = building_blocks.Intrinsic(
uri,
computation_types.FunctionType(
parameter=[
computation_types.at_clients(value_dtype),
computation_types.to_type(modulus_type)
],
result=computation_types.at_server(value_dtype)))
# First without secure intrinsics shouldn't modify anything.
reduced, modified = tree_transformations.replace_intrinsics_with_bodies(
comp)
self.assertFalse(modified)
self.assertGreater(_count_intrinsics(comp, uri), 0)
type_test_utils.assert_types_identical(comp.type_signature,
reduced.type_signature)
# Now replace bodies including secure intrinsics.
reduced, modified = tree_transformations.replace_secure_intrinsics_with_insecure_bodies(
comp)
self.assertTrue(modified)
# Inserting tensorflow, as we do here, does not preserve python containers
# currently.
type_test_utils.assert_types_equivalent(comp.type_signature,
reduced.type_signature)
self.assertGreater(
_count_intrinsics(reduced, intrinsic_defs.FEDERATED_SUM.uri), 0)
def test_federated_secure_select(self):
uri = intrinsic_defs.FEDERATED_SECURE_SELECT.uri
comp = building_blocks.Intrinsic(
uri,
computation_types.FunctionType(
[
computation_types.at_clients(tf.int32), # client_keys
computation_types.at_server(tf.int32), # max_key
computation_types.at_server(tf.float32), # server_state
computation_types.FunctionType([tf.float32, tf.int32],
tf.float32) # select_fn
],
computation_types.at_clients(
computation_types.SequenceType(tf.float32))))
self.assertGreater(_count_intrinsics(comp, uri), 0)
# First without secure intrinsics shouldn't modify anything.
reduced, modified = tree_transformations.replace_intrinsics_with_bodies(
comp)
self.assertFalse(modified)
self.assertGreater(_count_intrinsics(comp, uri), 0)
type_test_utils.assert_types_identical(comp.type_signature,
reduced.type_signature)
# Now replace bodies including secure intrinsics.
reduced, modified = tree_transformations.replace_secure_intrinsics_with_insecure_bodies(
comp)
self.assertTrue(modified)
type_test_utils.assert_types_identical(comp.type_signature,
reduced.type_signature)
self.assertGreater(
_count_intrinsics(reduced, intrinsic_defs.FEDERATED_SELECT.uri), 0)
def test_federated_secure_sum_bitwidth(self, value_dtype, bitwidth_type):
uri = intrinsic_defs.FEDERATED_SECURE_SUM_BITWIDTH.uri
comp = building_blocks.Intrinsic(
uri,
computation_types.FunctionType(
parameter=[
computation_types.at_clients(value_dtype),
computation_types.to_type(bitwidth_type)
],
result=computation_types.at_server(value_dtype)))
# First without secure intrinsics shouldn't modify anything.
reduced, modified = tree_transformations.replace_intrinsics_with_bodies(
comp)
self.assertFalse(modified)
self.assertGreater(_count_intrinsics(comp, uri), 0)
type_test_utils.assert_types_identical(comp.type_signature,
reduced.type_signature)
# Now replace bodies including secure intrinsics.
reduced, modified = tree_transformations.replace_secure_intrinsics_with_insecure_bodies(
comp)
self.assertTrue(modified)
type_test_utils.assert_types_identical(comp.type_signature,
reduced.type_signature)
self.assertGreater(
_count_intrinsics(reduced, intrinsic_defs.FEDERATED_AGGREGATE.uri),
0)
def test_generic_plus_reduces(self):
uri = intrinsic_defs.GENERIC_PLUS.uri
comp = building_blocks.Intrinsic(
uri,
computation_types.FunctionType([tf.float32, tf.float32],
tf.float32))
count_before_reduction = _count_intrinsics(comp, uri)
reduced, modified = tree_transformations.replace_intrinsics_with_bodies(
comp)
count_after_reduction = _count_intrinsics(reduced, uri)
self.assertTrue(modified)
type_test_utils.assert_types_identical(comp.type_signature,
reduced.type_signature)
self.assertGreater(count_before_reduction, 0)
self.assertEqual(count_after_reduction, 0)
tree_analysis.check_contains_only_reducible_intrinsics(reduced)
def test_federated_sum_reduces_to_aggregate(self):
uri = intrinsic_defs.FEDERATED_SUM.uri
comp = building_blocks.Intrinsic(
uri,
computation_types.FunctionType(
computation_types.at_clients(tf.float32),
computation_types.at_server(tf.float32)))
count_sum_before_reduction = _count_intrinsics(comp, uri)
reduced, modified = tree_transformations.replace_intrinsics_with_bodies(
comp)
count_sum_after_reduction = _count_intrinsics(reduced, uri)
count_aggregations = _count_intrinsics(
reduced, intrinsic_defs.FEDERATED_AGGREGATE.uri)
self.assertTrue(modified)
type_test_utils.assert_types_identical(comp.type_signature,
reduced.type_signature)
self.assertGreater(count_sum_before_reduction, 0)
self.assertEqual(count_sum_after_reduction, 0)
self.assertGreater(count_aggregations, 0)
def desugar_and_transform_to_native(comp):
"""Transform to native form and replace intrinsics with TensorFlow."""
# Turn on static grappler. The function inlining is critical for GPU support,
# otherwise variant placeholders that received datasets will be placed on GPUs
# which don't have kernels for datastes, causing TF to error.
grappler_config = tf.compat.v1.ConfigProto()
aggressive = grappler_config.graph_options.rewrite_options.AGGRESSIVE
rewrite_options = grappler_config.graph_options.rewrite_options
rewrite_options.memory_optimization = aggressive
rewrite_options.constant_folding = aggressive
rewrite_options.arithmetic_optimization = aggressive
rewrite_options.loop_optimization = aggressive
rewrite_options.function_optimization = aggressive
intrinsics_desugared_bb, _ = tree_transformations.replace_intrinsics_with_bodies(
comp.to_building_block())
# Desugaring intrinsics injects TF computations; transforming to native form
# adds TF cache IDs to them. It is crucial that these transformations execute
# in this order.
native_form = transform_to_native_form(
computation_impl.ConcreteComputation.from_building_block(
intrinsics_desugared_bb),
grappler_config=grappler_config)
return native_form
def get_broadcast_form_for_computation(
comp: computation_base.Computation,
grappler_config: tf.compat.v1.ConfigProto = _GRAPPLER_DEFAULT_CONFIG,
*,
tff_internal_preprocessing: Optional[BuildingBlockFn] = None,
) -> forms.BroadcastForm:
"""Constructs `tff.backends.mapreduce.BroadcastForm` given a computation.
Args:
comp: An instance of `tff.Computation` that is compatible with broadcast
form. Computations are only compatible if they take in a single value
placed at server, return a single value placed at clients, and do not
contain any aggregations.
grappler_config: An instance of `tf.compat.v1.ConfigProto` to configure
Grappler graph optimization of the Tensorflow graphs backing the resulting
`tff.backends.mapreduce.BroadcastForm`. These options are combined with a
set of defaults that aggressively configure Grappler. If
`grappler_config_proto` has
`graph_options.rewrite_options.disable_meta_optimizer=True`, Grappler is
bypassed.
tff_internal_preprocessing: An optional function to transform the AST of the
computation.
Returns:
An instance of `tff.backends.mapreduce.BroadcastForm` equivalent to the
provided `tff.Computation`.
"""
py_typecheck.check_type(comp, computation_base.Computation)
_check_function_signature_compatible_with_broadcast_form(
comp.type_signature)
py_typecheck.check_type(grappler_config, tf.compat.v1.ConfigProto)
grappler_config = _merge_grappler_config_with_default(grappler_config)
bb = comp.to_building_block()
if tff_internal_preprocessing:
bb = tff_internal_preprocessing(bb)
bb, _ = tree_transformations.replace_intrinsics_with_bodies(bb)
bb = _replace_lambda_body_with_call_dominant_form(bb)
tree_analysis.check_contains_only_reducible_intrinsics(bb)
aggregations = tree_analysis.find_aggregations_in_tree(bb)
if aggregations:
raise ValueError(
f'`get_broadcast_form_for_computation` called with computation '
f'containing {len(aggregations)} aggregations, but broadcast form '
'does not allow aggregation. Full list of aggregations:\n{aggregations}'
)
before_broadcast, after_broadcast = _split_ast_on_broadcast(bb)
compute_server_context = _extract_compute_server_context(
before_broadcast, grappler_config)
client_processing = _extract_client_processing(after_broadcast,
grappler_config)
compute_server_context, client_processing = (
computation_impl.ConcreteComputation.from_building_block(bb)
for bb in (compute_server_context, client_processing))
comp_param_names = structure.name_list_with_nones(
comp.type_signature.parameter)
server_data_label, client_data_label = comp_param_names
return forms.BroadcastForm(compute_server_context,
client_processing,
server_data_label=server_data_label,
client_data_label=client_data_label)
def check_iterative_process_compatible_with_map_reduce_form(
ip: iterative_process.IterativeProcess,
*,
tff_internal_preprocessing: Optional[BuildingBlockFn] = None,
) -> Tuple[building_blocks.ComputationBuildingBlock,
building_blocks.ComputationBuildingBlock]:
"""Tests compatibility with `tff.backends.mapreduce.MapReduceForm`.
Note: the conditions here are specified in the documentation for
`get_map_reduce_form_for_iterative_process`. Changes to this function should
be propagated to that documentation.
Args:
ip: An instance of `tff.templates.IterativeProcess` to check for
compatibility with `tff.backends.mapreduce.MapReduceForm`.
tff_internal_preprocessing: An optional function to transform the AST of the
computation.
Returns:
TFF-internal building-blocks representing the validated and simplified
`initialize` and `next` computations.
Raises:
TypeError: If the arguments are of the wrong types.
"""
py_typecheck.check_type(ip, iterative_process.IterativeProcess)
initialize_tree = ip.initialize.to_building_block()
next_tree = ip.next.to_building_block()
if tff_internal_preprocessing:
initialize_tree = tff_internal_preprocessing(initialize_tree)
next_tree = tff_internal_preprocessing(next_tree)
init_type = initialize_tree.type_signature
_check_type_is_no_arg_fn(init_type, '`initialize`', TypeError)
if (not init_type.result.is_federated()
or init_type.result.placement != placements.SERVER):
raise TypeError(
'Expected `initialize` to return a single federated value '
'placed at server (type `T@SERVER`), found return type:\n'
f'{init_type.result}')
next_type = next_tree.type_signature
_check_type_is_fn(next_type, '`next`', TypeError)
if not next_type.parameter.is_struct() or len(next_type.parameter) != 2:
raise TypeError(
'Expected `next` to take two arguments, found parameter '
f' type:\n{next_type.parameter}')
if not next_type.result.is_struct() or len(next_type.result) != 2:
raise TypeError('Expected `next` to return two values, found result '
f'type:\n{next_type.result}')
initialize_tree, _ = tree_transformations.replace_intrinsics_with_bodies(
initialize_tree)
next_tree, _ = tree_transformations.replace_intrinsics_with_bodies(
next_tree)
next_tree = _replace_lambda_body_with_call_dominant_form(next_tree)
tree_analysis.check_contains_only_reducible_intrinsics(initialize_tree)
tree_analysis.check_contains_only_reducible_intrinsics(next_tree)
tree_analysis.check_broadcast_not_dependent_on_aggregate(next_tree)
return initialize_tree, next_tree
def compile_to_mergeable_comp_form(
comp: computation_impl.ConcreteComputation
) -> mergeable_comp_execution_context.MergeableCompForm:
"""Compiles a computation with a single aggregation to `MergeableCompForm`.
Compilation proceeds by splitting on the lone aggregation, and using the
aggregation's internal functions to generate a semantically equivalent
instance of `mergeable_comp_execution_context.MergeableCompForm`.
Args:
comp: Instance of `computation_impl.ConcreteComputation` to compile. Assumed
to be representable as a computation with a single aggregation in its
body, so that for example two parallel aggregations are allowed, but
multiple dependent aggregations are disallowed. Additionally assumed to be
of functional type.
Returns:
A semantically equivalent instance of
`mergeable_comp_execution_context.MergeableCompForm`.
Raises:
TypeError: If `comp` is not a building block, or is not of functional TFF
type.
ValueError: If `comp` cannot be represented as a computation with at most
one aggregation in its body.
"""
original_return_type = comp.type_signature.result
building_block = comp.to_building_block()
lam = _ensure_lambda(building_block)
lowered_bb, _ = tree_transformations.replace_intrinsics_with_bodies(lam)
# We transform the body of this computation to easily preserve the top-level
# lambda required by force-aligning.
call_dominant_body_bb = transformations.to_call_dominant(lowered_bb.result)
call_dominant_bb = building_blocks.Lambda(lowered_bb.parameter_name,
lowered_bb.parameter_type,
call_dominant_body_bb)
# This check should not throw false positives because we just ensured we are
# in call-dominant form.
tree_analysis.check_aggregate_not_dependent_on_aggregate(call_dominant_bb)
before_agg, after_agg = transformations.force_align_and_split_by_intrinsics(
call_dominant_bb,
[building_block_factory.create_null_federated_aggregate()])
# Construct a report function which accepts the result of merge.
merge_fn_type = before_agg.type_signature.result[
'federated_aggregate_param'][3]
identity_report = computation_impl.ConcreteComputation.from_building_block(
building_block_factory.create_compiled_identity(merge_fn_type.result))
zero_comp, accumulate_comp, merge_comp, report_comp = _extract_federated_aggregate_computations(
before_agg)
before_agg_callable = computation_impl.ConcreteComputation.from_building_block(
before_agg)
after_agg_callable = computation_impl.ConcreteComputation.from_building_block(
after_agg)
if before_agg.type_signature.parameter is not None:
# TODO(b/147499373): If None-arguments were uniformly represented as empty
# tuples, we would be able to avoid this (and related) ugly casing.
@federated_computation.federated_computation(
before_agg.type_signature.parameter)
def up_to_merge_computation(arg):
federated_aggregate_args = before_agg_callable(
arg)['federated_aggregate_param']
value_to_aggregate = federated_aggregate_args[0]
zero = zero_comp()
return intrinsics.federated_aggregate(value_to_aggregate, zero,
accumulate_comp, merge_comp,
identity_report)
@federated_computation.federated_computation(
before_agg.type_signature.parameter,
computation_types.at_server(identity_report.type_signature.result))
def after_merge_computation(top_level_arg, merge_result):
reported_result = intrinsics.federated_map(report_comp,
merge_result)
return after_agg_callable(top_level_arg, [reported_result])
else:
@federated_computation.federated_computation()
def up_to_merge_computation():
federated_aggregate_args = before_agg_callable(
)['federated_aggregate_param']
value_to_aggregate = federated_aggregate_args[0]
zero = zero_comp()
return intrinsics.federated_aggregate(value_to_aggregate, zero,
accumulate_comp, merge_comp,
identity_report)
@federated_computation.federated_computation(
computation_types.at_server(identity_report.type_signature.result))
def after_merge_computation(merge_result):
reported_result = intrinsics.federated_map(report_comp,
merge_result)
return after_agg_callable([[reported_result]])
annotated_type_signature = computation_types.FunctionType(
after_merge_computation.type_signature.parameter, original_return_type)
after_merge_computation = computation_impl.ConcreteComputation.with_type(
after_merge_computation, annotated_type_signature)
return mergeable_comp_execution_context.MergeableCompForm(
up_to_merge=up_to_merge_computation,
merge=merge_comp,
after_merge=after_merge_computation)
def test_raises_on_none(self):
with self.assertRaises(TypeError):
tree_transformations.replace_intrinsics_with_bodies(None)