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)
