def federated_map_all_equal(fn, arg):
"""`federated_map` with the `all_equal` bit set in the `arg` and return."""
# TODO(b/113112108): Possibly lift the restriction that the mapped value
# must be placed at the clients after adding support for placement labels
# in the federated types, and expanding the type specification of the
# intrinsic this is based on to work with federated values of arbitrary
# placement.
arg = value_impl.to_value(arg, None)
arg = value_utils.ensure_federated_value(arg, placements.CLIENTS,
'value to be mapped')
fn = value_impl.to_value(fn,
None,
parameter_type_hint=arg.type_signature.member)
py_typecheck.check_type(fn, value_impl.Value)
py_typecheck.check_type(fn.type_signature, computation_types.FunctionType)
if not fn.type_signature.parameter.is_assignable_from(
arg.type_signature.member):
raise TypeError(
'The mapping function expects a parameter of type {}, but member '
'constituents of the mapped value are of incompatible type {}.'.
format(fn.type_signature.parameter, arg.type_signature.member))
comp = building_block_factory.create_federated_map_all_equal(
fn.comp, arg.comp)
comp = _bind_comp_as_reference(comp)
return value_impl.Value(comp)
def federated_secure_sum(self, value, bitwidth):
"""Implements `federated_secure_sum` 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 summed')
type_analysis.check_is_structure_of_integers(value.type_signature)
bitwidth_value = value_impl.to_value(bitwidth, None,
self._context_stack)
value_member_type = value.type_signature.member
bitwidth_type = bitwidth_value.type_signature
if not type_analysis.is_valid_bitwidth_type_for_value_type(
bitwidth_type, value_member_type):
raise TypeError(
'Expected `federated_secure_sum` parameter `bitwidth` to match '
'the structure of `value`, with one integer bitwidth per tensor in '
'`value`. Found `value` of `{}` and `bitwidth` of `{}`.'.
format(value_member_type, bitwidth_type))
if bitwidth_type.is_tensor() and value_member_type.is_struct():
bitwidth_value = value_impl.to_value(
structure.map_structure(lambda _: bitwidth, value_member_type),
None, self._context_stack)
value = value_impl.ValueImpl.get_comp(value)
bitwidth_value = value_impl.ValueImpl.get_comp(bitwidth_value)
comp = building_block_factory.create_federated_secure_sum(
value, bitwidth_value)
comp = self._bind_comp_as_reference(comp)
return value_impl.ValueImpl(comp, self._context_stack)
def federated_broadcast(value):
"""Broadcasts a federated value from the `tff.SERVER` to the `tff.CLIENTS`.
Args:
value: A value of a TFF federated type placed at the `tff.SERVER`, all
members of which are equal (the `tff.FederatedType.all_equal` property of
`value` is `True`).
Returns:
A representation of the result of broadcasting: a value of a TFF federated
type placed at the `tff.CLIENTS`, all members of which are equal.
Raises:
TypeError: If the argument is not a federated TFF value placed at the
`tff.SERVER`.
"""
value = value_impl.to_value(value, None)
value = value_utils.ensure_federated_value(value, placements.SERVER,
'value to be broadcasted')
if not value.type_signature.all_equal:
raise TypeError(
'The broadcasted value should be equal at all locations.')
comp = building_block_factory.create_federated_broadcast(value.comp)
comp = _bind_comp_as_reference(comp)
return value_impl.Value(comp)
def federated_mean(self, value, weight):
"""Implements `federated_mean` as defined in `api/intrinsics.py`."""
# TODO(b/113112108): Possibly relax the constraints on numeric types, and
# inject implicit casts where appropriate. For instance, we might want to
# allow `tf.int32` values as the input, and automatically cast them to
# `tf.float321 before invoking the average, thus producing a floating-point
# result.
# TODO(b/120439632): Possibly allow the weight to be either structured or
# non-scalar, e.g., for the case of averaging a convolutional layer, when
# we would want to use a different weight for every filter, and where it
# might be cumbersome for users to have to manually slice and assemble a
# variable.
value = value_impl.to_value(value, None, self._context_stack)
value = value_utils.ensure_federated_value(value,
placement_literals.CLIENTS,
'value to be averaged')
if not type_analysis.is_average_compatible(value.type_signature):
raise TypeError(
'The value type {} is not compatible with the average operator.'
.format(value.type_signature))
if weight is not None:
weight = value_impl.to_value(weight, None, self._context_stack)
weight = value_utils.ensure_federated_value(
weight, placement_literals.CLIENTS,
'weight to use in averaging')
py_typecheck.check_type(weight.type_signature.member,
computation_types.TensorType)
if weight.type_signature.member.shape.ndims != 0:
raise TypeError(
'The weight type {} is not a federated scalar.'.format(
weight.type_signature))
if not (weight.type_signature.member.dtype.is_integer
or weight.type_signature.member.dtype.is_floating):
raise TypeError(
'The weight type {} is not a federated integer or floating-point '
'tensor.'.format(weight.type_signature))
value = value_impl.ValueImpl.get_comp(value)
if weight is not None:
weight = value_impl.ValueImpl.get_comp(weight)
comp = building_block_factory.create_federated_mean(value, weight)
comp = self._bind_comp_as_reference(comp)
return value_impl.ValueImpl(comp, self._context_stack)
def federated_sum(self, value):
"""Implements `federated_sum` 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 summed')
type_analysis.check_is_sum_compatible(value.type_signature)
value = value_impl.ValueImpl.get_comp(value)
comp = building_block_factory.create_federated_sum(value)
comp = self._bind_comp_as_reference(comp)
return value_impl.ValueImpl(comp, self._context_stack)
def federated_collect(self, value):
"""Implements `federated_collect` 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 collected')
value = value_impl.ValueImpl.get_comp(value)
comp = building_block_factory.create_federated_collect(value)
comp = self._bind_comp_as_reference(comp)
return value_impl.ValueImpl(comp, self._context_stack)
def federated_map(self, fn, arg):
"""Implements `federated_map` as defined in `api/intrinsics.py`."""
# TODO(b/113112108): Possibly lift the restriction that the mapped value
# must be placed at the server or clients. Would occur after adding support
# for placement labels in the federated types, and expanding the type
# specification of the intrinsic this is based on to work with federated
# values of arbitrary placement.
arg = value_impl.to_value(arg, None, self._context_stack)
arg = value_utils.ensure_federated_value(arg,
label='value to be mapped')
fn = value_impl.to_value(fn,
None,
self._context_stack,
parameter_type_hint=arg.type_signature.member)
py_typecheck.check_type(fn, value_base.Value)
py_typecheck.check_type(fn.type_signature,
computation_types.FunctionType)
if not fn.type_signature.parameter.is_assignable_from(
arg.type_signature.member):
raise TypeError(
'The mapping function expects a parameter of type {}, but member '
'constituents of the mapped value are of incompatible type {}.'
.format(fn.type_signature.parameter,
arg.type_signature.member))
# TODO(b/144384398): Change structure to one that maps the placement type
# to the building_block function that fits it, in a way that allows the
# appropriate type checks.
if arg.type_signature.placement is placement_literals.SERVER:
if not arg.type_signature.all_equal:
raise TypeError(
'Arguments placed at {} should be equal at all locations.'.
format(placement_literals.SERVER))
fn = value_impl.ValueImpl.get_comp(fn)
arg = value_impl.ValueImpl.get_comp(arg)
comp = building_block_factory.create_federated_apply(fn, arg)
elif arg.type_signature.placement is placement_literals.CLIENTS:
fn = value_impl.ValueImpl.get_comp(fn)
arg = value_impl.ValueImpl.get_comp(arg)
comp = building_block_factory.create_federated_map(fn, arg)
else:
raise TypeError(
'Expected `arg` to have a type with a supported placement, '
'found {}.'.format(arg.type_signature.placement))
comp = self._bind_comp_as_reference(comp)
return value_impl.ValueImpl(comp, self._context_stack)
def federated_broadcast(self, value):
"""Implements `federated_broadcast` 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.SERVER,
'value to be broadcasted')
if not value.type_signature.all_equal:
raise TypeError(
'The broadcasted value should be equal at all locations.')
value = value_impl.ValueImpl.get_comp(value)
comp = building_block_factory.create_federated_broadcast(value)
comp = self._bind_comp_as_reference(comp)
return value_impl.ValueImpl(comp, self._context_stack)
def _federated_select(client_keys, max_key, server_val, select_fn, secure):
"""Internal helper for `federated_select` and `federated_secure_select`."""
client_keys = value_impl.to_value(client_keys, None)
_check_select_keys_type(client_keys.type_signature, secure)
max_key = value_impl.to_value(max_key, None)
expected_max_key_type = computation_types.at_server(tf.int32)
if not expected_max_key_type.is_assignable_from(max_key.type_signature):
_select_parameter_mismatch(
max_key.type_signature,
'a 32-bit unsigned integer placed at server',
'max_key',
secure,
expected_type=expected_max_key_type)
server_val = value_impl.to_value(server_val, None)
server_val = value_utils.ensure_federated_value(server_val,
label='server_val')
expected_server_val_type = computation_types.at_server(
computation_types.AbstractType('T'))
if (not server_val.type_signature.is_federated()
or not server_val.type_signature.placement.is_server()):
_select_parameter_mismatch(server_val.type_signature,
'a value placed at server',
'server_val',
secure,
expected_type=expected_server_val_type)
select_fn_param_type = computation_types.to_type(
[server_val.type_signature.member, tf.int32])
select_fn = value_impl.to_value(select_fn,
None,
parameter_type_hint=select_fn_param_type)
expected_select_fn_type = computation_types.FunctionType(
select_fn_param_type, computation_types.AbstractType('U'))
if (not select_fn.type_signature.is_function()
or not select_fn.type_signature.parameter.is_assignable_from(
select_fn_param_type)):
_select_parameter_mismatch(select_fn.type_signature,
'a function from state and key to result',
'select_fn',
secure,
expected_type=expected_select_fn_type)
comp = building_block_factory.create_federated_select(
client_keys.comp, max_key.comp, server_val.comp, select_fn.comp,
secure)
comp = _bind_comp_as_reference(comp)
return value_impl.Value(comp)
def federated_collect(value):
"""Returns a federated value from `tff.CLIENTS` as a `tff.SERVER` sequence.
Args:
value: A value of a TFF federated type placed at the `tff.CLIENTS`.
Returns:
A stream of the same type as the member constituents of `value` placed at
the `tff.SERVER`.
Raises:
TypeError: If the argument is not a federated TFF value placed at
`tff.CLIENTS`.
"""
value = value_impl.to_value(value, None)
value = value_utils.ensure_federated_value(value, placements.CLIENTS,
'value to be collected')
comp = building_block_factory.create_federated_collect(value.comp)
comp = _bind_comp_as_reference(comp)
return value_impl.Value(comp)
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_sum(value):
"""Computes a sum at `tff.SERVER` of a `value` placed on the `tff.CLIENTS`.
To sum integer values with stronger privacy properties, consider using
`tff.federated_secure_sum_bitwidth`.
Args:
value: A value of a TFF federated type placed at the `tff.CLIENTS`.
Returns:
A representation of the sum of the member constituents of `value` placed
on the `tff.SERVER`.
Raises:
TypeError: If the argument is not a federated TFF value placed at
`tff.CLIENTS`.
"""
value = value_impl.to_value(value, None)
value = value_utils.ensure_federated_value(value, placements.CLIENTS,
'value to be summed')
type_analysis.check_is_sum_compatible(value.type_signature)
comp = building_block_factory.create_federated_sum(value.comp)
comp = _bind_comp_as_reference(comp)
return value_impl.Value(comp)
def _(x):
x = value_impl.to_value(x, None, _context_stack)
value_utils.ensure_federated_value(x, placements.CLIENTS)
return x
def federated_secure_sum_bitwidth(value, bitwidth):
"""Computes a sum at `tff.SERVER` of a `value` placed on the `tff.CLIENTS`.
This function computes a sum such that it should not be possible for the
server to learn any clients individual value. The specific algorithm and
mechanism used to compute the secure sum may vary depending on the target
runtime environment the computation is compiled for or executed on. See
https://research.google/pubs/pub47246/ for more information.
Not all executors support `tff.federated_secure_sum_bitwidth()`; consult the
documentation for the specific executor or executor stack you plan on using
for the specific of how it's handled by that executor.
The `bitwidth` argument represents the bitwidth of the aggregand, that is the
bitwidth of the input `value`. The federated secure sum bitwidth (i.e., the
bitwidth of the *sum* of the input `value`s over all clients) will be a
function of this bitwidth and the number of participating clients.
Example:
```python
value = tff.federated_value(1, tff.CLIENTS)
result = tff.federated_secure_sum_bitwidth(value, 2)
value = tff.federated_value([1, 1], tff.CLIENTS)
result = tff.federated_secure_sum_bitwidth(value, [2, 4])
value = tff.federated_value([1, [1, 1]], tff.CLIENTS)
result = tff.federated_secure_sum_bitwidth(value, [2, [4, 8]])
```
Note: To sum non-integer values or to sum integers with fewer constraints and
weaker privacy properties, consider using `federated_sum`.
Args:
value: An integer value of a TFF federated type placed at the `tff.CLIENTS`,
in the range [0, 2^bitwidth - 1].
bitwidth: An integer or nested structure of integers matching the structure
of `value`. If integer `bitwidth` is used with a nested `value`, the same
integer is used for each tensor in `value`.
Returns:
A representation of the sum of the member constituents of `value` placed
on the `tff.SERVER`.
Raises:
TypeError: If the argument is not a federated TFF value placed at
`tff.CLIENTS`.
"""
value = value_impl.to_value(value, None)
value = value_utils.ensure_federated_value(value, placements.CLIENTS,
'value to be summed')
type_analysis.check_is_structure_of_integers(value.type_signature)
bitwidth_value = value_impl.to_value(bitwidth, None)
value_member_type = value.type_signature.member
bitwidth_type = bitwidth_value.type_signature
if not type_analysis.is_single_integer_or_matches_structure(
bitwidth_type, value_member_type):
raise TypeError(
'Expected `federated_secure_sum_bitwidth` parameter `bitwidth` to match '
'the structure of `value`, with one integer bitwidth per tensor in '
'`value`. Found `value` of `{}` and `bitwidth` of `{}`.'.format(
value_member_type, bitwidth_type))
if bitwidth_type.is_tensor() and value_member_type.is_struct():
bitwidth_value = value_impl.to_value(
structure.map_structure(lambda _: bitwidth, value_member_type),
None)
comp = building_block_factory.create_federated_secure_sum_bitwidth(
value.comp, bitwidth_value.comp)
comp = _bind_comp_as_reference(comp)
return value_impl.Value(comp)
def federated_secure_sum(value, max_input):
"""Computes a sum at `tff.SERVER` of a `value` placed on the `tff.CLIENTS`.
This function computes a sum such that it should not be possible for the
server to learn any clients individual value. The specific algorithm and
mechanism used to compute the secure sum may vary depending on the target
runtime environment the computation is compiled for or executed on. See
https://research.google/pubs/pub47246/ for more information.
Not all executors support `tff.federated_secure_sum()`; consult the
documentation for the specific executor or executor stack you plan on using
for the specific of how it's handled by that executor.
The `max_input` argument is the maximum value (inclusive) that may appear in
`value`. *Lower values may allow for improved communication efficiency.*
Attempting to return a `value` higher than `max_input` is invalid, and will
result in a failure at the given client.
Example:
```python
value = tff.federated_value(1, tff.CLIENTS)
result = tff.federated_secure_sum(value, 1)
value = tff.federated_value((1, 2), tff.CLIENTS)
result = tff.federated_secure_sum(value, (1, 2))
```
Note: To sum non-integer values or to sum integers with fewer constraints and
weaker privacy properties, consider using `federated_sum`.
Args:
value: An integer or nested structure of integers placed at `tff.CLIENTS`,
in the range `[0, max_input]`.
max_input: A Python integer or nested structure of integers matching the
structure of `value`. If integer `max_value` is used with a nested
`value`, the same integer is used for each tensor in `value`.
Returns:
A representation of the sum of the member constituents of `value` placed
on the `tff.SERVER`.
Raises:
TypeError: If the argument is not a federated TFF value placed at
`tff.CLIENTS`.
"""
value = value_impl.to_value(value, None)
value = value_utils.ensure_federated_value(value, placements.CLIENTS,
'value to be summed')
type_analysis.check_is_structure_of_integers(value.type_signature)
max_input_value = value_impl.to_value(max_input, None)
value_member_type = value.type_signature.member
max_input_type = max_input_value.type_signature
if not type_analysis.is_single_integer_or_matches_structure(
max_input_type, value_member_type):
raise TypeError(
'Expected `federated_secure_sum` parameter `max_input` to match '
'the structure of `value`, with one integer max per tensor in '
'`value`. Found `value` of `{}` and `max_input` of `{}`.'.format(
value_member_type, max_input_type))
if max_input_type.is_tensor() and value_member_type.is_struct():
max_input_value = value_impl.to_value(
structure.map_structure(lambda _: max_input, value_member_type),
None)
comp = building_block_factory.create_federated_secure_sum(
value.comp, max_input_value.comp)
comp = _bind_comp_as_reference(comp)
return value_impl.Value(comp)
def federated_secure_modular_sum(value, modulus):
"""Computes a modular sum at `tff.SERVER` of a `value` from `tff.CLIENTS`.
This function computes a sum such that it should not be possible for the
server to learn any clients individual value. The specific algorithm and
mechanism used to compute the secure sum may vary depending on the target
runtime environment the computation is compiled for or executed on. See
https://research.google/pubs/pub47246/ for more information.
Not all executors support `tff.federated_secure_modular_sum()`; consult the
documentation for the specific executor or executor stack you plan on using
for the specific of how it's handled by that executor.
The `modulus` argument is the modulus under which the client values are added.
The result of this function will be equivalent to `SUM(value) % modulus`.
*Lower values may allow for improved communication efficiency.*
Example:
```python
value = tff.federated_value(5, tff.CLIENTS)
result = tff.federated_secure_modular_sum(value, 3)
# `result == (5 * num_clients % 3)@SERVER`
value = tff.federated_value((3, 9), tff.CLIENTS)
result = tff.federated_secure_modular_sum(value, (100, 200))
# `result == (3 * num_clients % 100, 9 * num_clients % 100)@SERVER`
```
Note: To sum non-integer values or to sum integers with fewer constraints and
weaker privacy properties, consider using `federated_sum`.
Args:
value: An integer or nested structure of integers placed at `tff.CLIENTS`.
Values outside of the range [0, modulus-1] will be considered equivalent
to mod(value, modulus), i.e. they will be projected into the range
[0, modulus-1] as part of the modular summation.
modulus: A Python integer or nested structure of integers matching the
structure of `value`. If integer `modulus` is used with a nested `value`,
the same integer is used for each tensor in `value`.
Returns:
A representation of the modular sum of the member constituents of `value`
placed on the `tff.SERVER`. The resulting modular sum will be on the range
[0, modulus-1].
Raises:
TypeError: If the argument is not a federated TFF value placed at
`tff.CLIENTS`.
"""
value = value_impl.to_value(value, None)
value = value_utils.ensure_federated_value(value, placements.CLIENTS,
'value to be summed')
type_analysis.check_is_structure_of_integers(value.type_signature)
modulus_value = value_impl.to_value(modulus, None)
value_member_type = value.type_signature.member
modulus_type = modulus_value.type_signature
if not type_analysis.is_single_integer_or_matches_structure(
modulus_type, value_member_type):
raise TypeError(
'Expected `federated_secure_sum` parameter `modulus` to match '
'the structure of `value`, with one integer max per tensor in '
'`value`. Found `value` of `{}` and `modulus` of `{}`.'.format(
value_member_type, modulus_type))
if modulus_type.is_tensor() and value_member_type.is_struct():
modulus_value = value_impl.to_value(
structure.map_structure(lambda _: modulus, value_member_type),
None)
comp = building_block_factory.create_federated_secure_modular_sum(
value.comp, modulus_value.comp)
comp = _bind_comp_as_reference(comp)
return value_impl.Value(comp)
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 federated_mean(value, weight=None):
"""Computes a `tff.SERVER` mean of `value` placed on `tff.CLIENTS`.
For values `v_1, ..., v_k`, and weights `w_1, ..., w_k`, this means
`sum_{i=1}^k (w_i * v_i) / sum_{i=1}^k w_i`.
Args:
value: The value of which the mean is to be computed. Must be of a TFF
federated type placed at `tff.CLIENTS`. The value may be structured, e.g.,
its member constituents can be named tuples. The tensor types that the
value is composed of must be floating-point or complex.
weight: An optional weight, a TFF federated integer or floating-point tensor
value, also placed at `tff.CLIENTS`.
Returns:
A representation at the `tff.SERVER` of the mean of the member constituents
of `value`, optionally weighted with `weight` if specified (otherwise, the
member constituents contributed by all clients are equally weighted).
Raises:
TypeError: If `value` is not a federated TFF value placed at `tff.CLIENTS`,
or if `weight` is not a federated integer or a floating-point tensor with
the matching placement.
"""
# TODO(b/113112108): Possibly relax the constraints on numeric types, and
# inject implicit casts where appropriate. For instance, we might want to
# allow `tf.int32` values as the input, and automatically cast them to
# `tf.float321 before invoking the average, thus producing a floating-point
# result.
# TODO(b/120439632): Possibly allow the weight to be either structured or
# non-scalar, e.g., for the case of averaging a convolutional layer, when
# we would want to use a different weight for every filter, and where it
# might be cumbersome for users to have to manually slice and assemble a
# variable.
value = value_impl.to_value(value, None)
value = value_utils.ensure_federated_value(value, placements.CLIENTS,
'value to be averaged')
if not type_analysis.is_average_compatible(value.type_signature):
raise TypeError(
'The value type {} is not compatible with the average operator.'.
format(value.type_signature))
if weight is not None:
weight = value_impl.to_value(weight, None)
weight = value_utils.ensure_federated_value(
weight, placements.CLIENTS, 'weight to use in averaging')
py_typecheck.check_type(weight.type_signature.member,
computation_types.TensorType)
if weight.type_signature.member.shape.ndims != 0:
raise TypeError(
'The weight type {} is not a federated scalar.'.format(
weight.type_signature))
if not (weight.type_signature.member.dtype.is_integer
or weight.type_signature.member.dtype.is_floating):
raise TypeError(
'The weight type {} is not a federated integer or floating-point '
'tensor.'.format(weight.type_signature))
weight_comp = None if weight is None else weight.comp
comp = building_block_factory.create_federated_mean(
value.comp, weight_comp)
comp = _bind_comp_as_reference(comp)
return value_impl.Value(comp)
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.StructType(
[zero.type_signature, value.type_signature.member]))
merge = value_impl.to_value(
merge,
None,
self._context_stack,
parameter_type_hint=computation_types.StructType(
[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)
def _(x):
x = value_impl.to_value(x, None, _context_stack)
value_utils.ensure_federated_value(x)
return x
def _(x):
x = value_impl.to_value(x, None)
value_utils.ensure_federated_value(x)
return x
def _(x):
x = value_impl.to_value(x, None)
with self.assertRaises(TypeError):
value_utils.ensure_federated_value(x, placements.SERVER)
return x
def _(x):
x = value_impl.to_value(x, None, _context_stack)
with self.assertRaises(TypeError):
value_utils.ensure_federated_value(x)
return x
def federated_map(fn, arg):
"""Maps a federated value pointwise using a mapping function.
The function `fn` is applied separately across the group of devices
represented by the placement type of `arg`. For example, if `value` has
placement type `tff.CLIENTS`, then `fn` is applied to each client
individually. In particular, this operation does not alter the placement of
the federated value.
Args:
fn: A mapping function to apply pointwise to member constituents of `arg`.
The parameter of this function must be of the same type as the member
constituents of `arg`.
arg: A value of a TFF federated type (or a value that can be implicitly
converted into a TFF federated type, e.g., by zipping) placed at
`tff.CLIENTS` or `tff.SERVER`.
Returns:
A federated value with the same placement as `arg` that represents the
result of `fn` on the member constituent of `arg`.
Raises:
TypeError: If the arguments are not of the appropriate types.
"""
# TODO(b/113112108): Possibly lift the restriction that the mapped value
# must be placed at the server or clients. Would occur after adding support
# for placement labels in the federated types, and expanding the type
# specification of the intrinsic this is based on to work with federated
# values of arbitrary placement.
arg = value_impl.to_value(arg, None)
arg = value_utils.ensure_federated_value(arg, label='value to be mapped')
fn = value_impl.to_value(fn,
None,
parameter_type_hint=arg.type_signature.member)
py_typecheck.check_type(fn, value_impl.Value)
py_typecheck.check_type(fn.type_signature, computation_types.FunctionType)
if not fn.type_signature.parameter.is_assignable_from(
arg.type_signature.member):
raise TypeError(
'The mapping function expects a parameter of type {}, but member '
'constituents of the mapped value are of incompatible type {}.'.
format(fn.type_signature.parameter, arg.type_signature.member))
# TODO(b/144384398): Change structure to one that maps the placement type
# to the building_block function that fits it, in a way that allows the
# appropriate type checks.
if arg.type_signature.placement is placements.SERVER:
if not arg.type_signature.all_equal:
raise TypeError(
'Arguments placed at {} should be equal at all locations.'.
format(placements.SERVER))
comp = building_block_factory.create_federated_apply(fn.comp, arg.comp)
elif arg.type_signature.placement is placements.CLIENTS:
comp = building_block_factory.create_federated_map(fn.comp, arg.comp)
else:
raise TypeError(
'Expected `arg` to have a type with a supported placement, '
'found {}.'.format(arg.type_signature.placement))
comp = _bind_comp_as_reference(comp)
return value_impl.Value(comp)
