def bar(x):
arg = anonymous_tuple.AnonymousTuple([
(str(k), x[k]) if k % 2 == 0 else (None, x[k])
for k in range(n)
])
return tff.federated_zip(arg)
async def create_value(self, value, type_spec=None):
type_spec = computation_types.to_type(type_spec)
py_typecheck.check_type(type_spec, computation_types.Type)
if isinstance(value, intrinsic_defs.IntrinsicDef):
if not type_analysis.is_concrete_instance_of(
type_spec, value.type_signature): # pytype: disable=attribute-error
raise TypeError(
'Incompatible type {} used with intrinsic {}.'.format(
type_spec, value.uri)) # pytype: disable=attribute-error
else:
return CompositeValue(value, type_spec)
elif isinstance(value, pb.Computation):
which_computation = value.WhichOneof('computation')
if which_computation in ['tensorflow', 'lambda']:
return CompositeValue(value, type_spec)
elif which_computation == 'intrinsic':
intr = intrinsic_defs.uri_to_intrinsic_def(value.intrinsic.uri)
if intr is None:
raise ValueError(
'Encountered an unrecognized intrinsic "{}".'.format(
value.intrinsic.uri))
py_typecheck.check_type(intr, intrinsic_defs.IntrinsicDef)
return await self.create_value(intr, type_spec)
else:
raise NotImplementedError(
'Unimplemented computation type {}.'.format(
which_computation))
elif isinstance(type_spec, computation_types.NamedTupleType):
value_tuple = anonymous_tuple.from_container(value)
items = await asyncio.gather(*[
self.create_value(v, t)
for v, t in zip(value_tuple, type_spec)
])
type_elemnents_iter = anonymous_tuple.iter_elements(type_spec)
return self.create_tuple(
anonymous_tuple.AnonymousTuple(
(k, i) for (k, _), i in zip(type_elemnents_iter, items)))
elif isinstance(type_spec, computation_types.FederatedType):
if type_spec.placement == placement_literals.SERVER:
if not type_spec.all_equal:
raise ValueError(
'Expected an all equal value at the `SERVER` placement, '
'found {}.'.format(type_spec))
results = await self._parent_executor.create_value(
value, type_spec.member)
return CompositeValue(results, type_spec)
elif type_spec.placement == placement_literals.CLIENTS:
if type_spec.all_equal:
results = await asyncio.gather(*[
c.create_value(value, type_spec)
for c in self._child_executors
])
return CompositeValue(results, type_spec)
else:
py_typecheck.check_type(value, list)
cardinalities = await self._get_cardinalities()
total_clients = sum(cardinalities)
py_typecheck.check_len(value, total_clients)
results = []
offset = 0
for child, num_clients in zip(self._child_executors,
cardinalities):
new_offset = offset + num_clients
result = child.create_value(value[offset:new_offset],
type_spec)
results.append(result)
offset = new_offset
return CompositeValue(await asyncio.gather(*results),
type_spec)
else:
raise ValueError('Unexpected placement {}.'.format(
type_spec.placement))
else:
return CompositeValue(
await self._parent_executor.create_value(value, type_spec),
type_spec)
def test_list_structures_from_element_type_spec_with_anonymous_tuples(self):
self._test_list_structure(
computation_types.NamedTupleType([('a', tf.int32)]), [
anonymous_tuple.AnonymousTuple([('a', 1)]),
anonymous_tuple.AnonymousTuple([('a', 2)])
], 'OrderedDict([(\'a\', array([1,2],dtype=int32))])')
def stamp_parameter_in_graph(parameter_name, parameter_type, graph):
"""Stamps a parameter of a given type in the given tf.Graph instance.
Tensors are stamped as placeholders, sequences are stamped as data sets
constructed from string tensor handles, and named tuples are stamped by
independently stamping their elements.
Args:
parameter_name: The suggested (string) name of the parameter to use in
determining the names of the graph components to construct. The names that
will actually appear in the graph are not guaranteed to be based on this
suggested name, and may vary, e.g., due to existing naming conflicts, but
a best-effort attempt will be made to make them similar for ease of
debugging.
parameter_type: The type of the parameter to stamp. Must be either an
instance of computation_types.Type (or convertible to it), or None.
graph: The instance of tf.Graph to stamp in.
Returns:
A tuple (val, binding), where 'val' is a Python object (such as a dataset,
a placeholder, or a dictionary that represents a named tuple) that
represents the stamped parameter for use in the body of a Python function
that consumes this parameter, and the 'binding' is an instance of
TensorFlow.Binding that indicates how parts of the type signature relate
to the tensors and ops stamped into the graph.
Raises:
TypeError: If the arguments are of the wrong computation_types.
ValueError: If the parameter type cannot be stamped in a TensorFlow graph.
"""
py_typecheck.check_type(parameter_name, six.string_types)
py_typecheck.check_type(graph, tf.Graph)
if parameter_type is None:
return (None, None)
parameter_type = computation_types.to_type(parameter_type)
if isinstance(parameter_type, computation_types.TensorType):
with graph.as_default():
placeholder = tf.placeholder(dtype=parameter_type.dtype,
shape=parameter_type.shape,
name=parameter_name)
binding = pb.TensorFlow.Binding(tensor=pb.TensorFlow.TensorBinding(
tensor_name=placeholder.name))
return (placeholder, binding)
elif isinstance(parameter_type, computation_types.NamedTupleType):
element_name_value_pairs = []
element_bindings = []
for e in anonymous_tuple.to_elements(parameter_type):
e_val, e_binding = stamp_parameter_in_graph(
'{}_{}'.format(parameter_name, e[0]), e[1], graph)
element_name_value_pairs.append((e[0], e_val))
element_bindings.append(e_binding)
return (anonymous_tuple.AnonymousTuple(element_name_value_pairs),
pb.TensorFlow.Binding(tuple=pb.TensorFlow.NamedTupleBinding(
element=element_bindings)))
elif isinstance(parameter_type, computation_types.SequenceType):
with graph.as_default():
handle = tf.placeholder(tf.string, shape=[])
ds = make_dataset_from_string_handle(handle, parameter_type.element)
return (ds,
pb.TensorFlow.Binding(sequence=pb.TensorFlow.SequenceBinding(
iterator_string_handle_name=handle.name)))
else:
raise ValueError(
'Parameter type component {} cannot be stamped into a TensorFlow '
'graph.'.format(repr(parameter_type)))
class FunctionUtilsTest(test.TestCase, parameterized.TestCase):
def test_get_defun_argspec_with_typed_non_eager_defun(self):
# In a tf.function with a defined input signature, **kwargs or default
# values are not allowed, but *args are, and the input signature may overlap
# with *args.
fn = tf.function(lambda x, y, *z: None, (
tf.TensorSpec(None, tf.int32),
tf.TensorSpec(None, tf.bool),
tf.TensorSpec(None, tf.float32),
tf.TensorSpec(None, tf.float32),
))
self.assertEqual(
collections.OrderedDict(
function_utils.get_signature(fn).parameters),
collections.OrderedDict(
x=inspect.Parameter('x',
inspect.Parameter.POSITIONAL_OR_KEYWORD),
y=inspect.Parameter('y',
inspect.Parameter.POSITIONAL_OR_KEYWORD),
z=inspect.Parameter('z', inspect.Parameter.VAR_POSITIONAL),
))
def test_get_defun_argspec_with_untyped_non_eager_defun(self):
# In a tf.function with no input signature, the same restrictions as in a
# typed eager function apply.
fn = tf.function(lambda x, y, *z: None)
self.assertEqual(
collections.OrderedDict(
function_utils.get_signature(fn).parameters),
collections.OrderedDict(
x=inspect.Parameter('x',
inspect.Parameter.POSITIONAL_OR_KEYWORD),
y=inspect.Parameter('y',
inspect.Parameter.POSITIONAL_OR_KEYWORD),
z=inspect.Parameter('z', inspect.Parameter.VAR_POSITIONAL),
))
def test_get_signature_with_class_instance_method(self):
class C:
def __init__(self, x):
self._x = x
def foo(self, y):
return self._x * y
c = C(5)
signature = function_utils.get_signature(c.foo)
self.assertEqual(
signature.parameters,
collections.OrderedDict(y=inspect.Parameter(
'y', inspect.Parameter.POSITIONAL_OR_KEYWORD)))
def test_get_signature_with_class_property(self):
class C:
@property
def x(self):
return 99
c = C()
with self.assertRaises(TypeError):
function_utils.get_signature(c.x)
def test_as_wrapper_with_classmethod(self):
class C:
@classmethod
def foo(cls, x):
return x * 2
signature = function_utils.get_signature(C.foo)
self.assertEqual(
signature.parameters,
collections.OrderedDict(x=inspect.Parameter(
'x', inspect.Parameter.POSITIONAL_OR_KEYWORD)))
# pyformat: disable
@parameterized.parameters(
itertools.product(
# Values of 'fn' to test.
[
lambda: None, lambda a: None, lambda a, b: None,
lambda *a: None, lambda **a: None, lambda *a, **b: None,
lambda a, *b: None, lambda a, **b: None,
lambda a, b, **c: None, lambda a, b=10: None,
lambda a, b=10, c=20: None, lambda a, b=10, *c: None,
lambda a, b=10, **c: None, lambda a, b=10, *c, **d: None,
lambda a, b, c=10, *d: None, lambda a=10, b=20, c=30, **d: None
],
# Values of 'args' to test.
[[], [1], [1, 2], [1, 2, 3], [1, 2, 3, 4]],
# Values of 'kwargs' to test.
[{}, {
'b': 100
}, {
'name': 'foo'
}, {
'b': 100,
'name': 'foo'
}]))
# pyformat: enable
def test_get_callargs_for_signature(self, fn, args, kwargs):
signature = function_utils.get_signature(fn)
expected_error = None
try:
signature = inspect.signature(fn)
bound_arguments = signature.bind(*args, **kwargs)
expected_callargs = bound_arguments.arguments
except TypeError as e:
expected_error = e
expected_callargs = None
result_callargs = None
if expected_error is None:
try:
bound_args = signature.bind(*args, **kwargs).arguments
self.assertEqual(bound_args, expected_callargs)
except (TypeError, AssertionError) as test_err:
raise AssertionError(
'With signature `{!s}`, args {!s}, kwargs {!s}, expected bound '
'args {!s} and error {!s}, tested function returned {!s} and the '
'test has failed with message: {!s}'.format(
signature, args, kwargs, expected_callargs,
expected_error, result_callargs, test_err))
else:
with self.assertRaises(TypeError):
_ = signature.bind(*args, **kwargs)
# pyformat: disable
# pylint: disable=g-complex-comprehension
@parameterized.parameters(
(function_utils.get_signature(params[0]), ) + params[1:]
for params in [
(lambda a: None, [tf.int32], {}),
(lambda a, b=True: None, [tf.int32, tf.bool], {}),
(lambda a, b=True: None, [tf.int32], {
'b': tf.bool
}),
(lambda a, b=True: None, [tf.bool], {
'b': tf.bool
}),
(lambda a=10, b=True: None, [tf.int32], {
'b': tf.bool
}),
])
# pylint: enable=g-complex-comprehension
# pyformat: enable
def test_is_signature_compatible_with_types_true(self, signature, args,
kwargs):
self.assertTrue(
function_utils.is_signature_compatible_with_types(
signature, *[computation_types.to_type(a) for a in args],
**{k: computation_types.to_type(v)
for k, v in kwargs.items()}))
# pyformat: disable
# pylint: disable=g-complex-comprehension
@parameterized.parameters(
(function_utils.get_signature(params[0]), ) + params[1:]
for params in [
(lambda a=True: None, [tf.int32], {}),
(lambda a=10, b=True: None, [tf.bool], {
'b': tf.bool
}),
])
# pylint: enable=g-complex-comprehension
# pyformat: enable
def test_is_signature_compatible_with_types_false(self, signature, args,
kwargs):
self.assertFalse(
function_utils.is_signature_compatible_with_types(
signature, *[computation_types.to_type(a) for a in args],
**{k: computation_types.to_type(v)
for k, v in kwargs.items()}))
# pyformat: disable
@parameterized.parameters(
(tf.int32, False), ([tf.int32, tf.int32], True),
([tf.int32, ('b', tf.int32)], True), ([('a', tf.int32),
('b', tf.int32)], True),
([('a', tf.int32), tf.int32], False),
(anonymous_tuple.AnonymousTuple([(None, 1), ('a', 2)]), True),
(anonymous_tuple.AnonymousTuple([('a', 1), (None, 2)]), False))
# pyformat: enable
def test_is_argument_tuple(self, arg, expected_result):
self.assertEqual(function_utils.is_argument_tuple(arg),
expected_result)
# pyformat: disable
@parameterized.parameters(
(anonymous_tuple.AnonymousTuple([(None, 1)]), [1], {}),
(anonymous_tuple.AnonymousTuple([(None, 1), ('a', 2)]), [1], {
'a': 2
}))
# pyformat: enable
def test_unpack_args_from_anonymous_tuple(self, tuple_with_args,
expected_args, expected_kwargs):
self.assertEqual(
function_utils.unpack_args_from_tuple(tuple_with_args),
(expected_args, expected_kwargs))
# pyformat: disable
@parameterized.parameters(
([tf.int32], [tf.int32], {}), ([('a', tf.int32)], [], {
'a': tf.int32
}), ([tf.int32, tf.bool], [tf.int32, tf.bool], {}),
([tf.int32, ('b', tf.bool)], [tf.int32], {
'b': tf.bool
}), ([('a', tf.int32), ('b', tf.bool)], [], {
'a': tf.int32,
'b': tf.bool
}))
# pyformat: enable
def test_unpack_args_from_tuple_type(self, tuple_with_args, expected_args,
expected_kwargs):
args, kwargs = function_utils.unpack_args_from_tuple(tuple_with_args)
self.assertEqual(len(args), len(expected_args))
for idx, arg in enumerate(args):
self.assertTrue(
type_analysis.are_equivalent_types(
arg, computation_types.to_type(expected_args[idx])))
self.assertEqual(set(kwargs.keys()), set(expected_kwargs.keys()))
for k, v in kwargs.items():
self.assertTrue(
type_analysis.are_equivalent_types(
computation_types.to_type(v), expected_kwargs[k]))
def test_pack_args_into_anonymous_tuple_without_type_spec(self):
self.assertEqual(
function_utils.pack_args_into_anonymous_tuple([1], {'a': 10}),
anonymous_tuple.AnonymousTuple([(None, 1), ('a', 10)]))
self.assertIn(
function_utils.pack_args_into_anonymous_tuple([1, 2], {
'a': 10,
'b': 20
}), [
anonymous_tuple.AnonymousTuple([
(None, 1),
(None, 2),
('a', 10),
('b', 20),
]),
anonymous_tuple.AnonymousTuple([
(None, 1),
(None, 2),
('b', 20),
('a', 10),
])
])
self.assertIn(
function_utils.pack_args_into_anonymous_tuple([], {
'a': 10,
'b': 20
}), [
anonymous_tuple.AnonymousTuple([('a', 10), ('b', 20)]),
anonymous_tuple.AnonymousTuple([('b', 20), ('a', 10)])
])
self.assertEqual(
function_utils.pack_args_into_anonymous_tuple([1], {}),
anonymous_tuple.AnonymousTuple([(None, 1)]))
# pyformat: disable
@parameterized.parameters(
([1], {}, [tf.int32], [(None, 1)]),
([1, True], {}, [tf.int32, tf.bool], [(None, 1), (None, True)]),
([1, True], {}, [('x', tf.int32),
('y', tf.bool)], [('x', 1), ('y', True)]), ([1], {
'y': True
}, [('x', tf.int32), ('y', tf.bool)], [('x', 1),
('y', True)]),
([], {
'x': 1,
'y': True
}, [('x', tf.int32), ('y', tf.bool)], [('x', 1), ('y', True)]),
([], collections.OrderedDict([('y', True), ('x', 1)]), [
('x', tf.int32), ('y', tf.bool)
], [('x', 1), ('y', True)]))
# pyformat: enable
def test_pack_args_into_anonymous_tuple_with_type_spec_expect_success(
self, args, kwargs, type_spec, elements):
self.assertEqual(
function_utils.pack_args_into_anonymous_tuple(
args, kwargs, type_spec, NoopIngestContextForTest()),
anonymous_tuple.AnonymousTuple(elements))
# pyformat: disable
@parameterized.parameters(([1], {}, [(tf.bool)]), ([], {
'x': 1,
'y': True
}, [(tf.int32), (tf.bool)]))
# pyformat: enable
def test_pack_args_into_anonymous_tuple_with_type_spec_expect_failure(
self, args, kwargs, type_spec):
with self.assertRaises(TypeError):
function_utils.pack_args_into_anonymous_tuple(
args, kwargs, type_spec, NoopIngestContextForTest())
# pyformat: disable
@parameterized.parameters(
(None, [], {}, 'None'), (tf.int32, [1], {}, '1'),
([tf.int32, tf.bool], [1, True], {}, '<1,True>'),
([('x', tf.int32), ('y', tf.bool)], [1, True], {}, '<x=1,y=True>'),
([('x', tf.int32), ('y', tf.bool)], [1], {
'y': True
}, '<x=1,y=True>'),
([tf.int32, tf.bool],
[anonymous_tuple.AnonymousTuple([(None, 1),
(None, True)])], {}, '<1,True>'))
# pyformat: enable
def test_pack_args(self, parameter_type, args, kwargs,
expected_value_string):
self.assertEqual(
str(
function_utils.pack_args(parameter_type, args, kwargs,
NoopIngestContextForTest())),
expected_value_string)
# pyformat: disable
@parameterized.parameters(
(1, lambda: 10, None, None, None, 10),
(2, lambda x=1: x + 10, None, None, None, 11),
(3, lambda x=1: x + 10, tf.int32, None, 20, 30),
(4, lambda x, y: x + y, [tf.int32, tf.int32], None,
anonymous_tuple.AnonymousTuple([('x', 5), ('y', 6)]), 11),
(5, lambda *args: str(args), [tf.int32, tf.int32], True,
anonymous_tuple.AnonymousTuple([(None, 5), (None, 6)]), '(5, 6)'),
(6, lambda *args: str(args), [
('x', tf.int32), ('y', tf.int32)
], False, anonymous_tuple.AnonymousTuple([
('x', 5), ('y', 6)
]), '(AnonymousTuple([(\'x\', 5), (\'y\', 6)]),)'),
(
7,
lambda x: str(x), # pylint: disable=unnecessary-lambda
[tf.int32],
None,
anonymous_tuple.AnonymousTuple([(None, 10)]),
'[10]'))
# pyformat: enable
def test_wrap_as_zero_or_one_arg_callable(self, unused_index, fn,
parameter_type, unpack, arg,
expected_result):
wrapped_fn = function_utils.wrap_as_zero_or_one_arg_callable(
fn, parameter_type, unpack)
actual_result = wrapped_fn(arg) if parameter_type else wrapped_fn()
self.assertEqual(actual_result, expected_result)
def test_returns_false_with_named_tuple_type_spec(self):
value = anonymous_tuple.AnonymousTuple([('a', 0.0)])
type_spec = computation_types.NamedTupleType([('a', tf.float32)])
self.assertFalse(
type_analysis.is_anon_tuple_with_py_container(value, type_spec))
async def create_value(self, value, type_spec=None):
type_spec = computation_types.to_type(type_spec)
py_typecheck.check_type(type_spec, computation_types.Type)
if isinstance(value, intrinsic_defs.IntrinsicDef):
if not type_utils.is_concrete_instance_of(type_spec,
value.type_signature): # pytype: disable=attribute-error
raise TypeError(
'Incompatible type {} used with intrinsic {}.'.format(
type_spec, value.uri)) # pytype: disable=attribute-error
else:
return CompositeValue(value, type_spec)
elif isinstance(value, pb.Computation):
which_computation = value.WhichOneof('computation')
if which_computation in ['tensorflow', 'lambda']:
return CompositeValue(value, type_spec)
elif which_computation == 'intrinsic':
intr = intrinsic_defs.uri_to_intrinsic_def(value.intrinsic.uri)
if intr is None:
raise ValueError(
'Encountered an unrecognized intrinsic "{}".'.format(
value.intrinsic.uri))
py_typecheck.check_type(intr, intrinsic_defs.IntrinsicDef)
return await self.create_value(intr, type_spec)
else:
raise NotImplementedError(
'Unimplemented computation type {}.'.format(
which_computation))
elif isinstance(type_spec, computation_types.NamedTupleType):
v_el = anonymous_tuple.to_elements(
anonymous_tuple.from_container(value))
t_el = anonymous_tuple.to_elements(type_spec)
items = await asyncio.gather(*[
self.create_value(v, t) for (_, v), (_, t) in zip(v_el, t_el)
])
return self.create_tuple(
anonymous_tuple.AnonymousTuple([
(k, i) for (k, _), i in zip(t_el, items)
]))
elif isinstance(type_spec, computation_types.FederatedType):
if type_spec.placement == placement_literals.SERVER:
if type_spec.all_equal:
return CompositeValue(
await self._parent_executor.create_value(
value, type_spec.member), type_spec)
else:
raise ValueError(
'A non-all_equal value on the server is unexpected.')
elif type_spec.placement == placement_literals.CLIENTS:
if type_spec.all_equal:
return CompositeValue(
await asyncio.gather(*[
c.create_value(value, type_spec)
for c in self._child_executors
]), type_spec)
else:
py_typecheck.check_type(value, list)
if self._cardinalities is None:
self._cardinalities = asyncio.ensure_future(
self._get_cardinalities())
cardinalities = await self._cardinalities
py_typecheck.check_len(cardinalities,
len(self._child_executors))
count = sum(cardinalities)
py_typecheck.check_len(value, count)
result = []
offset = 0
for c, n in zip(self._child_executors, cardinalities):
new_offset = offset + n
# The slice opporator is not supported on all the types `value`
# supports.
# pytype: disable=unsupported-operands
result.append(
c.create_value(value[offset:new_offset],
type_spec))
# pytype: enable=unsupported-operands
offset = new_offset
return CompositeValue(await asyncio.gather(*result),
type_spec)
else:
raise ValueError('Unexpected placement {}.'.format(
type_spec.placement))
else:
return CompositeValue(
await self._parent_executor.create_value(value, type_spec),
type_spec)
def test_getitem_key_builtin_attribute_raises(self):
v = [('foo', 10), ('bar', 20)]
x = anonymous_tuple.AnonymousTuple(v)
with self.assertRaises(AttributeError):
_ = x['__getattr__']
def test_getitem_bad_bounds(self):
v = [(None, i) for i in range(0, 50, 10)]
x = anonymous_tuple.AnonymousTuple(v)
with self.assertRaises(IndexError):
_ = x[10]
async def _child_fn(ex, x, y):
py_typecheck.check_type(x, executor_value_base.ExecutorValue)
py_typecheck.check_type(y, executor_value_base.ExecutorValue)
return await ex.create_call(
await ex.create_value(zip_comp, zip_type), await ex.create_tuple(
anonymous_tuple.AnonymousTuple([(keys[0], x), (keys[1], y)])))
def test_construction_from_tuple(self):
v = (('a', 1), ('b', 2), (None, 3))
x = anonymous_tuple.AnonymousTuple(v)
self.assertSequenceEqual(anonymous_tuple.to_elements(x), v)
def comp():
return anonymous_tuple.AnonymousTuple([('a', 1), ('b', 2)])
async def _evaluate(self, comp, scope=None):
"""Evaluates or partially evaluates `comp` in `scope`.
Args:
comp: An instance of `pb.Computation` to process.
scope: An optional `LambdaExecutorScope` to process it in, or `None` if
there's no surrounding scope (the computation is self-contained).
Returns:
An instance of `LambdaExecutorValue` that isn't unprocessed (i.e., the
internal representation directly in it isn't simply `comp` or any
other instance of`pb.Computation`). The result, however, does not have,
and often won't be processed completely; it suffices for this function
to make only partial progress.
"""
py_typecheck.check_type(comp, pb.Computation)
if scope is not None:
py_typecheck.check_type(scope, LambdaExecutorScope)
which_computation = comp.WhichOneof('computation')
if which_computation in [
'tensorflow', 'intrinsic', 'data', 'placement'
]:
return LambdaExecutorValue(
await self._target_executor.create_value(
comp,
type_utils.get_function_type(
type_serialization.deserialize_type(comp.type))))
elif which_computation == 'lambda':
def _make_comp_fn(scope, result, name, type_spec):
async def _comp_fn(arg):
return await self._evaluate(
result, LambdaExecutorScope({name: arg}, scope))
return LambdaExecutorValue(_comp_fn, type_spec=type_spec)
comp_lambda = getattr(comp, 'lambda')
type_spec = type_utils.get_function_type(
type_serialization.deserialize_type(comp.type))
return _make_comp_fn(scope, comp_lambda.result,
comp_lambda.parameter_name, type_spec)
elif which_computation == 'reference':
return scope.resolve_reference(comp.reference.name)
elif which_computation == 'call':
if comp.call.argument.WhichOneof('computation') is not None:
arg = LambdaExecutorValue(comp.call.argument, scope=scope)
else:
arg = None
return await self.create_call(LambdaExecutorValue(
comp.call.function, scope=scope),
arg=arg)
elif which_computation == 'selection':
which_selection = comp.selection.WhichOneof('selection')
return await self.create_selection(
await self.create_call(
LambdaExecutorValue(comp.selection.source, scope=scope)),
**{which_selection: getattr(comp.selection, which_selection)})
elif which_computation == 'tuple':
names = [
str(e.name) if e.name else None for e in comp.tuple.element
]
values = []
for e in comp.tuple.element:
val = LambdaExecutorValue(e.value, scope=scope)
if (isinstance(val.type_signature,
computation_types.FunctionType)
and val.type_signature.parameter is None):
val = self.create_call(val)
else:
async def _async_identity(x):
return x
val = _async_identity(val)
values.append(val)
values = await asyncio.gather(*values)
return await self.create_tuple(
anonymous_tuple.AnonymousTuple(list(zip(names, values))))
elif which_computation == 'block':
for loc in comp.block.local:
value = LambdaExecutorValue(loc.value, scope)
scope = LambdaExecutorScope({loc.name: value}, scope)
return await self._evaluate(comp.block.result, scope)
else:
raise NotImplementedError(
'Unsupported computation type "{}".'.format(which_computation))
class FuncUtilsTest(test.TestCase, parameterized.TestCase):
def test_is_defun(self):
self.assertTrue(
function_utils.is_defun(function.Defun()(lambda x: None)))
self.assertTrue(
function_utils.is_defun(function.Defun(tf.int32)(lambda x: None)))
self.assertFalse(function_utils.is_defun(function.Defun))
self.assertFalse(function_utils.is_defun(lambda x: None))
self.assertFalse(function_utils.is_defun(None))
def test_get_defun_argspec_with_typed_non_eager_defun(self):
# In a non-eager defun with a defined input signature, **kwargs or default
# values are not allowed, but *args are, and the input signature may
# overlap with *args.
self.assertEqual(
function_utils.get_argspec(
function.Defun(tf.int32, tf.bool, tf.float32,
tf.float32)(lambda x, y, *z: None)),
inspect.ArgSpec(args=['x', 'y'],
varargs='z',
keywords=None,
defaults=None))
def test_get_defun_argspec_with_untyped_non_eager_defun(self):
# In a non-eager defun with no input signature, the same restrictions as in
# a typed defun apply.
self.assertEqual(
function_utils.get_argspec(
function.Defun()(lambda x, y, *z: None)),
inspect.ArgSpec(args=['x', 'y'],
varargs='z',
keywords=None,
defaults=None))
# pyformat: disable
@parameterized.parameters(
itertools.product(
# Values of 'fn' to test.
[
lambda: None, lambda a: None, lambda a, b: None,
lambda *a: None, lambda **a: None, lambda *a, **b: None,
lambda a, *b: None, lambda a, **b: None,
lambda a, b, **c: None, lambda a, b=10: None,
lambda a, b=10, c=20: None, lambda a, b=10, *c: None,
lambda a, b=10, **c: None, lambda a, b=10, *c, **d: None,
lambda a, b, c=10, *d: None, lambda a=10, b=20, c=30, **d: None
],
# Values of 'args' to test.
[[], [1], [1, 2], [1, 2, 3], [1, 2, 3, 4]],
# Values of 'kwargs' to test.
[{}, {
'b': 100
}, {
'name': 'foo'
}, {
'b': 100,
'name': 'foo'
}]))
# pyformat: enable
def test_get_callargs_for_argspec(self, fn, args, kwargs):
argspec = inspect.getargspec(fn)
expected_error = None
try:
expected_callargs = inspect.getcallargs(fn, *args, **kwargs)
except TypeError as e:
expected_error = e
expected_callargs = None
try:
if expected_error is None:
result_callargs = function_utils.get_callargs_for_argspec(
argspec, *args, **kwargs)
self.assertEqual(result_callargs, expected_callargs)
else:
with self.assertRaises(TypeError):
result_callargs = function_utils.get_callargs_for_argspec(
argspec, *args, **kwargs)
except (TypeError, AssertionError) as test_err:
raise AssertionError(
'With argspec {}, args {}, kwargs {}, expected callargs {} and '
'error {}, tested function returned {} and the test has failed '
'with message: {}'.format(str(argspec), str(args), str(kwargs),
str(expected_callargs),
str(expected_error),
str(result_callargs), str(test_err)))
# pyformat: disable
@parameterized.parameters(
(inspect.getargspec(params[0]), ) + params[1:] for params in
[(lambda a: None, [tf.int32], {},
True), (
lambda a=True: None,
[tf.int32], {},
False), (lambda a, b=True: None, [tf.int32, tf.bool], {}, True),
(lambda a, b=True: None, [tf.int32], {
'b': tf.bool
}, True), (lambda a, b=True: None, [tf.bool], {
'b': tf.bool
}, True), (lambda a=10, b=True: None, [tf.int32], {
'b': tf.bool
}, True), (lambda a=10, b=True: None, [tf.bool], {
'b': tf.bool
}, False)])
# pyformat: enable
def test_is_argspec_compatible_with_types(self, argspec, args, kwargs,
expected_result):
self.assertEqual(
function_utils.is_argspec_compatible_with_types(
argspec, *[computation_types.to_type(a) for a in args], **{
k: computation_types.to_type(v)
for k, v in six.iteritems(kwargs)
}), expected_result)
# pyformat: disable
@parameterized.parameters(
(tf.int32, False), ([tf.int32, tf.int32], True),
([tf.int32, ('b', tf.int32)], True), ([('a', tf.int32),
('b', tf.int32)], True),
([('a', tf.int32), tf.int32], False),
(anonymous_tuple.AnonymousTuple([(None, 1), ('a', 2)]), True),
(anonymous_tuple.AnonymousTuple([('a', 1), (None, 2)]), False))
# pyformat: enable
def test_is_argument_tuple(self, arg, expected_result):
self.assertEqual(function_utils.is_argument_tuple(arg),
expected_result)
# pyformat: disable
@parameterized.parameters(
(anonymous_tuple.AnonymousTuple([(None, 1)]), [1], {}),
(anonymous_tuple.AnonymousTuple([(None, 1), ('a', 2)]), [1], {
'a': 2
}))
# pyformat: enable
def test_unpack_args_from_anonymous_tuple(self, tuple_with_args,
expected_args, expected_kwargs):
self.assertEqual(
function_utils.unpack_args_from_tuple(tuple_with_args),
(expected_args, expected_kwargs))
# pyformat: disable
@parameterized.parameters(
([tf.int32], [tf.int32], {}), ([('a', tf.int32)], [], {
'a': tf.int32
}), ([tf.int32, tf.bool], [tf.int32, tf.bool], {}),
([tf.int32, ('b', tf.bool)], [tf.int32], {
'b': tf.bool
}), ([('a', tf.int32), ('b', tf.bool)], [], {
'a': tf.int32,
'b': tf.bool
}))
# pyformat: enable
def test_unpack_args_from_tuple_type(self, tuple_with_args, expected_args,
expected_kwargs):
args, kwargs = function_utils.unpack_args_from_tuple(tuple_with_args)
self.assertEqual(len(args), len(expected_args))
for idx, arg in enumerate(args):
self.assertTrue(
type_utils.are_equivalent_types(
arg, computation_types.to_type(expected_args[idx])))
self.assertEqual(set(kwargs.keys()), set(expected_kwargs.keys()))
for k, v in six.iteritems(kwargs):
self.assertTrue(
type_utils.are_equivalent_types(computation_types.to_type(v),
expected_kwargs[k]))
def test_pack_args_into_anonymous_tuple_without_type_spec(self):
self.assertEqual(
function_utils.pack_args_into_anonymous_tuple([1], {'a': 10}),
anonymous_tuple.AnonymousTuple([(None, 1), ('a', 10)]))
self.assertIn(
function_utils.pack_args_into_anonymous_tuple([1, 2], {
'a': 10,
'b': 20
}), [
anonymous_tuple.AnonymousTuple([
(None, 1),
(None, 2),
('a', 10),
('b', 20),
]),
anonymous_tuple.AnonymousTuple([
(None, 1),
(None, 2),
('b', 20),
('a', 10),
])
])
self.assertIn(
function_utils.pack_args_into_anonymous_tuple([], {
'a': 10,
'b': 20
}), [
anonymous_tuple.AnonymousTuple([('a', 10), ('b', 20)]),
anonymous_tuple.AnonymousTuple([('b', 20), ('a', 10)])
])
self.assertEqual(
function_utils.pack_args_into_anonymous_tuple([1], {}),
anonymous_tuple.AnonymousTuple([(None, 1)]))
# pyformat: disable
@parameterized.parameters(
([1], {}, [tf.int32], [(None, 1)]),
([1, True], {}, [tf.int32, tf.bool], [(None, 1), (None, True)]),
([1, True], {}, [('x', tf.int32),
('y', tf.bool)], [('x', 1), ('y', True)]), ([1], {
'y': True
}, [('x', tf.int32), ('y', tf.bool)], [('x', 1),
('y', True)]),
([], {
'x': 1,
'y': True
}, [('x', tf.int32), ('y', tf.bool)], [('x', 1), ('y', True)]),
([], collections.OrderedDict([('y', True), ('x', 1)]), [
('x', tf.int32), ('y', tf.bool)
], [('x', 1), ('y', True)]))
# pyformat: enable
def test_pack_args_into_anonymous_tuple_with_type_spec_expect_success(
self, args, kwargs, type_spec, elements):
self.assertEqual(
function_utils.pack_args_into_anonymous_tuple(
args, kwargs, type_spec, NoopIngestContextForTest()),
anonymous_tuple.AnonymousTuple(elements))
# pyformat: disable
@parameterized.parameters(([1], {}, [(tf.bool)]), ([], {
'x': 1,
'y': True
}, [(tf.int32), (tf.bool)]))
# pyformat: enable
def test_pack_args_into_anonymous_tuple_with_type_spec_expect_failure(
self, args, kwargs, type_spec):
with self.assertRaises(TypeError):
function_utils.pack_args_into_anonymous_tuple(
args, kwargs, type_spec, NoopIngestContextForTest())
# pyformat: disable
@parameterized.parameters(
(None, [], {}, 'None'), (tf.int32, [1], {}, '1'),
([tf.int32, tf.bool], [1, True], {}, '<1,True>'),
([('x', tf.int32), ('y', tf.bool)], [1, True], {}, '<x=1,y=True>'),
([('x', tf.int32), ('y', tf.bool)], [1], {
'y': True
}, '<x=1,y=True>'),
([tf.int32, tf.bool],
[anonymous_tuple.AnonymousTuple([(None, 1),
(None, True)])], {}, '<1,True>'))
# pyformat: enable
def test_pack_args(self, parameter_type, args, kwargs,
expected_value_string):
self.assertEqual(
str(
function_utils.pack_args(parameter_type, args, kwargs,
NoopIngestContextForTest())),
expected_value_string)
# pyformat: disable
@parameterized.parameters(
(1, lambda: 10, None, None, None, 10),
(2, lambda x=1: x + 10, None, None, None, 11),
(3, lambda x=1: x + 10, tf.int32, None, 20, 30),
(4, lambda x, y: x + y, [tf.int32, tf.int32], None,
anonymous_tuple.AnonymousTuple([('x', 5), ('y', 6)]), 11),
(5, lambda *args: str(args), [tf.int32, tf.int32], True,
anonymous_tuple.AnonymousTuple([(None, 5), (None, 6)]), '(5, 6)'),
(6, lambda *args: str(args), [
('x', tf.int32), ('y', tf.int32)
], False, anonymous_tuple.AnonymousTuple(
[('x', 5), ('y', 6)]), '(AnonymousTuple([(x, 5), (y, 6)]),)'),
(
7,
lambda x: str(x), # pylint: disable=unnecessary-lambda
[tf.int32],
None,
anonymous_tuple.AnonymousTuple([(None, 10)]),
'<10>'))
# pyformat: enable
def test_wrap_as_zero_or_one_arg_callable(self, unused_index, fn,
parameter_type, unpack, arg,
expected_result):
wrapped_fn = function_utils.wrap_as_zero_or_one_arg_callable(
fn, parameter_type, unpack)
actual_result = wrapped_fn(arg) if parameter_type else wrapped_fn()
self.assertEqual(actual_result, expected_result)
def test_polymorphic_function(self):
class ContextForTest(context_base.Context):
def ingest(self, val, type_spec):
return val
def invoke(self, comp, arg):
return 'name={},type={},arg={}'.format(
comp.name, str(comp.type_signature.parameter), str(arg))
class TestFunction(function_utils.ConcreteFunction):
def __init__(self, name, parameter_type):
self._name = name
super(TestFunction, self).__init__(
computation_types.FunctionType(parameter_type, tf.string),
context_stack_impl.context_stack)
@property
def name(self):
return self._name
class TestFunctionFactory(object):
def __init__(self):
self._count = 0
def __call__(self, parameter_type):
self._count = self._count + 1
return TestFunction(str(self._count), parameter_type)
with context_stack_impl.context_stack.install(ContextForTest()):
fn = function_utils.PolymorphicFunction(TestFunctionFactory())
self.assertEqual(fn(10), 'name=1,type=<int32>,arg=<10>')
self.assertEqual(fn(20, x=True),
'name=2,type=<int32,x=bool>,arg=<20,x=True>')
self.assertEqual(fn(True), 'name=3,type=<bool>,arg=<True>')
self.assertEqual(fn(30, x=40),
'name=4,type=<int32,x=int32>,arg=<30,x=40>')
self.assertEqual(fn(50), 'name=1,type=<int32>,arg=<50>')
self.assertEqual(fn(0, x=False),
'name=2,type=<int32,x=bool>,arg=<0,x=False>')
self.assertEqual(fn(False), 'name=3,type=<bool>,arg=<False>')
self.assertEqual(fn(60, x=70),
'name=4,type=<int32,x=int32>,arg=<60,x=70>')
def test_concrete_function(self):
class ContextForTest(context_base.Context):
def ingest(self, val, type_spec):
return val
def invoke(self, comp, arg):
return comp.invoke_fn(arg)
class TestFunction(function_utils.ConcreteFunction):
def __init__(self, type_signature, invoke_fn):
super(TestFunction,
self).__init__(type_signature,
context_stack_impl.context_stack)
self._invoke_fn = invoke_fn
def invoke_fn(self, arg):
return self._invoke_fn(arg)
with context_stack_impl.context_stack.install(ContextForTest()):
fn = TestFunction(
computation_types.FunctionType(tf.int32, tf.bool),
lambda x: x > 10)
self.assertEqual(fn(5), False)
self.assertEqual(fn(15), True)
fn = TestFunction(
computation_types.FunctionType([('x', tf.int32),
('y', tf.int32)], tf.bool),
lambda arg: arg.x > arg.y)
self.assertEqual(fn(5, 10), False)
self.assertEqual(fn(10, 5), True)
self.assertEqual(fn(y=10, x=5), False)
self.assertEqual(fn(y=5, x=10), True)
self.assertEqual(fn(10, y=5), True)
async def create_value(self, value, type_spec=None):
"""A coroutine that creates embedded value from `value` of type `type_spec`.
See the `FederatingExecutorValue` for detailed information about the
`value`s and `type_spec`s that can be embedded using `create_value`.
Args:
value: An object that represents the value to embed within the executor.
type_spec: An optional `tff.Type` of the value represented by this object,
or something convertible to it.
Returns:
An instance of `FederatingExecutorValue` that represents the embedded
value.
Raises:
TypeError: If the `value` and `type_spec` do not match.
ValueError: If `value` is not a kind recognized by the
`FederatingExecutor`.
"""
type_spec = computation_types.to_type(type_spec)
if isinstance(type_spec, computation_types.FederatedType):
self._check_executor_compatible_with_placement(type_spec.placement)
elif (isinstance(type_spec, computation_types.FunctionType) and
isinstance(type_spec.result, computation_types.FederatedType)):
self._check_executor_compatible_with_placement(
type_spec.result.placement)
if isinstance(value, intrinsic_defs.IntrinsicDef):
if not type_utils.is_concrete_instance_of(type_spec,
value.type_signature):
raise TypeError(
'Incompatible type {} used with intrinsic {}.'.format(
type_spec, value.uri))
return FederatingExecutorValue(value, type_spec)
elif isinstance(value, placement_literals.PlacementLiteral):
if type_spec is None:
type_spec = computation_types.PlacementType()
else:
py_typecheck.check_type(type_spec,
computation_types.PlacementType)
return FederatingExecutorValue(value, type_spec)
elif isinstance(value, computation_impl.ComputationImpl):
return await self.create_value(
computation_impl.ComputationImpl.get_proto(value),
type_utils.reconcile_value_with_type_spec(value, type_spec))
elif isinstance(value, pb.Computation):
if type_spec is None:
type_spec = type_serialization.deserialize_type(value.type)
which_computation = value.WhichOneof('computation')
if which_computation in ['lambda', 'tensorflow']:
return FederatingExecutorValue(value, type_spec)
elif which_computation == 'reference':
raise ValueError(
'Encountered an unexpected unbound references "{}".'.
format(value.reference.name))
elif which_computation == 'intrinsic':
intr = intrinsic_defs.uri_to_intrinsic_def(value.intrinsic.uri)
if intr is None:
raise ValueError(
'Encountered an unrecognized intrinsic "{}".'.format(
value.intrinsic.uri))
py_typecheck.check_type(intr, intrinsic_defs.IntrinsicDef)
return await self.create_value(intr, type_spec)
elif which_computation == 'placement':
return await self.create_value(
placement_literals.uri_to_placement_literal(
value.placement.uri), type_spec)
elif which_computation == 'call':
parts = [value.call.function]
if value.call.argument.WhichOneof('computation'):
parts.append(value.call.argument)
parts = await asyncio.gather(
*[self.create_value(x) for x in parts])
return await self.create_call(
parts[0], parts[1] if len(parts) > 1 else None)
elif which_computation == 'tuple':
element_values = await asyncio.gather(
*[self.create_value(x.value) for x in value.tuple.element])
return await self.create_tuple(
anonymous_tuple.AnonymousTuple(
(e.name if e.name else None, v)
for e, v in zip(value.tuple.element, element_values)))
elif which_computation == 'selection':
which_selection = value.selection.WhichOneof('selection')
if which_selection == 'name':
name = value.selection.name
index = None
elif which_selection != 'index':
raise ValueError(
'Unrecognized selection type: "{}".'.format(
which_selection))
else:
index = value.selection.index
name = None
return await self.create_selection(await self.create_value(
value.selection.source),
index=index,
name=name)
else:
raise ValueError(
'Unsupported computation building block of type "{}".'.
format(which_computation))
else:
py_typecheck.check_type(type_spec, computation_types.Type)
if isinstance(type_spec, computation_types.FunctionType):
raise ValueError(
'Encountered a value of a functional TFF type {} and Python type '
'{} that is not of one of the recognized representations.'.
format(type_spec, py_typecheck.type_string(type(value))))
elif isinstance(type_spec, computation_types.FederatedType):
children = self._target_executors.get(type_spec.placement)
self._check_value_compatible_with_placement(
value, type_spec.placement, type_spec.all_equal)
if type_spec.all_equal:
value = [value for _ in children]
child_vals = await asyncio.gather(*[
c.create_value(v, type_spec.member)
for v, c in zip(value, children)
])
return FederatingExecutorValue(child_vals, type_spec)
else:
child = self._target_executors.get(None)
if not child or len(child) > 1:
raise ValueError(
'Executor is not configured for unplaced values.')
else:
return FederatingExecutorValue(
await child[0].create_value(value, type_spec),
type_spec)
def test_construction_from_ordereddict(self):
v = collections.OrderedDict(a=1, b=2, c=3)
x = anonymous_tuple.AnonymousTuple(v.items())
self.assertSequenceEqual(anonymous_tuple.to_elements(x),
list(v.items()))
def test_returns_true(self):
value = anonymous_tuple.AnonymousTuple([('a', 0.0)])
type_spec = computation_types.NamedTupleTypeWithPyContainerType(
[('a', tf.float32)], dict)
self.assertTrue(
type_analysis.is_anon_tuple_with_py_container(value, type_spec))
def test_is_same_structure_check_types(self):
self.assertTrue(
anonymous_tuple.is_same_structure(
anonymous_tuple.AnonymousTuple([('a', 10)]),
anonymous_tuple.AnonymousTuple([('a', 20)])))
self.assertTrue(
anonymous_tuple.is_same_structure(
anonymous_tuple.AnonymousTuple([
('a', 10),
('b', anonymous_tuple.AnonymousTuple([('z', 5)])),
]),
anonymous_tuple.AnonymousTuple([
('a', 20),
('b', anonymous_tuple.AnonymousTuple([('z', 50)])),
])))
self.assertFalse(
anonymous_tuple.is_same_structure(
anonymous_tuple.AnonymousTuple([('x', {
'y': 4
})]), anonymous_tuple.AnonymousTuple([('x', {
'y': 5,
'z': 6
})])))
self.assertTrue(
anonymous_tuple.is_same_structure(
anonymous_tuple.AnonymousTuple([('x', {
'y': 5
})]), anonymous_tuple.AnonymousTuple([('x', {
'y': 6
})])))
with self.assertRaises(TypeError):
anonymous_tuple.is_same_structure(
{'x': 5.0}, # not an AnonymousTuple
anonymous_tuple.AnonymousTuple([('x', 5.0)]))
async def create_value(self, value, type_spec=None):
type_spec = computation_types.to_type(type_spec)
if isinstance(value, intrinsic_defs.IntrinsicDef):
if not type_utils.is_concrete_instance_of(type_spec,
value.type_signature):
raise TypeError(
'Incompatible type {} used with intrinsic {}.'.format(
type_spec, value.uri))
else:
return FederatedExecutorValue(value, type_spec)
if isinstance(value, placement_literals.PlacementLiteral):
if type_spec is not None:
py_typecheck.check_type(type_spec,
computation_types.PlacementType)
return FederatedExecutorValue(value,
computation_types.PlacementType())
elif isinstance(value, computation_impl.ComputationImpl):
return await self.create_value(
computation_impl.ComputationImpl.get_proto(value),
type_utils.reconcile_value_with_type_spec(value, type_spec))
elif isinstance(value, pb.Computation):
if type_spec is None:
type_spec = type_serialization.deserialize_type(value.type)
which_computation = value.WhichOneof('computation')
if which_computation in ['tensorflow', 'lambda']:
return FederatedExecutorValue(value, type_spec)
elif which_computation == 'reference':
raise ValueError(
'Encountered an unexpected unbound references "{}".'.
format(value.reference.name))
elif which_computation == 'intrinsic':
intr = intrinsic_defs.uri_to_intrinsic_def(value.intrinsic.uri)
if intr is None:
raise ValueError(
'Encountered an unrecognized intrinsic "{}".'.format(
value.intrinsic.uri))
py_typecheck.check_type(intr, intrinsic_defs.IntrinsicDef)
return await self.create_value(intr, type_spec)
elif which_computation == 'placement':
return await self.create_value(
placement_literals.uri_to_placement_literal(
value.placement.uri), type_spec)
elif which_computation == 'call':
parts = [value.call.function]
if value.call.argument.WhichOneof('computation'):
parts.append(value.call.argument)
parts = await asyncio.gather(
*[self.create_value(x) for x in parts])
return await self.create_call(
parts[0], parts[1] if len(parts) > 1 else None)
elif which_computation == 'tuple':
element_values = await asyncio.gather(
*[self.create_value(x.value) for x in value.tuple.element])
return await self.create_tuple(
anonymous_tuple.AnonymousTuple([
(e.name if e.name else None, v)
for e, v in zip(value.tuple.element, element_values)
]))
elif which_computation == 'selection':
which_selection = value.selection.WhichOneof('selection')
if which_selection == 'name':
name = value.selection.name
index = None
elif which_selection != 'index':
raise ValueError(
'Unrecognized selection type: "{}".'.format(
which_selection))
else:
index = value.selection.index
name = None
return await self.create_selection(await self.create_value(
value.selection.source),
index=index,
name=name)
else:
raise ValueError(
'Unsupported computation building block of type "{}".'.
format(which_computation))
else:
py_typecheck.check_type(type_spec, computation_types.Type)
if isinstance(type_spec, computation_types.FunctionType):
raise ValueError(
'Encountered a value of a functional TFF type {} and Python type '
'{} that is not of one of the recognized representations.'.
format(type_spec, py_typecheck.type_string(type(value))))
elif isinstance(type_spec, computation_types.FederatedType):
children = self._target_executors.get(type_spec.placement)
if not children:
raise ValueError(
'Placement "{}" is not configured in this executor.'.
format(type_spec.placement))
py_typecheck.check_type(children, list)
if not type_spec.all_equal:
py_typecheck.check_type(value,
(list, tuple, set, frozenset))
if not isinstance(value, list):
value = list(value)
elif isinstance(value, list):
raise ValueError(
'An all_equal value should be passed directly, not as a list.'
)
else:
value = [value for _ in children]
if len(value) != len(children):
raise ValueError(
'Federated value contains {} items, but the placement {} in this '
'executor is configured with {} participants.'.format(
len(value), type_spec.placement, len(children)))
child_vals = await asyncio.gather(*[
c.create_value(v, type_spec.member)
for v, c in zip(value, children)
])
return FederatedExecutorValue(child_vals, type_spec)
else:
child = self._target_executors.get(None)
if not child or len(child) > 1:
raise RuntimeError(
'Executor is not configured for unplaced values.')
else:
return FederatedExecutorValue(
await child[0].create_value(value, type_spec),
type_spec)
def test_map_structure(self):
x = anonymous_tuple.AnonymousTuple([
('a', 10),
('b',
anonymous_tuple.AnonymousTuple([
('x', anonymous_tuple.AnonymousTuple([('p', 40)])),
('y', 30),
('z', anonymous_tuple.AnonymousTuple([('q', 50), ('r', 60)])),
])),
('c', 20),
])
y = anonymous_tuple.AnonymousTuple([
('a', 1),
('b',
anonymous_tuple.AnonymousTuple([
('x', anonymous_tuple.AnonymousTuple([('p', 4)])),
('y', 3),
('z', anonymous_tuple.AnonymousTuple([('q', 5), ('r', 6)])),
])),
('c', 2),
])
self.assertEqual(
anonymous_tuple.map_structure(lambda x, y: x + y, x, y),
anonymous_tuple.AnonymousTuple([
('a', 11),
('b',
anonymous_tuple.AnonymousTuple([
('x', anonymous_tuple.AnonymousTuple([('p', 44)])),
('y', 33),
('z',
anonymous_tuple.AnonymousTuple([('q', 55), ('r', 66)])),
])),
('c', 22),
]))
def assemble_result_from_graph(type_spec, binding, output_map):
"""Assembles a result stamped into a `tf.Graph` given type signature/binding.
This method does roughly the opposite of `capture_result_from_graph`, in that
whereas `capture_result_from_graph` starts with a single structured object
made up of tensors and computes its type and bindings, this method starts
with the type/bindings and constructs a structured object made up of tensors.
Args:
type_spec: The type signature of the result to assemble, an instance of
`types.Type` or something convertible to it.
binding: The binding that relates the type signature to names of tensors in
the graph, an instance of `pb.TensorFlow.Binding`.
output_map: The mapping from tensor names that appear in the binding to
actual stamped tensors (possibly renamed during import).
Returns:
The assembled result, a Python object that is composed of tensors, possibly
nested within Python structures such as anonymous tuples.
Raises:
TypeError: If the argument or any of its parts are of an uexpected type.
ValueError: If the arguments are invalid or inconsistent witch other, e.g.,
the type and binding don't match, or the tensor is not found in the map.
"""
type_spec = computation_types.to_type(type_spec)
py_typecheck.check_type(type_spec, computation_types.Type)
py_typecheck.check_type(binding, pb.TensorFlow.Binding)
py_typecheck.check_type(output_map, dict)
for k, v in six.iteritems(output_map):
py_typecheck.check_type(k, six.string_types)
if not tf.contrib.framework.is_tensor(v):
raise TypeError(
'Element with key {} in the output map is {}, not a tensor.'.
format(k, py_typecheck.type_string(type(v))))
binding_oneof = binding.WhichOneof('binding')
if isinstance(type_spec, computation_types.TensorType):
if binding_oneof != 'tensor':
raise ValueError(
'Expected a tensor binding, found {}.'.format(binding_oneof))
elif binding.tensor.tensor_name not in output_map:
raise ValueError(
'Tensor named {} not found in the output map.'.format(
binding.tensor.tensor_name))
else:
return output_map[binding.tensor.tensor_name]
elif isinstance(type_spec, computation_types.NamedTupleType):
if binding_oneof != 'tuple':
raise ValueError(
'Expected a tuple binding, found {}.'.format(binding_oneof))
else:
type_elements = anonymous_tuple.to_elements(type_spec)
if len(binding.tuple.element) != len(type_elements):
raise ValueError(
'Mismatching tuple sizes in type ({}) and binding ({}).'.
format(len(type_elements), len(binding.tuple.element)))
result_elements = []
for (element_name,
element_type), element_binding in zip(type_elements,
binding.tuple.element):
element_object = assemble_result_from_graph(
element_type, element_binding, output_map)
result_elements.append((element_name, element_object))
return anonymous_tuple.AnonymousTuple(result_elements)
elif isinstance(type_spec, computation_types.SequenceType):
if binding_oneof != 'sequence':
raise ValueError(
'Expected a sequence binding, found {}.'.format(binding_oneof))
else:
handle = output_map[binding.sequence.iterator_string_handle_name]
return make_dataset_from_string_handle(handle, type_spec.element)
else:
raise ValueError('Unsupported type \'{}\'.'.format(str(type_spec)))
def test_construction_from_generator_expression(self):
x = anonymous_tuple.AnonymousTuple(
(name, i) for i, name in enumerate(('a', 'b', None)))
self.assertSequenceEqual(anonymous_tuple.to_elements(x), [('a', 0),
('b', 1),
(None, 2)])
def test_pack_args_into_anonymous_tuple_with_type_spec_expect_success(
self, args, kwargs, type_spec, elements):
self.assertEqual(
function_utils.pack_args_into_anonymous_tuple(
args, kwargs, type_spec, NoopIngestContextForTest()),
anonymous_tuple.AnonymousTuple(elements))
def test_from_container_with_anonymous_tuple(self):
x = anonymous_tuple.from_container(
anonymous_tuple.AnonymousTuple([('a', 10), ('b', 20)]))
self.assertIs(x, x)
async def _compute_tuple(anon_tuple):
elements = anonymous_tuple.to_elements(anon_tuple)
keys = [k for k, _ in elements]
vals = await asyncio.gather(
*[_compute_element(v) for _, v in elements])
return anonymous_tuple.AnonymousTuple(zip(keys, vals))
def test_construction_from_iter_elements(self):
x = anonymous_tuple.AnonymousTuple((('a', 1), ('b', 2), (None, 3)))
self.assertSequenceEqual(
anonymous_tuple.AnonymousTuple(anonymous_tuple.iter_elements(x)),
x)
def pack_args_into_anonymous_tuple(args, kwargs, type_spec=None, context=None):
"""Packs positional and keyword arguments into an anonymous tuple.
If 'type_spec' is not None, it must be a tuple type or something that's
convertible to it by computation_types.to_type(). The assignment of arguments
to fields of the tuple follows the same rule as during function calls. If
'type_spec' is None, the positional arguments precede any of the keyword
arguments, and the ordering of the keyword arguments matches the ordering in
which they appear in kwargs. If the latter is an OrderedDict, the ordering
will be preserved. On the other hand, if the latter is an ordinary unordered
dict, the ordering is arbitrary.
Args:
args: Positional arguments.
kwargs: Keyword arguments.
type_spec: The optional type specification (either an instance of
computation_types.NamedTupleType or something convertible to it), or None
if there's no type. Used to drive the arrangements of args into fields of
the constructed anonymous tuple, as noted in the description.
context: The optional context (an instance of `context_base.Context`) in
which the arguments are being packed. Required if and only if the
`type_spec` is not `None`.
Returns:
An anoymous tuple containing all the arguments.
Raises:
TypeError: if the arguments are of the wrong computation_types.
"""
type_spec = computation_types.to_type(type_spec)
if not type_spec:
return anonymous_tuple.AnonymousTuple([(None, arg) for arg in args] +
list(six.iteritems(kwargs)))
else:
py_typecheck.check_type(type_spec, computation_types.NamedTupleType)
py_typecheck.check_type(context, context_base.Context)
context = context # type: context_base.Context
if not is_argument_tuple(type_spec): # pylint: disable=attribute-error
raise TypeError(
'Parameter type {} does not have a structure of an argument tuple, '
'and cannot be populated from multiple positional and keyword '
'arguments'.format(type_spec))
else:
result_elements = []
positions_used = set()
keywords_used = set()
for index, (name, elem_type) in enumerate(
anonymous_tuple.to_elements(type_spec)):
if index < len(args):
if name is not None and name in kwargs:
raise TypeError('Argument {} specified twice.'.format(name))
else:
arg_value = args[index]
result_elements.append((name, context.ingest(arg_value, elem_type)))
positions_used.add(index)
elif name is not None and name in kwargs:
arg_value = kwargs[name]
result_elements.append((name, context.ingest(arg_value, elem_type)))
keywords_used.add(name)
elif name:
raise TypeError('Argument named {} is missing.'.format(name))
else:
raise TypeError('Argument at position {} is missing.'.format(index))
positions_missing = set(range(len(args))).difference(positions_used)
if positions_missing:
raise TypeError(
'Positional arguments at {} not used.'.format(positions_missing))
keywords_missing = set(kwargs.keys()).difference(keywords_used)
if keywords_missing:
raise TypeError(
'Keyword arguments at {} not used.'.format(keywords_missing))
return anonymous_tuple.AnonymousTuple(result_elements)
def to_representation_for_type(value,
tf_function_cache,
type_spec=None,
device=None):
"""Verifies or converts the `value` to an eager object matching `type_spec`.
WARNING: This function is only partially implemented. It does not support
data sets at this point.
The output of this function is always an eager tensor, eager dataset, a
representation of a TensorFlow computation, or a nested structure of those
that matches `type_spec`, and when `device` has been specified, everything
is placed on that device on a best-effort basis.
TensorFlow computations are represented here as zero- or one-argument Python
callables that accept their entire argument bundle as a single Python object.
Args:
value: The raw representation of a value to compare against `type_spec` and
potentially to be converted.
tf_function_cache: A cache obeying `dict` semantics that can be used to look
up previously embedded TensorFlow functions.
type_spec: An instance of `tff.Type`, can be `None` for values that derive
from `typed_object.TypedObject`.
device: The optional device to place the value on (for tensor-level values).
Returns:
Either `value` itself, or a modified version of it.
Raises:
TypeError: If the `value` is not compatible with `type_spec`.
"""
type_spec = type_utils.reconcile_value_with_type_spec(value, type_spec)
if isinstance(value, computation_base.Computation):
return to_representation_for_type(
computation_impl.ComputationImpl.get_proto(value),
tf_function_cache, type_spec, device)
elif isinstance(value, pb.Computation):
key = (value.SerializeToString(), str(type_spec), device)
cached_fn = tf_function_cache.get(key)
if cached_fn:
return cached_fn
embedded_fn = embed_tensorflow_computation(value, type_spec, device)
tf_function_cache[key] = embedded_fn
return embedded_fn
elif isinstance(type_spec, computation_types.NamedTupleType):
type_elem = anonymous_tuple.to_elements(type_spec)
value_elem = (anonymous_tuple.to_elements(
anonymous_tuple.from_container(value)))
result_elem = []
if len(type_elem) != len(value_elem):
raise TypeError(
'Expected a {}-element tuple, found {} elements.'.format(
len(type_elem), len(value_elem)))
for (t_name, el_type), (v_name, el_val) in zip(type_elem, value_elem):
if t_name != v_name:
raise TypeError(
'Mismatching element names in type vs. value: {} vs. {}.'.
format(t_name, v_name))
el_repr = to_representation_for_type(el_val, tf_function_cache,
el_type, device)
result_elem.append((t_name, el_repr))
return anonymous_tuple.AnonymousTuple(result_elem)
elif device is not None:
py_typecheck.check_type(device, str)
with tf.device(device):
return to_representation_for_type(value,
tf_function_cache,
type_spec=type_spec,
device=None)
elif isinstance(value, EagerValue):
return value.internal_representation
elif isinstance(value, executor_value_base.ExecutorValue):
raise TypeError(
'Cannot accept a value embedded within a non-eager executor.')
elif isinstance(type_spec, computation_types.TensorType):
if not tf.is_tensor(value):
value = tf.convert_to_tensor(value, dtype=type_spec.dtype)
elif hasattr(value, 'read_value'):
# a tf.Variable-like result, get a proper tensor.
value = value.read_value()
value_type = (computation_types.TensorType(value.dtype.base_dtype,
value.shape))
if not type_analysis.is_assignable_from(type_spec, value_type):
raise TypeError(
'The apparent type {} of a tensor {} does not match the expected '
'type {}.'.format(value_type, value, type_spec))
return value
elif isinstance(type_spec, computation_types.SequenceType):
if isinstance(value, list):
value = tensorflow_utils.make_data_set_from_elements(
None, value, type_spec.element)
py_typecheck.check_type(
value, type_conversions.TF_DATASET_REPRESENTATION_TYPES)
element_type = computation_types.to_type(value.element_spec)
value_type = computation_types.SequenceType(element_type)
type_analysis.check_assignable_from(type_spec, value_type)
return value
else:
raise TypeError('Unexpected type {}.'.format(type_spec))
def test_list_structures_from_element_type_spec_with_empty_anon_tuples(self):
self._test_list_structure(
computation_types.NamedTupleType([]), [
anonymous_tuple.AnonymousTuple([]),
anonymous_tuple.AnonymousTuple([])
], 'OrderedDict()')
def bar(x):
arg = anonymous_tuple.AnonymousTuple(
_make_test_tuple(x, k) for k in range(n))
val = intrinsics.federated_zip(arg)
self.assertIsInstance(val, value_base.Value)
return val
