def sequence_reduce(self, value, zero, op):
"""Implements `sequence_reduce` as defined in `api/intrinsics.py`."""
value = value_impl.to_value(value, None, self._context_stack)
zero = value_impl.to_value(zero, None, self._context_stack)
op = value_impl.to_value(op, None, self._context_stack)
# Check if the value is a federated sequence that should be reduced
# under a `federated_map`.
if value.type_signature.is_federated():
value_sequence_type = value.type_signature.member
needs_map = True
else:
value_sequence_type = value.type_signature
needs_map = False
value_sequence_type.check_sequence()
element_type = value_sequence_type.element
op_type_expected = type_factory.reduction_op(zero.type_signature,
element_type)
if not op_type_expected.is_assignable_from(op.type_signature):
raise TypeError('Expected an operator of type {}, got {}.'.format(
op_type_expected, op.type_signature))
value = value_impl.ValueImpl.get_comp(value)
zero = value_impl.ValueImpl.get_comp(zero)
op = value_impl.ValueImpl.get_comp(op)
if not needs_map:
comp = building_block_factory.create_sequence_reduce(
value, zero, op)
comp = self._bind_comp_as_reference(comp)
return value_impl.ValueImpl(comp, self._context_stack)
else:
ref = building_blocks.Reference('arg', value_sequence_type)
call = building_block_factory.create_sequence_reduce(ref, zero, op)
fn = building_blocks.Lambda(ref.name, ref.type_signature, call)
fn_impl = value_impl.ValueImpl(fn, self._context_stack)
return self.federated_map(fn_impl, value)
def parse_federated_aggregate_argument_types(type_spec):
"""Verifies and parses `type_spec` into constituents.
Args:
type_spec: An instance of `computation_types.StructType`.
Returns:
A tuple of (value_type, zero_type, accumulate_type, merge_type, report_type)
for the 5 type constituents.
"""
py_typecheck.check_type(type_spec, computation_types.StructType)
py_typecheck.check_len(type_spec, 5)
value_type = type_spec[0]
py_typecheck.check_type(value_type, computation_types.FederatedType)
item_type = value_type.member
zero_type = type_spec[1]
accumulate_type = type_spec[2]
accumulate_type.check_equivalent_to(
type_factory.reduction_op(zero_type, item_type))
merge_type = type_spec[3]
merge_type.check_equivalent_to(type_factory.binary_op(zero_type))
report_type = type_spec[4]
py_typecheck.check_type(report_type, computation_types.FunctionType)
report_type.parameter.check_equivalent_to(zero_type)
return value_type, zero_type, accumulate_type, merge_type, report_type
def test_reduction_op(self):
result_type = computation_types.TensorType(tf.float32)
element_type = computation_types.TensorType(tf.int32)
actual_type = type_factory.reduction_op(result_type, element_type)
expected_type = computation_types.FunctionType(
computation_types.StructType([result_type, element_type]), result_type)
self.assertEqual(actual_type, expected_type)
def create_dummy_intrinsic_def_federated_reduce():
value = intrinsic_defs.FEDERATED_REDUCE
type_signature = computation_types.FunctionType([
computation_types.at_clients(tf.float32),
tf.float32,
type_factory.reduction_op(tf.float32, tf.float32),
], computation_types.at_server(tf.float32))
return value, type_signature
def create_dummy_intrinsic_def_federated_aggregate():
value = intrinsic_defs.FEDERATED_AGGREGATE
type_signature = computation_types.FunctionType([
type_factory.at_clients(tf.float32),
tf.float32,
type_factory.reduction_op(tf.float32, tf.float32),
type_factory.binary_op(tf.float32),
computation_types.FunctionType(tf.float32, tf.float32),
], type_factory.at_server(tf.float32))
return value, type_signature
def sequence_reduce(self, value, zero, op):
"""Implements `sequence_reduce` as defined in `api/intrinsics.py`."""
value = value_impl.to_value(value, None, self._context_stack)
zero = value_impl.to_value(zero, None, self._context_stack)
op = value_impl.to_value(op, None, self._context_stack)
if isinstance(value.type_signature, computation_types.SequenceType):
element_type = value.type_signature.element
else:
py_typecheck.check_type(value.type_signature,
computation_types.FederatedType)
py_typecheck.check_type(value.type_signature.member,
computation_types.SequenceType)
element_type = value.type_signature.member.element
op_type_expected = type_factory.reduction_op(zero.type_signature,
element_type)
if not type_analysis.is_assignable_from(op_type_expected,
op.type_signature):
raise TypeError('Expected an operator of type {}, got {}.'.format(
op_type_expected, op.type_signature))
value = value_impl.ValueImpl.get_comp(value)
zero = value_impl.ValueImpl.get_comp(zero)
op = value_impl.ValueImpl.get_comp(op)
if isinstance(value.type_signature, computation_types.SequenceType):
return value_impl.ValueImpl(
building_block_factory.create_sequence_reduce(value, zero, op),
self._context_stack)
else:
value_type = computation_types.SequenceType(element_type)
intrinsic_type = computation_types.FunctionType((
value_type,
zero.type_signature,
op.type_signature,
), op.type_signature.result)
intrinsic = building_blocks.Intrinsic(
intrinsic_defs.SEQUENCE_REDUCE.uri, intrinsic_type)
ref = building_blocks.Reference('arg', value_type)
tup = building_blocks.Tuple((ref, zero, op))
call = building_blocks.Call(intrinsic, tup)
fn = building_blocks.Lambda(ref.name, ref.type_signature, call)
fn_impl = value_impl.ValueImpl(fn, self._context_stack)
if value.type_signature.placement in [
placement_literals.SERVER, placement_literals.CLIENTS
]:
return self.federated_map(fn_impl, value)
else:
raise TypeError('Unsupported placement {}.'.format(
value.type_signature.placement))
async def compute_federated_reduce(
self, arg: FederatedResolvingStrategyValue
) -> FederatedResolvingStrategyValue:
self._check_arg_is_structure(arg)
if len(arg.internal_representation) != 3:
raise ValueError(
'Expected 3 elements in the `federated_reduce()` argument tuple, '
'found {}.'.format(len(arg.internal_representation)))
val_type = arg.type_signature[0]
py_typecheck.check_type(val_type, computation_types.FederatedType)
item_type = val_type.member
zero_type = arg.type_signature[1]
op_type = arg.type_signature[2]
op_type.check_equivalent_to(
type_factory.reduction_op(zero_type, item_type))
val = arg.internal_representation[0]
py_typecheck.check_type(val, list)
child = self._target_executors[placement_literals.SERVER][0]
async def _move(v):
return await child.create_value(await v.compute(), item_type)
item_futures = asyncio.as_completed([_move(v) for v in val])
zero = await child.create_value(
await (await self._executor.create_selection(arg,
index=1)).compute(),
zero_type)
op = await child.create_value(arg.internal_representation[2], op_type)
result = zero
for item_future in item_futures:
item = await item_future
result = await child.create_call(
op, await child.create_struct(
structure.Struct([(None, result), (None, item)])))
return FederatedResolvingStrategyValue([result],
computation_types.FederatedType(
result.type_signature,
placement_literals.SERVER,
all_equal=True))
async def _compute_intrinsic_federated_reduce(self, arg):
self._check_arg_is_anonymous_tuple(arg)
if len(arg.internal_representation) != 3:
raise ValueError(
'Expected 3 elements in the `federated_reduce()` argument tuple, '
'found {}.'.format(len(arg.internal_representation)))
val_type = arg.type_signature[0]
py_typecheck.check_type(val_type, computation_types.FederatedType)
item_type = val_type.member
zero_type = arg.type_signature[1]
op_type = arg.type_signature[2]
type_analysis.check_equivalent_types(
op_type, type_factory.reduction_op(zero_type, item_type))
val = arg.internal_representation[0]
py_typecheck.check_type(val, list)
child = self._target_executors[placement_literals.SERVER][0]
async def _move(v):
return await child.create_value(await v.compute(), item_type)
items = await asyncio.gather(*[_move(v) for v in val])
zero = await child.create_value(
await (await self.create_selection(arg, index=1)).compute(),
zero_type)
op = await child.create_value(arg.internal_representation[2], op_type)
result = zero
for item in items:
result = await child.create_call(
op, await child.create_tuple(
anonymous_tuple.AnonymousTuple([(None, result),
(None, item)])))
return FederatingExecutorValue([result],
computation_types.FederatedType(
result.type_signature,
placement_literals.SERVER,
all_equal=True))
def federated_reduce(self, value, zero, op):
"""Implements `federated_reduce` as defined in `api/intrinsics.py`."""
value = value_impl.to_value(value, None, self._context_stack)
value = value_utils.ensure_federated_value(value,
placement_literals.CLIENTS,
'value to be reduced')
zero = value_impl.to_value(zero, None, self._context_stack)
if type_analysis.contains_federated_types(zero.type_signature):
raise TypeError(
'`zero` may not contain a federated type, found type:\n' +
str(zero.type_signature))
op = value_impl.to_value(
op,
None,
self._context_stack,
parameter_type_hint=computation_types.StructType(
[zero.type_signature, value.type_signature.member]))
op.type_signature.check_function()
if not op.type_signature.result.is_assignable_from(
zero.type_signature):
raise TypeError(
'`zero` must be assignable to the result type from `op`:\n',
computation_types.type_mismatch_error_message(
zero.type_signature, op.type_signature.result,
computation_types.TypeRelation.ASSIGNABLE))
op_type_expected = type_factory.reduction_op(
op.type_signature.result, value.type_signature.member)
if not op_type_expected.is_assignable_from(op.type_signature):
raise TypeError('Expected an operator of type {}, got {}.'.format(
op_type_expected, op.type_signature))
value = value_impl.ValueImpl.get_comp(value)
zero = value_impl.ValueImpl.get_comp(zero)
op = value_impl.ValueImpl.get_comp(op)
comp = building_block_factory.create_federated_reduce(value, zero, op)
comp = self._bind_comp_as_reference(comp)
return value_impl.ValueImpl(comp, self._context_stack)
def federated_reduce(self, value, zero, op):
"""Implements `federated_reduce` as defined in `api/intrinsics.py`."""
# TODO(b/113112108): Since in most cases, it can be assumed that CLIENTS is
# a non-empty collective (or else, the computation fails), specifying zero
# at this level of the API should probably be optional. TBD.
value = value_impl.to_value(value, None, self._context_stack)
value = value_utils.ensure_federated_value(value,
placement_literals.CLIENTS,
'value to be reduced')
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.
op = value_impl.to_value(
op,
None,
self._context_stack,
parameter_type_hint=computation_types.NamedTupleType(
[zero.type_signature, value.type_signature.member]))
py_typecheck.check_type(op, value_base.Value)
py_typecheck.check_type(op.type_signature,
computation_types.FunctionType)
op_type_expected = type_factory.reduction_op(
zero.type_signature, value.type_signature.member)
if not type_analysis.is_assignable_from(op_type_expected,
op.type_signature):
raise TypeError('Expected an operator of type {}, got {}.'.format(
op_type_expected, op.type_signature))
value = value_impl.ValueImpl.get_comp(value)
zero = value_impl.ValueImpl.get_comp(zero)
op = value_impl.ValueImpl.get_comp(op)
comp = building_block_factory.create_federated_reduce(value, zero, op)
return value_impl.ValueImpl(comp, self._context_stack)
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,
placement_literals.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,
parameter_type_hint=computation_types.NamedTupleType(
[zero.type_signature, value.type_signature.member]))
merge = value_impl.to_value(
merge,
None,
self._context_stack,
parameter_type_hint=computation_types.NamedTupleType(
[zero.type_signature, zero.type_signature]))
report = value_impl.to_value(
report,
None,
self._context_stack,
parameter_type_hint=zero.type_signature)
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 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))
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)
comp = self._bind_comp_as_reference(comp)
return value_impl.ValueImpl(comp, self._context_stack)
# @federated_computation
# def federated_aggregate(x, zero, accumulate, merge, report):
# a = generic_partial_reduce(x, zero, accumulate, INTERMEDIATE_AGGREGATORS)
# b = generic_reduce(a, zero, merge, SERVER)
# c = generic_map(report, b)
# return c
#
# Actual implementations might vary.
#
# Type signature: <{T}@CLIENTS,U,(<U,T>->U),(<U,U>->U),(U->R)> -> R@SERVER
FEDERATED_AGGREGATE = IntrinsicDef(
'FEDERATED_AGGREGATE', 'federated_aggregate',
computation_types.FunctionType(parameter=[
type_factory.at_clients(computation_types.AbstractType('T')),
computation_types.AbstractType('U'),
type_factory.reduction_op(computation_types.AbstractType('U'),
computation_types.AbstractType('T')),
type_factory.binary_op(computation_types.AbstractType('U')),
computation_types.FunctionType(computation_types.AbstractType('U'),
computation_types.AbstractType('R'))
],
result=type_factory.at_server(
computation_types.AbstractType('R'))))
# Applies a given function to a value on the server.
#
# Type signature: <(T->U),T@SERVER> -> U@SERVER
FEDERATED_APPLY = IntrinsicDef(
'FEDERATED_APPLY', 'federated_apply',
computation_types.FunctionType(parameter=[
computation_types.FunctionType(computation_types.AbstractType('T'),
computation_types.AbstractType('U')),
def test_reduction_op(self):
self.assertEqual(
str(type_factory.reduction_op(tf.float32, tf.int32)),
'(<float32,int32> -> float32)')
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)
def _make_sequence_reduce_type(element_type, accumulator_type):
return computation_types.FunctionType(parameter=[
computation_types.SequenceType(element_type), accumulator_type,
type_factory.reduction_op(accumulator_type, element_type)
],
result=accumulator_type)
