def _create_complex_computation():
tensor_type = computation_types.TensorType(tf.int32)
compiled = building_block_factory.create_compiled_identity(
tensor_type, 'a')
federated_type = computation_types.FederatedType(tf.int32,
placements.SERVER)
arg_ref = building_blocks.Reference('arg', federated_type)
bindings = []
results = []
def _bind(name, value):
bindings.append((name, value))
return building_blocks.Reference(name, value.type_signature)
for i in range(2):
called_federated_broadcast = building_block_factory.create_federated_broadcast(
arg_ref)
called_federated_map = building_block_factory.create_federated_map(
compiled, _bind(f'broadcast_{i}', called_federated_broadcast))
called_federated_mean = building_block_factory.create_federated_mean(
_bind(f'map_{i}', called_federated_map), None)
results.append(_bind(f'mean_{i}', called_federated_mean))
result = building_blocks.Struct(results)
block = building_blocks.Block(bindings, result)
return building_blocks.Lambda('arg', tf.int32, block)
def create_whimsy_called_federated_mean(value_type=tf.float32,
weights_type=None):
fed_value_type = computation_types.at_clients(value_type)
values = building_blocks.Data('values', fed_value_type)
if weights_type is not None:
fed_weights_type = computation_types.at_clients(weights_type)
weights = building_blocks.Data('weights', fed_weights_type)
else:
weights = None
return building_block_factory.create_federated_mean(values, weights)
def federated_mean(self, value, weight):
"""Implements `federated_mean` as defined in `api/intrinsics.py`.
Args:
value: As in `api/intrinsics.py`.
weight: As in `api/intrinsics.py`.
Returns:
As in `api/intrinsics.py`.
Raises:
TypeError: As 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_utils.check_federated_value_placement(value, placements.CLIENTS,
'value to be averaged')
if not type_utils.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)
value_utils.check_federated_value_placement(
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))
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)
return value_impl.ValueImpl(comp, self._context_stack)
def _create_complex_computation():
compiled = building_block_factory.create_compiled_identity(tf.int32, 'a')
federated_type = computation_types.FederatedType(tf.int32,
placements.SERVER)
ref = building_blocks.Reference('b', federated_type)
called_federated_broadcast = building_block_factory.create_federated_broadcast(
ref)
called_federated_map = building_block_factory.create_federated_map(
compiled, called_federated_broadcast)
called_federated_mean = building_block_factory.create_federated_mean(
called_federated_map, None)
tup = building_blocks.Tuple([called_federated_mean, called_federated_mean])
return building_blocks.Lambda('b', tf.int32, tup)
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)
