def test_returns_sequence_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.SequenceType(tf.int32)
arg = computation_building_blocks.Data('y', arg_type)
comp = computation_constructing_utils.create_sequence_map(fn, arg)
self.assertEqual(comp.tff_repr, 'sequence_map(<(x -> x),y>)')
self.assertEqual(str(comp.type_signature), 'int32*')
def create_dummy_called_sequence_map(parameter_name, parameter_type=tf.int32):
r"""Returns a dummy called sequence map.
Call
/ \
sequence_map data
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.SequenceType(parameter_type)
arg = computation_building_blocks.Data('data', arg_type)
return computation_constructing_utils.create_sequence_map(fn, arg)
def sequence_map(self, fn, arg):
"""Implements `sequence_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`.
"""
fn = value_impl.to_value(fn, None, self._context_stack)
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.SequenceType):
fn = value_impl.ValueImpl.get_comp(fn)
arg = value_impl.ValueImpl.get_comp(arg)
return computation_constructing_utils.create_sequence_map(fn, arg)
elif isinstance(arg.type_signature, computation_types.FederatedType):
parameter_type = computation_types.SequenceType(
fn.type_signature.parameter)
result_type = computation_types.SequenceType(
fn.type_signature.result)
intrinsic_type = computation_types.FunctionType(
(fn.type_signature, parameter_type), result_type)
intrinsic = computation_building_blocks.Intrinsic(
intrinsic_defs.SEQUENCE_MAP.uri, intrinsic_type)
intrinsic_impl = value_impl.ValueImpl(intrinsic,
self._context_stack)
local_fn = value_utils.get_curried(intrinsic_impl)(fn)
if arg.type_signature.placement is placements.SERVER:
return self.federated_apply(local_fn, arg)
elif arg.type_signature.placement is placements.CLIENTS:
return self.federated_map(local_fn, arg)
else:
raise TypeError('Unsupported placement {}.'.format(
arg.type_signature.placement))
else:
raise TypeError(
'Cannot apply `tff.sequence_map()` to a value of type {}.'.
format(arg.type_signature))
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_sequence_map(fn, None)
def test_raises_type_error_with_none_fn(self):
arg_type = computation_types.SequenceType(tf.int32)
arg = computation_building_blocks.Data('y', arg_type)
with self.assertRaises(TypeError):
computation_constructing_utils.create_sequence_map(None, arg)