def test_returns_federated_apply(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.SERVER,
True)
arg = computation_building_blocks.Data('y', arg_type)
comp = computation_constructing_utils.create_federated_apply(fn, arg)
self.assertEqual(comp.tff_repr, 'federated_apply(<(x -> x),y>)')
self.assertEqual(str(comp.type_signature), 'int32@SERVER')
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 = computation_building_blocks.Data('data', arg_type)
return computation_constructing_utils.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)
type_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 = computation_constructing_utils.create_federated_apply(fn, arg)
return value_impl.ValueImpl(comp, self._context_stack)
def test_raises_type_error_with_none_arg(self):
ref = computation_building_blocks.Reference('x', tf.int32)
fn = computation_building_blocks.Lambda(ref.name, ref.type_signature,
ref)
with self.assertRaises(TypeError):
computation_constructing_utils.create_federated_apply(fn, None)
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_apply(None, arg)