def test_raises_type_error_with_nonfederated_arg(self):
ref = computation_building_blocks.Reference('x', tf.int32)
fn = computation_building_blocks.Lambda(ref.name, ref.type_signature,
ref)
arg = computation_building_blocks.Data('y', tf.int32)
with self.assertRaises(TypeError):
computation_constructing_utils.create_federated_map(fn, arg)
def test_raises_type_error_with_nonfunctional_fn(self):
fn = computation_building_blocks.Reference('x', tf.int32)
arg_type = computation_types.FederatedType(tf.bool, placements.CLIENTS,
False)
arg = computation_building_blocks.Data('y', arg_type)
with self.assertRaises(TypeError):
computation_constructing_utils.create_federated_map(fn, arg)
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 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.
# 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)
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)
fn = value_impl.ValueImpl.get_comp(fn)
arg = value_impl.ValueImpl.get_comp(arg)
comp = computation_constructing_utils.create_federated_map(fn, arg)
return value_impl.ValueImpl(comp, self._context_stack)
def test_returns_federated_map(self):
ref = computation_building_blocks.Reference('x', tf.int32)
fn = computation_building_blocks.Lambda(ref.name, ref.type_signature,
ref)
arg_type = computation_types.FederatedType(tf.int32,
placements.CLIENTS, False)
arg = computation_building_blocks.Data('y', arg_type)
comp = computation_constructing_utils.create_federated_map(fn, arg)
self.assertEqual(comp.tff_repr, 'federated_map(<(x -> x),y>)')
self.assertEqual(str(comp.type_signature), '{int32}@CLIENTS')
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 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, 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)
type_utils.check_federated_value_placement(arg, placements.CLIENTS,
'value to be mapped')
# 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))
fn = value_impl.ValueImpl.get_comp(fn)
arg = value_impl.ValueImpl.get_comp(arg)
comp = computation_constructing_utils.create_federated_map(fn, arg)
return value_impl.ValueImpl(comp, self._context_stack)
def create_dummy_called_federated_map(parameter_name, parameter_type=tf.int32):
r"""Returns a dummy called federated map.
Call
/ \
federated_map 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.CLIENTS)
arg = computation_building_blocks.Data('data', arg_type)
return computation_constructing_utils.create_federated_map(fn, arg)
def test_raises_type_error_with_none_fn(self):
arg = computation_building_blocks.Data('y', tf.int32)
with self.assertRaises(TypeError):
computation_constructing_utils.create_federated_map(None, arg)