def test_aggregate_with_selection_from_block_by_name_results_in_single_aggregate( self): data = building_blocks.Reference( 'a', computation_types.FederatedType(tf.int32, placement_literals.CLIENTS)) tup_of_data = building_blocks.Tuple([('a', data), ('b', data)]) block_holding_tup = building_blocks.Block([], tup_of_data) index_0_from_block = building_blocks.Selection( source=block_holding_tup, name='a') index_1_from_block = building_blocks.Selection( source=block_holding_tup, name='b') result = building_blocks.Data('aggregation_result', tf.int32) zero = building_blocks.Data('zero', tf.int32) accumulate = building_blocks.Lambda('accumulate_param', [tf.int32, tf.int32], result) merge = building_blocks.Lambda('merge_param', [tf.int32, tf.int32], result) report = building_blocks.Lambda('report_param', tf.int32, result) called_intrinsic0 = building_block_factory.create_federated_aggregate( index_0_from_block, zero, accumulate, merge, report) called_intrinsic1 = building_block_factory.create_federated_aggregate( index_1_from_block, zero, accumulate, merge, report) calls = building_blocks.Tuple((called_intrinsic0, called_intrinsic1)) comp = calls deduped_and_merged_comp, deduped_modified = transformations.dedupe_and_merge_tuple_intrinsics( comp, intrinsic_defs.FEDERATED_AGGREGATE.uri) self.assertTrue(deduped_modified) fed_agg = [] def _find_called_federated_aggregate(comp): if (isinstance(comp, building_blocks.Call) and isinstance(comp.function, building_blocks.Intrinsic) and comp.function.uri == intrinsic_defs.FEDERATED_AGGREGATE.uri): fed_agg.append(comp.function) return comp, False transformation_utils.transform_postorder( deduped_and_merged_comp, _find_called_federated_aggregate) self.assertLen(fed_agg, 1) self.assertEqual( fed_agg[0].type_signature.parameter[0].compact_representation(), '{<int32>}@CLIENTS')
def create_dummy_called_federated_aggregate(accumulate_parameter_name, merge_parameter_name, report_parameter_name): r"""Returns a dummy called federated aggregate. Call / \ federated_aggregate Tuple | [data, data, Lambda(x), Lambda(x), Lambda(x)] | | | data data data Args: accumulate_parameter_name: The name of the accumulate parameter. merge_parameter_name: The name of the merge parameter. report_parameter_name: The name of the report parameter. """ value_type = computation_types.FederatedType(tf.int32, placements.CLIENTS) value = building_blocks.Data('data', value_type) zero = building_blocks.Data('data', tf.float32) accumulate_type = computation_types.NamedTupleType((tf.float32, tf.int32)) accumulate_result = building_blocks.Data('data', tf.float32) accumulate = building_blocks.Lambda(accumulate_parameter_name, accumulate_type, accumulate_result) merge_type = computation_types.NamedTupleType((tf.float32, tf.float32)) merge_result = building_blocks.Data('data', tf.float32) merge = building_blocks.Lambda(merge_parameter_name, merge_type, merge_result) report_result = building_blocks.Data('data', tf.bool) report = building_blocks.Lambda(report_parameter_name, tf.float32, report_result) return building_block_factory.create_federated_aggregate( value, zero, accumulate, merge, report)
def federated_secure_sum(arg): py_typecheck.check_type(arg, building_blocks.ComputationBuildingBlock) summand_arg = building_blocks.Selection(arg, index=0) summand_type = summand_arg.type_signature.member max_input_arg = building_blocks.Selection(arg, index=1) max_input_type = max_input_arg.type_signature # Add the max_value as a second value in the zero, so it can be read during # `accumulate` to ensure client summands are valid. The value will be # later dropped in `report`. # # While accumulating summands, we'll assert each summand is less than or # equal to max_input. Otherwise the comptuation should issue an error. summation_zero = building_block_factory.create_generic_constant( summand_type, 0) aggregation_zero = building_blocks.Struct([summation_zero, max_input_arg], container_type=tuple) def assert_less_equal_max_and_add(summation_and_max_input, summand): summation, original_max_input = summation_and_max_input max_input = _ensure_structure(original_max_input, max_input_type, summand_type) # Assert that all coordinates in all tensors are less than the secure sum # allowed max input value. def assert_all_coordinates_less_equal(x, m): return tf.Assert( tf.reduce_all( tf.less_equal(tf.cast(x, tf.int64), tf.cast(m, tf.int64))), [ 'client value larger than maximum specified for secure sum', x, 'not less than or equal to', m ]) assert_ops = structure.flatten( structure.map_structure(assert_all_coordinates_less_equal, summand, max_input)) with tf.control_dependencies(assert_ops): return structure.map_structure(tf.add, summation, summand), original_max_input assert_less_equal_and_add = building_block_factory.create_tensorflow_binary_operator( assert_less_equal_max_and_add, operand_type=aggregation_zero.type_signature, second_operand_type=summand_type) def nested_plus(a, b): return structure.map_structure(tf.add, a, b) plus_op = building_block_factory.create_tensorflow_binary_operator( nested_plus, operand_type=aggregation_zero.type_signature) # In the `report` function we take the summation and drop the second element # of the struct (which was holding the max_value). drop_max_value_op = building_block_factory.create_tensorflow_unary_operator( lambda x: type_conversions.type_to_py_container(x[0], summand_type), aggregation_zero.type_signature) return building_block_factory.create_federated_aggregate( summand_arg, aggregation_zero, assert_less_equal_and_add, plus_op, drop_max_value_op)
def federated_aggregate(self, value, zero, accumulate, merge, report): """Implements `federated_aggregate` as defined in `api/intrinsics.py`. Args: value: As in `api/intrinsics.py`. zero: As in `api/intrinsics.py`. accumulate: As in `api/intrinsics.py`. merge: As in `api/intrinsics.py`. report: As in `api/intrinsics.py`. Returns: As in `api/intrinsics.py`. Raises: TypeError: As in `api/intrinsics.py`. """ value = value_impl.to_value(value, None, self._context_stack) value_utils.check_federated_value_placement(value, placements.CLIENTS, 'value to be aggregated') zero = value_impl.to_value(zero, None, self._context_stack) py_typecheck.check_type(zero, value_base.Value) # TODO(b/113112108): We need a check here that zero does not have federated # constituents. accumulate = value_impl.to_value(accumulate, None, self._context_stack) merge = value_impl.to_value(merge, None, self._context_stack) report = value_impl.to_value(report, None, self._context_stack) for op in [accumulate, merge, report]: py_typecheck.check_type(op, value_base.Value) py_typecheck.check_type(op.type_signature, computation_types.FunctionType) accumulate_type_expected = type_factory.reduction_op( zero.type_signature, value.type_signature.member) merge_type_expected = type_factory.reduction_op( zero.type_signature, zero.type_signature) report_type_expected = computation_types.FunctionType( zero.type_signature, report.type_signature.result) for op_name, op, type_expected in [ ('accumulate', accumulate, accumulate_type_expected), ('merge', merge, merge_type_expected), ('report', report, report_type_expected) ]: if not type_utils.is_assignable_from(type_expected, op.type_signature): raise TypeError( 'Expected parameter `{}` to be of type {}, but received {} instead.' .format(op_name, type_expected, op.type_signature)) value = value_impl.ValueImpl.get_comp(value) zero = value_impl.ValueImpl.get_comp(zero) accumulate = value_impl.ValueImpl.get_comp(accumulate) merge = value_impl.ValueImpl.get_comp(merge) report = value_impl.ValueImpl.get_comp(report) comp = building_block_factory.create_federated_aggregate( value, zero, accumulate, merge, report) return value_impl.ValueImpl(comp, self._context_stack)
def federated_sum(x): py_typecheck.check_type(x, building_blocks.ComputationBuildingBlock) operand_type = x.type_signature.member zero = building_block_factory.create_generic_constant(operand_type, 0) plus_op = building_block_factory.create_tensorflow_binary_operator_with_upcast( tf.add, computation_types.StructType([operand_type, operand_type])) identity = building_block_factory.create_identity(operand_type) return building_block_factory.create_federated_aggregate( x, zero, plus_op, plus_op, identity)
def federated_aggregate(self, value, zero, accumulate, merge, report): """Implements `federated_aggregate` as defined in `api/intrinsics.py`.""" value = value_impl.to_value(value, None, self._context_stack) value = value_utils.ensure_federated_value(value, placements.CLIENTS, 'value to be aggregated') zero = value_impl.to_value(zero, None, self._context_stack) py_typecheck.check_type(zero, value_base.Value) accumulate = value_impl.to_value(accumulate, None, self._context_stack) merge = value_impl.to_value(merge, None, self._context_stack) report = value_impl.to_value(report, None, self._context_stack) for op in [accumulate, merge, report]: py_typecheck.check_type(op, value_base.Value) py_typecheck.check_type(op.type_signature, computation_types.FunctionType) if not type_utils.is_assignable_from( accumulate.type_signature.parameter[0], zero.type_signature): raise TypeError('Expected `zero` to be assignable to type {}, ' 'but was of incompatible type {}.'.format( accumulate.type_signature.parameter[0], zero.type_signature)) accumulate_type_expected = type_factory.reduction_op( accumulate.type_signature.result, value.type_signature.member) merge_type_expected = type_factory.reduction_op( accumulate.type_signature.result, accumulate.type_signature.result) report_type_expected = computation_types.FunctionType( merge.type_signature.result, report.type_signature.result) for op_name, op, type_expected in [ ('accumulate', accumulate, accumulate_type_expected), ('merge', merge, merge_type_expected), ('report', report, report_type_expected) ]: if not type_utils.is_assignable_from(type_expected, op.type_signature): raise TypeError( 'Expected parameter `{}` to be of type {}, but received {} instead.' .format(op_name, type_expected, op.type_signature)) value = value_impl.ValueImpl.get_comp(value) zero = value_impl.ValueImpl.get_comp(zero) accumulate = value_impl.ValueImpl.get_comp(accumulate) merge = value_impl.ValueImpl.get_comp(merge) report = value_impl.ValueImpl.get_comp(report) comp = building_block_factory.create_federated_aggregate( value, zero, accumulate, merge, report) return value_impl.ValueImpl(comp, self._context_stack)
def test_does_not_find_aggregate_dependent_on_broadcast(self): broadcast = test_utils.create_dummy_called_federated_broadcast() value_type = broadcast.type_signature zero = building_blocks.Data('zero', value_type.member) accumulate_result = building_blocks.Data('accumulate_result', value_type.member) accumulate = building_blocks.Lambda('accumulate_parameter', [value_type.member, value_type.member], accumulate_result) merge_result = building_blocks.Data('merge_result', value_type.member) merge = building_blocks.Lambda('merge_parameter', [value_type.member, value_type.member], merge_result) report_result = building_blocks.Data('report_result', value_type.member) report = building_blocks.Lambda('report_parameter', value_type.member, report_result) aggregate_dependent_on_broadcast = building_block_factory.create_federated_aggregate( broadcast, zero, accumulate, merge, report) tree_analysis.check_broadcast_not_dependent_on_aggregate( aggregate_dependent_on_broadcast)
def create_whimsy_called_federated_aggregate( accumulate_parameter_name='acc_param', merge_parameter_name='merge_param', report_parameter_name='report_param', value_type=tf.int32): r"""Returns a whimsy called federated aggregate. Call / \ federated_aggregate Tuple | [data, data, Lambda(x), Lambda(x), Lambda(x)] | | | data data data Args: accumulate_parameter_name: The name of the accumulate parameter. merge_parameter_name: The name of the merge parameter. report_parameter_name: The name of the report parameter. value_type: The TFF type of the value to be aggregated, placed at CLIENTS. """ federated_value_type = computation_types.FederatedType( value_type, placements.CLIENTS) value = building_blocks.Data('data', federated_value_type) zero = building_blocks.Data('data', value_type) accumulate_type = computation_types.StructType((value_type, value_type)) accumulate_result = building_blocks.Data('data', value_type) accumulate = building_blocks.Lambda(accumulate_parameter_name, accumulate_type, accumulate_result) merge_type = computation_types.StructType((value_type, value_type)) merge_result = building_blocks.Data('data', value_type) merge = building_blocks.Lambda(merge_parameter_name, merge_type, merge_result) report_result = building_blocks.Data('data', value_type) report = building_blocks.Lambda(report_parameter_name, value_type, report_result) return building_block_factory.create_federated_aggregate( value, zero, accumulate, merge, report)
def federated_aggregate(value, zero, accumulate, merge, report) -> value_impl.Value: """Aggregates `value` from `tff.CLIENTS` to `tff.SERVER`. This generalized aggregation function admits multi-layered architectures that involve one or more intermediate stages to handle scalable aggregation across a very large number of participants. The multi-stage aggregation process is defined as follows: * Clients are organized into groups. Within each group, a set of all the member constituents of `value` contributed by clients in the group are first reduced using reduction operator `accumulate` with `zero` as the zero in the algebra. If members of `value` are of type `T`, and `zero` (the result of reducing an empty set) is of type `U`, the reduction operator `accumulate` used at this stage should be of type `(<U,T> -> U)`. The result of this stage is a set of items of type `U`, one item for each group of clients. * Next, the `U`-typed items generated by the preceding stage are merged using the binary commutative associative operator `merge` of type `(<U,U> -> U)`. The result of this stage is a single top-level `U` that emerges at the root of the hierarchy at the `tff.SERVER`. Actual implementations may structure this step as a cascade of multiple layers. * Finally, the `U`-typed result of the reduction performed in the preceding stage is projected into the result value using `report` as the mapping function (for example, if the structures being merged consist of counters, this final step might include computing their ratios). Args: value: A value of a TFF federated type placed at `tff.CLIENTS` to aggregate. zero: The zero of type `U` in the algebra of reduction operators, as described above. accumulate: The reduction operator to use in the first stage of the process. If `value` is of type `{T}@CLIENTS`, and `zero` is of type `U`, this operator should be of type `(<U,T> -> U)`. merge: The reduction operator to employ in the second stage of the process. Must be of type `(<U,U> -> U)`, where `U` is as defined above. report: The projection operator to use at the final stage of the process to compute the final result of aggregation. If the intended result to be returned by `tff.federated_aggregate` is of type `R@SERVER`, this operator must be of type `(U -> R)`. Returns: A representation on the `tff.SERVER` of the result of aggregating `value` using the multi-stage process described above. Raises: TypeError: If the arguments are not of the types specified above. """ value = value_impl.to_value(value, None) value = value_utils.ensure_federated_value(value, placements.CLIENTS, 'value to be aggregated') zero = value_impl.to_value(zero, None) py_typecheck.check_type(zero, value_impl.Value) accumulate = value_impl.to_value( accumulate, None, parameter_type_hint=computation_types.StructType( [zero.type_signature, value.type_signature.member])) merge = value_impl.to_value( merge, None, parameter_type_hint=computation_types.StructType( [accumulate.type_signature.result] * 2)) report = value_impl.to_value( report, None, parameter_type_hint=merge.type_signature.result) for op in [accumulate, merge, report]: py_typecheck.check_type(op, value_impl.Value) py_typecheck.check_type(op.type_signature, computation_types.FunctionType) if not accumulate.type_signature.parameter[0].is_assignable_from( zero.type_signature): raise TypeError('Expected `zero` to be assignable to type {}, ' 'but was of incompatible type {}.'.format( accumulate.type_signature.parameter[0], zero.type_signature)) accumulate_type_expected = type_factory.reduction_op( accumulate.type_signature.result, value.type_signature.member) merge_type_expected = type_factory.reduction_op( accumulate.type_signature.result, accumulate.type_signature.result) report_type_expected = computation_types.FunctionType( merge.type_signature.result, report.type_signature.result) for op_name, op, type_expected in [ ('accumulate', accumulate, accumulate_type_expected), ('merge', merge, merge_type_expected), ('report', report, report_type_expected) ]: if not type_expected.is_assignable_from(op.type_signature): raise TypeError( 'Expected parameter `{}` to be of type {}, but received {} instead.' .format(op_name, type_expected, op.type_signature)) comp = building_block_factory.create_federated_aggregate( value.comp, zero.comp, accumulate.comp, merge.comp, report.comp) comp = _bind_comp_as_reference(comp) return value_impl.Value(comp)