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 type_utils.is_assignable_from(fn.type_signature.parameter,
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))
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 placements.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(
'The argument should be placed at {} or {}, placed at {} instead.'
.format(placements.SERVER, placements.CLIENTS,
arg.type_signature.placement))
return value_impl.ValueImpl(comp, self._context_stack)
def create_dummy_called_federated_apply(parameter_name,
parameter_type=tf.int32):
r"""Returns a dummy called federated apply.
Call
/ \
federated_apply Tuple
|
[Lambda(x), data]
|
Ref(x)
Args:
parameter_name: The name of the parameter.
parameter_type: The type of the parameter.
"""
fn = create_identity_function(parameter_name, parameter_type)
arg_type = computation_types.FederatedType(parameter_type, placements.SERVER)
arg = building_blocks.Data('data', arg_type)
return building_block_factory.create_federated_apply(fn, arg)
def federated_apply(self, fn, arg):
"""Implements `federated_apply` as defined in `api/intrinsics.py`.
Args:
fn: As in `api/intrinsics.py`.
arg: As in `api/intrinsics.py`.
Returns:
As in `api/intrinsics.py`.
Raises:
TypeError: As in `api/intrinsics.py`.
"""
fn = value_impl.to_value(fn, None, self._context_stack)
py_typecheck.check_type(fn, value_base.Value)
py_typecheck.check_type(fn.type_signature,
computation_types.FunctionType)
arg = value_impl.to_value(arg, None, self._context_stack)
if isinstance(arg.type_signature, computation_types.NamedTupleType):
if len(anonymous_tuple.to_elements(arg.type_signature)) >= 2:
# We've been passed a value which the user expects to be zipped.
arg = self.federated_zip(arg)
value_utils.check_federated_value_placement(arg, placements.SERVER,
'the argument')
if not arg.type_signature.all_equal:
raise TypeError('The argument should be equal at all locations.')
if not type_utils.is_assignable_from(fn.type_signature.parameter,
arg.type_signature.member):
raise TypeError(
'The function to apply expects a parameter of type {}, but member '
'constituents of the argument are of an incompatible type {}.'.
format(fn.type_signature.parameter, arg.type_signature.member))
fn = value_impl.ValueImpl.get_comp(fn)
arg = value_impl.ValueImpl.get_comp(arg)
comp = building_block_factory.create_federated_apply(fn, arg)
return value_impl.ValueImpl(comp, self._context_stack)
def test_compiles_lambda_under_federated_comp_to_tf(self):
ref_to_x = building_blocks.Reference(
'x', computation_types.StructType([tf.int32, tf.float32]))
identity_lambda = building_blocks.Lambda(ref_to_x.name,
ref_to_x.type_signature,
ref_to_x)
federated_data = building_blocks.Data(
'a',
computation_types.FederatedType(
computation_types.StructType([tf.int32, tf.float32]),
placements.SERVER))
applied = building_block_factory.create_federated_apply(
identity_lambda, federated_data)
transformed = compiler.compile_local_subcomputations_to_tensorflow(
applied)
self.assertIsInstance(transformed, building_blocks.Call)
self.assertIsInstance(transformed.function, building_blocks.Intrinsic)
self.assertIsInstance(transformed.argument[0],
building_blocks.CompiledComputation)
self.assertEqual(transformed.argument[1], federated_data)
self.assertEqual(transformed.argument[0].type_signature,
identity_lambda.type_signature)
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)
def federated_map(self, fn, arg):
"""Implements `federated_map` as defined in `api/intrinsics.py`.
Args:
fn: As in `api/intrinsics.py`.
arg: As in `api/intrinsics.py`.
Returns:
As in `api/intrinsics.py`.
Raises:
TypeError: As 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)
if isinstance(arg.type_signature, computation_types.NamedTupleType):
if len(anonymous_tuple.to_elements(arg.type_signature)) >= 2:
# We've been passed a value which the user expects to be zipped.
named_type_signatures = anonymous_tuple.to_elements(arg.type_signature)
_, first_type_signature = named_type_signatures[0]
for _, type_signature in named_type_signatures:
py_typecheck.check_type(type_signature,
computation_types.FederatedType)
if type_signature.placement is not first_type_signature.placement:
raise TypeError(
'You cannot apply federated_map on nested values with mixed '
'placements (was given a nested value of type {}).'.format(
arg.type_signature))
arg = self.federated_zip(arg)
py_typecheck.check_type(arg.type_signature, computation_types.FederatedType)
# TODO(b/113112108): Add support for polymorphic templates auto-instantiated
# here based on the actual type of the argument.
fn = value_impl.to_value(fn, None, self._context_stack)
py_typecheck.check_type(fn, value_base.Value)
py_typecheck.check_type(fn.type_signature, computation_types.FunctionType)
if not type_utils.is_assignable_from(fn.type_signature.parameter,
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))
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 placements.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(
'The argument should be placed at {} or {}, placed at {} instead.'
.format(placements.SERVER, placements.CLIENTS,
arg.type_signature.placement))
return value_impl.ValueImpl(comp, self._context_stack)
