def test_removes_selected_intrinsic_leaving_remaining_intrinsic(self): federated_aggregate = building_block_test_utils.create_whimsy_called_federated_aggregate( accumulate_parameter_name='a', merge_parameter_name='b', report_parameter_name='c') federated_secure_sum_bitwidth = building_block_test_utils.create_whimsy_called_federated_secure_sum_bitwidth( ) called_intrinsics = building_blocks.Struct([ federated_aggregate, federated_secure_sum_bitwidth, ]) comp = building_blocks.Lambda('d', tf.int32, called_intrinsics) null_aggregate = building_block_factory.create_null_federated_aggregate( ) secure_sum_bitwidth_uri = federated_secure_sum_bitwidth.function.uri aggregate_uri = null_aggregate.function.uri before, after = transformations.force_align_and_split_by_intrinsics( comp, [null_aggregate]) self.assertTrue( tree_analysis.contains_called_intrinsic(comp, secure_sum_bitwidth_uri)) self.assertTrue( tree_analysis.contains_called_intrinsic(comp, aggregate_uri)) self.assertFalse( tree_analysis.contains_called_intrinsic(before, aggregate_uri)) self.assertFalse( tree_analysis.contains_called_intrinsic(after, aggregate_uri)) self.assertTrue( tree_analysis.contains_called_intrinsic(before, secure_sum_bitwidth_uri) or tree_analysis.contains_called_intrinsic( after, secure_sum_bitwidth_uri))
def test_splits_on_selected_intrinsic_aggregate(self): federated_aggregate = building_block_test_utils.create_whimsy_called_federated_aggregate( accumulate_parameter_name='a', merge_parameter_name='b', report_parameter_name='c') called_intrinsics = building_blocks.Struct([federated_aggregate]) comp = building_blocks.Lambda('d', tf.int32, called_intrinsics) call = building_block_factory.create_null_federated_aggregate() self.assert_splits_on(comp, call)
def test_splits_even_when_selected_intrinsic_is_not_present(self): federated_aggregate = compiler_test_utils.create_whimsy_called_federated_aggregate( accumulate_parameter_name='a', merge_parameter_name='b', report_parameter_name='c') called_intrinsics = building_blocks.Struct([federated_aggregate]) comp = building_blocks.Lambda('d', tf.int32, called_intrinsics) transformations.force_align_and_split_by_intrinsics(comp, [ building_block_factory.create_null_federated_aggregate(), building_block_factory.create_null_federated_secure_sum_bitwidth(), ])
def test_splits_on_two_intrinsics(self): federated_aggregate = building_block_test_utils.create_whimsy_called_federated_aggregate( accumulate_parameter_name='a', merge_parameter_name='b', report_parameter_name='c') federated_secure_sum_bitwidth = building_block_test_utils.create_whimsy_called_federated_secure_sum_bitwidth( ) called_intrinsics = building_blocks.Struct([ federated_aggregate, federated_secure_sum_bitwidth, ]) comp = building_blocks.Lambda('d', tf.int32, called_intrinsics) self.assert_splits_on(comp, [ building_block_factory.create_null_federated_aggregate(), building_block_factory.create_null_federated_secure_sum_bitwidth() ])
def _split_ast_on_aggregate(bb): """Splits an AST on reduced aggregation intrinsics. Args: bb: An AST containing `federated_aggregate` or `federated_secure_sum_bitwidth` aggregations. Returns: Two ASTs, the first of which maps comp's input to the arguments to `federated_aggregate` and `federated_secure_sum_bitwidth`, and the second of which maps comp's input and the output of `federated_aggregate` and `federated_secure_sum_bitwidth` to comp's output. """ return transformations.force_align_and_split_by_intrinsics( bb, [ building_block_factory.create_null_federated_aggregate(), building_block_factory.create_null_federated_secure_sum_bitwidth() ])
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)