def test_with_function_type(self):
t1 = computation_types.FunctionType(tf.int32, tf.bool)
t2 = computation_types.FunctionType(tf.int32, tf.bool)
t3 = computation_types.FunctionType(tf.int32, tf.int32)
t4 = computation_types.TensorType(tf.int32)
self.assertTrue(type_analysis.is_assignable_from(t1, t1))
self.assertTrue(type_analysis.is_assignable_from(t1, t2))
self.assertFalse(type_analysis.is_assignable_from(t1, t3))
self.assertFalse(type_analysis.is_assignable_from(t1, t4))
def test_with_sequence_type(self):
self.assertTrue(
type_analysis.is_assignable_from(
computation_types.SequenceType(tf.int32),
computation_types.SequenceType(tf.int32)))
self.assertFalse(
type_analysis.is_assignable_from(
computation_types.SequenceType(tf.int32),
computation_types.SequenceType(tf.bool)))
def __init__(self, initialize_fn, next_fn):
"""Creates a `tff.templates.IterativeProcess`.
Args:
initialize_fn: a no-arg `tff.Computation` that creates the initial state
of the chained computation.
next_fn: a `tff.Computation` that defines an iterated function. If
`initialize_fn` returns a type _T_, then `next_fn` must return a type
_U_ which is compatible with _T_ or multiple values where the first type
is _U_, and accept either a single argument of type _U_ or multiple
arguments where the first argument must be of type _U_.
Raises:
TypeError: `initialize_fn` and `next_fn` are not compatible function
types.
"""
py_typecheck.check_type(initialize_fn, computation_base.Computation)
if initialize_fn.type_signature.parameter is not None:
raise TypeError(
'initialize_fn must be a no-arg tff.Computation, but found parameter '
'{}'.format(initialize_fn.type_signature))
initialize_result_type = initialize_fn.type_signature.result
py_typecheck.check_type(next_fn, computation_base.Computation)
if isinstance(next_fn.type_signature.parameter,
computation_types.NamedTupleType):
next_first_param_type = next_fn.type_signature.parameter[0]
else:
next_first_param_type = next_fn.type_signature.parameter
if not type_analysis.is_assignable_from(next_first_param_type,
initialize_result_type):
raise TypeError(
'The return type of initialize_fn must be assignable '
'to the first parameter of next_fn, but found\n'
'initialize_fn.type_signature.result=\n{}\n'
'next_fn.type_signature.parameter[0]=\n{}'.format(
initialize_result_type, next_first_param_type))
next_result_type = next_fn.type_signature.result
if not type_analysis.is_assignable_from(next_first_param_type,
next_result_type):
# This might be multiple output next_fn, check if the first argument might
# be the state. If still not the right type, raise an error.
if isinstance(next_result_type, computation_types.NamedTupleType):
next_result_type = next_result_type[0]
if next_first_param_type != next_result_type:
raise TypeError(
'The return type of next_fn must be assignable to the '
'first parameter, but found\n'
'next_fn.type_signature.parameter[0]=\n{}\n'
'actual next_result_type=\n{}'.format(
next_first_param_type, next_result_type))
self._initialize_fn = initialize_fn
self._next_fn = next_fn
def _unpack_and_call(fn, arg_types, kwarg_types, arg):
"""An interceptor function that unpacks 'arg' before calling `fn`.
The function verifies the actual parameters before it forwards the
call as a last-minute check.
Args:
fn: The function or defun to invoke.
arg_types: The list of positional argument types (guaranteed to all be
instances of computation_types.Types).
kwarg_types: The dictionary of keyword argument types (guaranteed to
all be instances of computation_types.Types).
arg: The argument to unpack.
Returns:
The result of invoking `fn` on the unpacked arguments.
Raises:
TypeError: if types don't match.
"""
py_typecheck.check_type(
arg, (anonymous_tuple.AnonymousTuple, value_base.Value))
args = []
for idx, expected_type in enumerate(arg_types):
element_value = arg[idx]
actual_type = type_conversions.infer_type(element_value)
if not type_analysis.is_assignable_from(
expected_type, actual_type):
raise TypeError(
'Expected element at position {} to be of type {}, found {}.'
.format(idx, expected_type, actual_type))
if isinstance(element_value,
anonymous_tuple.AnonymousTuple):
element_value = type_conversions.type_to_py_container(
element_value, expected_type)
args.append(element_value)
kwargs = {}
for name, expected_type in kwarg_types.items():
element_value = getattr(arg, name)
actual_type = type_conversions.infer_type(element_value)
if not type_analysis.is_assignable_from(
expected_type, actual_type):
raise TypeError(
'Expected element named {} to be of type {}, found {}.'
.format(name, expected_type, actual_type))
if type_analysis.is_anon_tuple_with_py_container(
element_value, expected_type):
element_value = type_conversions.type_to_py_container(
element_value, expected_type)
kwargs[name] = element_value
return fn(*args, **kwargs)
def test_with_tensor_type_and_tensor_type(self):
t = computation_types.TensorType(tf.int32, [10])
self.assertFalse(
type_analysis.is_assignable_from(
t, computation_types.TensorType(tf.int32)))
self.assertFalse(
type_analysis.is_assignable_from(
t, computation_types.TensorType(tf.int32, [5])))
self.assertFalse(
type_analysis.is_assignable_from(
t, computation_types.TensorType(tf.int32, [10, 10])))
self.assertTrue(
type_analysis.is_assignable_from(
t, computation_types.TensorType(tf.int32, 10)))
def create_chained_calls(functions, arg):
r"""Creates a chain of `n` calls.
Call
/ \
Comp ...
\
Call
/ \
Comp Comp
The first functional computation in `functions` must have a parameter type
that is assignable from the type of `arg`, each other functional computation
in `functions` must have a parameter type that is assignable from the previous
functional computations result type.
Args:
functions: A Python list of functional computations.
arg: A `building_blocks.ComputationBuildingBlock`.
Returns:
A `building_blocks.Call`.
"""
for fn in functions:
if not type_analysis.is_assignable_from(fn.parameter_type,
arg.type_signature):
raise TypeError(
'The parameter of the function is of type {}, and the argument is of '
'an incompatible type {}.'.format(str(fn.parameter_type),
str(arg.type_signature)))
call = building_blocks.Call(fn, arg)
arg = call
return call
async def create_call(self, comp, arg=None):
py_typecheck.check_type(comp, ReferenceResolvingExecutorValue)
py_typecheck.check_type(comp.type_signature, computation_types.FunctionType)
param_type = comp.type_signature.parameter
if param_type is None:
py_typecheck.check_none(arg)
else:
py_typecheck.check_type(arg, ReferenceResolvingExecutorValue)
arg_type = arg.type_signature # pytype: disable=attribute-error
if not type_analysis.is_assignable_from(param_type, arg_type):
raise TypeError('ReferenceResolvingExecutor asked to create call with '
'incompatible type specifications. Function '
'takes an argument of type {}, but was supplied '
'an argument of type {}'.format(param_type, arg_type))
comp_repr = comp.internal_representation
if isinstance(comp_repr, executor_value_base.ExecutorValue):
# `comp` represents a function in the target executor, so we convert the
# argument to a value inside the target executor and `create_call` on
# the target executor.
delegated_arg = await self._embed_value_in_target_exec(
arg) if arg is not None else None
return ReferenceResolvingExecutorValue(await
self._target_executor.create_call(
comp_repr, delegated_arg))
elif isinstance(comp_repr, ScopedLambda):
return await comp_repr.invoke(self, arg)
else:
raise TypeError(
'Unexpected type to ReferenceResolvingExecutor create_call: {}'
.format(type(comp_repr)))
def federated_map_all_equal(self, fn, arg):
"""`federated_map` with the `all_equal` bit set in the `arg` and return."""
# 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)
arg = value_utils.ensure_federated_value(arg,
placement_literals.CLIENTS,
'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_analysis.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 = building_block_factory.create_federated_map_all_equal(fn, arg)
return value_impl.ValueImpl(comp, self._context_stack)
def _call(fn, parameter_type, arg):
arg_type = type_conversions.infer_type(arg)
if not type_analysis.is_assignable_from(
parameter_type, arg_type):
raise TypeError(
'Expected an argument of type {}, found {}.'.format(
parameter_type, arg_type))
if type_analysis.is_anon_tuple_with_py_container(
arg, parameter_type):
arg = type_conversions.type_to_py_container(
arg, parameter_type)
return fn(arg)
def _create_chained_dummy_federated_maps(functions, arg):
py_typecheck.check_type(arg, building_blocks.ComputationBuildingBlock)
for fn in functions:
py_typecheck.check_type(fn, building_blocks.ComputationBuildingBlock)
if not type_analysis.is_assignable_from(fn.parameter_type,
arg.type_signature.member):
raise TypeError(
'The parameter of the function is of type {}, and the argument is of '
'an incompatible type {}.'.format(
str(fn.parameter_type), str(arg.type_signature.member)))
call = building_block_factory.create_federated_map(fn, arg)
arg = call
return call
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_analysis.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 placement_literals.SERVER:
if not arg.type_signature.all_equal:
raise TypeError(
'Arguments placed at {} should be equal at all locations.'.
format(placement_literals.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 placement_literals.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(placement_literals.SERVER, placement_literals.CLIENTS,
arg.type_signature.placement))
return value_impl.ValueImpl(comp, self._context_stack)
def sequence_reduce(self, value, zero, op):
"""Implements `sequence_reduce` as defined in `api/intrinsics.py`."""
value = value_impl.to_value(value, None, self._context_stack)
zero = value_impl.to_value(zero, None, self._context_stack)
op = value_impl.to_value(op, None, self._context_stack)
if isinstance(value.type_signature, computation_types.SequenceType):
element_type = value.type_signature.element
else:
py_typecheck.check_type(value.type_signature,
computation_types.FederatedType)
py_typecheck.check_type(value.type_signature.member,
computation_types.SequenceType)
element_type = value.type_signature.member.element
op_type_expected = type_factory.reduction_op(zero.type_signature,
element_type)
if not type_analysis.is_assignable_from(op_type_expected,
op.type_signature):
raise TypeError('Expected an operator of type {}, got {}.'.format(
op_type_expected, op.type_signature))
value = value_impl.ValueImpl.get_comp(value)
zero = value_impl.ValueImpl.get_comp(zero)
op = value_impl.ValueImpl.get_comp(op)
if isinstance(value.type_signature, computation_types.SequenceType):
return value_impl.ValueImpl(
building_block_factory.create_sequence_reduce(value, zero, op),
self._context_stack)
else:
value_type = computation_types.SequenceType(element_type)
intrinsic_type = computation_types.FunctionType((
value_type,
zero.type_signature,
op.type_signature,
), op.type_signature.result)
intrinsic = building_blocks.Intrinsic(
intrinsic_defs.SEQUENCE_REDUCE.uri, intrinsic_type)
ref = building_blocks.Reference('arg', value_type)
tup = building_blocks.Tuple((ref, zero, op))
call = building_blocks.Call(intrinsic, tup)
fn = building_blocks.Lambda(ref.name, ref.type_signature, call)
fn_impl = value_impl.ValueImpl(fn, self._context_stack)
if value.type_signature.placement in [
placement_literals.SERVER, placement_literals.CLIENTS
]:
return self.federated_map(fn_impl, value)
else:
raise TypeError('Unsupported placement {}.'.format(
value.type_signature.placement))
def is_signature_compatible_with_types(signature: inspect.Signature, *args,
**kwargs) -> bool:
"""Determines if functions matching signature accept `args` and `kwargs`.
Args:
signature: An instance of `inspect.Signature` to verify agains the
arguments.
*args: Zero or more positional arguments, all of which must be instances of
computation_types.Type or something convertible to it by
computation_types.to_type().
**kwargs: Zero or more keyword arguments, all of which must be instances of
computation_types.Type or something convertible to it by
computation_types.to_type().
Returns:
`True` or `False`, depending on the outcome of the test.
Raises:
TypeError: if the arguments are of the wrong computation_types.
"""
try:
bound_args = signature.bind(*args, **kwargs)
except TypeError:
return False
# If we have no defaults then `bind` will have raised `TypeError` if the
# signature was not compatible with *args and **kwargs.
if all(p.default is inspect.Parameter.empty
for p in signature.parameters.values()):
return True
# Otherwise we need to check the defaults against the types that were given to
# ensure they are compatible.
for p in signature.parameters.values():
if p.default is inspect.Parameter.empty or p.default is None:
# No default value or optional.
continue
arg_value = bound_args.arguments.get(p.name, p.default)
if arg_value is p.default:
continue
arg_type = computation_types.to_type(arg_value)
default_type = type_conversions.infer_type(p.default)
if not type_analysis.is_assignable_from(arg_type, default_type):
return False
return True
def test_with_named_tuple_type(self):
t1 = computation_types.NamedTupleType([tf.int32, ('a', tf.bool)])
t2 = computation_types.NamedTupleType([tf.int32, ('a', tf.bool)])
t3 = computation_types.NamedTupleType([tf.int32, ('b', tf.bool)])
t4 = computation_types.NamedTupleType([tf.int32, ('a', tf.string)])
t5 = computation_types.NamedTupleType([tf.int32])
t6 = computation_types.NamedTupleType([tf.int32, tf.bool])
self.assertTrue(type_analysis.is_assignable_from(t1, t2))
self.assertFalse(type_analysis.is_assignable_from(t1, t3))
self.assertFalse(type_analysis.is_assignable_from(t1, t4))
self.assertFalse(type_analysis.is_assignable_from(t1, t5))
self.assertTrue(type_analysis.is_assignable_from(t1, t6))
self.assertFalse(type_analysis.is_assignable_from(t6, t1))
def _wrap_sequence_as_value(elements, element_type, context_stack):
"""Wraps `elements` as a TFF sequence with elements of type `element_type`.
Args:
elements: Python object to the wrapped as a TFF sequence value.
element_type: An instance of `Type` that determines the type of elements of
the sequence.
context_stack: The context stack to use.
Returns:
An instance of `tff.Value`.
Raises:
TypeError: If `elements` and `element_type` are of incompatible types.
"""
# TODO(b/113116813): Add support for other representations of sequences.
py_typecheck.check_type(elements, list)
py_typecheck.check_type(context_stack, context_stack_base.ContextStack)
# Checks that the types of all the individual elements are compatible with the
# requested type of the sequence as a while.
for elem in elements:
elem_type = type_conversions.infer_type(elem)
if not type_analysis.is_assignable_from(element_type, elem_type):
raise TypeError(
'Expected all sequence elements to be {}, found {}.'.format(
element_type, elem_type))
# Defines a no-arg function that builds a `tf.data.Dataset` from the elements.
def _create_dataset_from_elements():
return tensorflow_utils.make_data_set_from_elements(
tf.compat.v1.get_default_graph(), elements, element_type)
# Wraps the dataset as a value backed by a no-argument TensorFlow computation.
tf_comp, _ = tensorflow_serialization.serialize_py_fn_as_tf_computation(
_create_dataset_from_elements, None, context_stack)
return ValueImpl(
building_blocks.Call(building_blocks.CompiledComputation(tf_comp)),
context_stack)
def federated_reduce(self, value, zero, op):
"""Implements `federated_reduce` as defined in `api/intrinsics.py`."""
# TODO(b/113112108): Since in most cases, it can be assumed that CLIENTS is
# a non-empty collective (or else, the computation fails), specifying zero
# at this level of the API should probably be optional. TBD.
value = value_impl.to_value(value, None, self._context_stack)
value = value_utils.ensure_federated_value(value,
placement_literals.CLIENTS,
'value to be reduced')
zero = value_impl.to_value(zero, None, self._context_stack)
py_typecheck.check_type(zero, value_base.Value)
# TODO(b/113112108): We need a check here that zero does not have federated
# constituents.
op = value_impl.to_value(
op,
None,
self._context_stack,
parameter_type_hint=computation_types.NamedTupleType(
[zero.type_signature, value.type_signature.member]))
py_typecheck.check_type(op, value_base.Value)
py_typecheck.check_type(op.type_signature,
computation_types.FunctionType)
op_type_expected = type_factory.reduction_op(
zero.type_signature, value.type_signature.member)
if not type_analysis.is_assignable_from(op_type_expected,
op.type_signature):
raise TypeError('Expected an operator of type {}, got {}.'.format(
op_type_expected, op.type_signature))
value = value_impl.ValueImpl.get_comp(value)
zero = value_impl.ValueImpl.get_comp(zero)
op = value_impl.ValueImpl.get_comp(op)
comp = building_block_factory.create_federated_reduce(value, zero, op)
return value_impl.ValueImpl(comp, self._context_stack)
def test_with_federated_type(self):
t1 = computation_types.FederatedType(tf.int32,
placement_literals.CLIENTS)
self.assertTrue(type_analysis.is_assignable_from(t1, t1))
t2 = computation_types.FederatedType(tf.int32,
placement_literals.CLIENTS,
all_equal=True)
self.assertTrue(type_analysis.is_assignable_from(t1, t2))
self.assertTrue(type_analysis.is_assignable_from(t2, t2))
self.assertFalse(type_analysis.is_assignable_from(t2, t1))
t3 = computation_types.FederatedType(
computation_types.TensorType(tf.int32, [10]),
placement_literals.CLIENTS)
t4 = computation_types.FederatedType(
computation_types.TensorType(tf.int32, [None]),
placement_literals.CLIENTS)
self.assertTrue(type_analysis.is_assignable_from(t4, t3))
self.assertFalse(type_analysis.is_assignable_from(t3, t4))
t5 = computation_types.FederatedType(
computation_types.TensorType(tf.int32, [10]),
placement_literals.SERVER)
self.assertFalse(type_analysis.is_assignable_from(t3, t5))
self.assertFalse(type_analysis.is_assignable_from(t5, t3))
t6 = computation_types.FederatedType(computation_types.TensorType(
tf.int32, [10]),
placement_literals.CLIENTS,
all_equal=True)
self.assertTrue(type_analysis.is_assignable_from(t3, t6))
self.assertTrue(type_analysis.is_assignable_from(t4, t6))
self.assertFalse(type_analysis.is_assignable_from(t6, t3))
self.assertFalse(type_analysis.is_assignable_from(t6, t4))
def test_with_placement_type(self):
t1 = computation_types.PlacementType()
t2 = computation_types.PlacementType()
self.assertTrue(type_analysis.is_assignable_from(t1, t1))
self.assertTrue(type_analysis.is_assignable_from(t1, t2))
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 to_value(
arg: Any,
type_spec,
context_stack: context_stack_base.ContextStack,
parameter_type_hint=None,
) -> ValueImpl:
"""Converts the argument into an instance of `tff.Value`.
The types of non-`tff.Value` arguments that are currently convertible to
`tff.Value` include the following:
* Lists, tuples, anonymous tuples, named tuples, and dictionaries, all
of which are converted into instances of `tff.Tuple`.
* Placement literals, converted into instances of `tff.Placement`.
* Computations.
* Python constants of type `str`, `int`, `float`, `bool`
* Numpy objects inherting from `np.ndarray` or `np.generic` (the parent
of numpy scalar types)
Args:
arg: Either an instance of `tff.Value`, or an argument convertible to
`tff.Value`. The argument must not be `None`.
type_spec: An optional `computation_types.Type` or value convertible to it
by `computation_types.to_type` which specifies the desired type signature
of the resulting value. This allows for disambiguating the target type
(e.g., when two TFF types can be mapped to the same Python
representations), or `None` if none available, in which case TFF tries to
determine the type of the TFF value automatically.
context_stack: The context stack to use.
parameter_type_hint: An optional `computation_types.Type` or value
convertible to it by `computation_types.to_type` which specifies an
argument type to use in the case that `arg` is a
`function_utils.PolymorphicFunction`.
Returns:
An instance of `tff.Value` corresponding to the given `arg`, and of TFF type
matching the `type_spec` if specified (not `None`).
Raises:
TypeError: if `arg` is of an unsupported type, or of a type that does not
match `type_spec`. Raises explicit error message if TensorFlow constructs
are encountered, as TensorFlow code should be sealed away from TFF
federated context.
"""
py_typecheck.check_type(context_stack, context_stack_base.ContextStack)
if type_spec is not None:
type_spec = computation_types.to_type(type_spec)
type_analysis.check_well_formed(type_spec)
if isinstance(arg, ValueImpl):
result = arg
elif isinstance(arg, building_blocks.ComputationBuildingBlock):
result = ValueImpl(arg, context_stack)
elif isinstance(arg, placement_literals.PlacementLiteral):
result = ValueImpl(building_blocks.Placement(arg), context_stack)
elif isinstance(
arg,
(computation_base.Computation, function_utils.PolymorphicFunction)):
if isinstance(arg, function_utils.PolymorphicFunction):
if parameter_type_hint is None:
raise TypeError(
'Polymorphic computations cannot be converted to TFF values '
'without a type hint. Consider explicitly specifying the '
'argument types of a computation before passing it to a '
'function that requires a TFF value (such as a TFF intrinsic '
'like `federated_map`). If you are a TFF developer and think '
'this should be supported, consider providing `parameter_type_hint` '
'as an argument to the encompassing `to_value` conversion.'
)
parameter_type_hint = computation_types.to_type(
parameter_type_hint)
type_analysis.check_well_formed(parameter_type_hint)
arg = arg.fn_for_argument_type(parameter_type_hint)
py_typecheck.check_type(arg, computation_base.Computation)
result = ValueImpl(
building_blocks.CompiledComputation(
computation_impl.ComputationImpl.get_proto(arg)),
context_stack)
elif type_spec is not None and isinstance(type_spec,
computation_types.SequenceType):
result = _wrap_sequence_as_value(arg, type_spec.element, context_stack)
elif isinstance(arg, anonymous_tuple.AnonymousTuple):
result = ValueImpl(
building_blocks.Tuple([
(k, ValueImpl.get_comp(to_value(v, None, context_stack)))
for k, v in anonymous_tuple.iter_elements(arg)
]), context_stack)
elif py_typecheck.is_named_tuple(arg):
result = to_value(arg._asdict(), None, context_stack) # pytype: disable=attribute-error
elif py_typecheck.is_attrs(arg):
result = to_value(
attr.asdict(arg,
dict_factory=collections.OrderedDict,
recurse=False), None, context_stack)
elif isinstance(arg, dict):
if isinstance(arg, collections.OrderedDict):
items = arg.items()
else:
items = sorted(arg.items())
value = building_blocks.Tuple([
(k, ValueImpl.get_comp(to_value(v, None, context_stack)))
for k, v in items
])
result = ValueImpl(value, context_stack)
elif isinstance(arg, (tuple, list)):
result = ValueImpl(
building_blocks.Tuple([
ValueImpl.get_comp(to_value(x, None, context_stack))
for x in arg
]), context_stack)
elif isinstance(arg, tensorflow_utils.TENSOR_REPRESENTATION_TYPES):
result = _wrap_constant_as_value(arg, context_stack)
elif isinstance(arg, (tf.Tensor, tf.Variable)):
raise TypeError(
'TensorFlow construct {} has been encountered in a federated '
'context. TFF does not support mixing TF and federated orchestration '
'code. Please wrap any TensorFlow constructs with '
'`tff.tf_computation`.'.format(arg))
else:
raise TypeError(
'Unable to interpret an argument of type {} as a TFF value.'.
format(py_typecheck.type_string(type(arg))))
py_typecheck.check_type(result, ValueImpl)
if (type_spec is not None and not type_analysis.is_assignable_from(
type_spec, result.type_signature)):
raise TypeError(
'The supplied argument maps to TFF type {}, which is incompatible with '
'the requested type {}.'.format(result.type_signature, type_spec))
return result
def federated_aggregate(self, value, zero, accumulate, merge, report):
"""Implements `federated_aggregate` as defined in `api/intrinsics.py`."""
value = value_impl.to_value(value, None, self._context_stack)
value = value_utils.ensure_federated_value(value,
placement_literals.CLIENTS,
'value to be aggregated')
zero = value_impl.to_value(zero, None, self._context_stack)
py_typecheck.check_type(zero, value_base.Value)
accumulate = value_impl.to_value(
accumulate,
None,
self._context_stack,
parameter_type_hint=computation_types.NamedTupleType(
[zero.type_signature, value.type_signature.member]))
merge = value_impl.to_value(
merge,
None,
self._context_stack,
parameter_type_hint=computation_types.NamedTupleType(
[zero.type_signature, zero.type_signature]))
report = value_impl.to_value(report,
None,
self._context_stack,
parameter_type_hint=zero.type_signature)
for op in [accumulate, merge, report]:
py_typecheck.check_type(op, value_base.Value)
py_typecheck.check_type(op.type_signature,
computation_types.FunctionType)
if not type_analysis.is_assignable_from(
accumulate.type_signature.parameter[0], zero.type_signature):
raise TypeError('Expected `zero` to be assignable to type {}, '
'but was of incompatible type {}.'.format(
accumulate.type_signature.parameter[0],
zero.type_signature))
accumulate_type_expected = type_factory.reduction_op(
accumulate.type_signature.result, value.type_signature.member)
merge_type_expected = type_factory.reduction_op(
accumulate.type_signature.result, accumulate.type_signature.result)
report_type_expected = computation_types.FunctionType(
merge.type_signature.result, report.type_signature.result)
for op_name, op, type_expected in [
('accumulate', accumulate, accumulate_type_expected),
('merge', merge, merge_type_expected),
('report', report, report_type_expected)
]:
if not type_analysis.is_assignable_from(type_expected,
op.type_signature):
raise TypeError(
'Expected parameter `{}` to be of type {}, but received {} instead.'
.format(op_name, type_expected, op.type_signature))
value = value_impl.ValueImpl.get_comp(value)
zero = value_impl.ValueImpl.get_comp(zero)
accumulate = value_impl.ValueImpl.get_comp(accumulate)
merge = value_impl.ValueImpl.get_comp(merge)
report = value_impl.ValueImpl.get_comp(report)
comp = building_block_factory.create_federated_aggregate(
value, zero, accumulate, merge, report)
return value_impl.ValueImpl(comp, self._context_stack)
def test_with_tensor_type_with_undefined_dims(self):
t1 = computation_types.TensorType(tf.int32, [None])
t2 = computation_types.TensorType(tf.int32, [10])
self.assertTrue(type_analysis.is_assignable_from(t1, t2))
self.assertFalse(type_analysis.is_assignable_from(t2, t1))
def serialize_value(value, type_spec=None):
"""Serializes a value into `executor_pb2.Value`.
Args:
value: A value to be serialized.
type_spec: Optional type spec, a `tff.Type` or something convertible to it.
Returns:
A tuple `(value_proto, ret_type_spec)` where `value_proto` is an instance
of `executor_pb2.Value` with the serialized content of `value`, and the
returned `ret_type_spec` is an instance of `tff.Type` that represents the
TFF type of the serialized value.
Raises:
TypeError: If the arguments are of the wrong types.
ValueError: If the value is malformed.
"""
type_spec = computation_types.to_type(type_spec)
if isinstance(value, computation_pb2.Computation):
type_spec = type_utils.reconcile_value_type_with_type_spec(
type_serialization.deserialize_type(value.type), type_spec)
return executor_pb2.Value(computation=value), type_spec
elif isinstance(value, computation_impl.ComputationImpl):
return serialize_value(
computation_impl.ComputationImpl.get_proto(value),
type_utils.reconcile_value_with_type_spec(value, type_spec))
elif isinstance(type_spec, computation_types.TensorType):
return serialize_tensor_value(value, type_spec)
elif isinstance(type_spec, computation_types.NamedTupleType):
type_elements = anonymous_tuple.to_elements(type_spec)
val_elements = anonymous_tuple.to_elements(
anonymous_tuple.from_container(value))
tup_elems = []
for (e_name, e_type), (_, e_val) in zip(type_elements, val_elements):
e_proto, _ = serialize_value(e_val, e_type)
tup_elems.append(
executor_pb2.Value.Tuple.Element(
name=e_name if e_name else None, value=e_proto))
result_proto = (executor_pb2.Value(tuple=executor_pb2.Value.Tuple(
element=tup_elems)))
return result_proto, type_spec
elif isinstance(type_spec, computation_types.SequenceType):
if not isinstance(value,
type_conversions.TF_DATASET_REPRESENTATION_TYPES):
raise TypeError(
'Cannot serialize Python type {!s} as TFF type {!s}.'.format(
py_typecheck.type_string(type(value)),
type_spec if type_spec is not None else 'unknown'))
value_type = computation_types.SequenceType(
computation_types.to_type(value.element_spec))
if not type_analysis.is_assignable_from(type_spec, value_type):
raise TypeError(
'Cannot serialize dataset with elements of type {!s} as TFF type {!s}.'
.format(value_type,
type_spec if type_spec is not None else 'unknown'))
return serialize_sequence_value(value), type_spec
elif isinstance(type_spec, computation_types.FederatedType):
if type_spec.all_equal:
value = [value]
else:
py_typecheck.check_type(value, list)
items = []
for v in value:
it, it_type = serialize_value(v, type_spec.member)
type_analysis.check_assignable_from(type_spec.member, it_type)
items.append(it)
result_proto = executor_pb2.Value(
federated=executor_pb2.Value.Federated(
type=type_serialization.serialize_type(type_spec).federated,
value=items))
return result_proto, type_spec
else:
raise ValueError(
'Unable to serialize value with Python type {} and {} TFF type.'.
format(str(py_typecheck.type_string(type(value))),
str(type_spec) if type_spec is not None else 'unknown'))
def deserialize_and_call_tf_computation(computation_proto, arg, graph):
"""Deserializes a TF computation and inserts it into `graph`.
This method performs an action that can be considered roughly the opposite of
what `tensorflow_serialization.serialize_py_fn_as_tf_computation` does. At
the moment, it simply imports the graph in the current context. A future
implementation may rely on different mechanisms. The caller should not be
concerned with the specifics of the implementation. At this point, the method
is expected to only be used within the body of another TF computation (within
an instance of `tf_computation_context.TensorFlowComputationContext` at the
top of the stack), and potentially also in certain types of interpreted
execution contexts (TBD).
Args:
computation_proto: An instance of `pb.Computation` with the `computation`
one of equal to `tensorflow` to be deserialized and called.
arg: The argument to invoke the computation with, or None if the computation
does not specify a parameter type and does not expects one.
graph: The graph to stamp into.
Returns:
A tuple (init_op, result) where:
init_op: String name of an op to initialize the graph.
result: The results to be fetched from TensorFlow. Depending on
the type of the result, this can be `tf.Tensor` or `tf.data.Dataset`
instances, or a nested structure (such as an
`anonymous_tuple.AnonymousTuple`).
Raises:
TypeError: If the arguments are of the wrong types.
ValueError: If `computation_proto` is not a TensorFlow computation proto.
"""
py_typecheck.check_type(computation_proto, pb.Computation)
computation_oneof = computation_proto.WhichOneof('computation')
if computation_oneof != 'tensorflow':
raise ValueError(
'Expected a TensorFlow computation, got {}.'.format(computation_oneof))
py_typecheck.check_type(graph, tf.Graph)
with graph.as_default():
type_spec = type_serialization.deserialize_type(computation_proto.type)
if type_spec.parameter is None:
if arg is None:
input_map = None
else:
raise TypeError(
'The computation declared no parameters; encountered an unexpected '
'argument {}.'.format(arg))
elif arg is None:
raise TypeError(
'The computation declared a parameter of type {}, but the argument '
'was not supplied.'.format(type_spec.parameter))
else:
arg_type, arg_binding = tensorflow_utils.capture_result_from_graph(
arg, graph)
if not type_analysis.is_assignable_from(type_spec.parameter, arg_type):
raise TypeError(
'The computation declared a parameter of type {}, but the argument '
'is of a mismatching type {}.'.format(type_spec.parameter,
arg_type))
else:
input_map = {
k: graph.get_tensor_by_name(v)
for k, v in tensorflow_utils.compute_map_from_bindings(
computation_proto.tensorflow.parameter, arg_binding).items()
}
return_elements = tensorflow_utils.extract_tensor_names_from_binding(
computation_proto.tensorflow.result)
orig_init_op_name = computation_proto.tensorflow.initialize_op
if orig_init_op_name:
return_elements.append(orig_init_op_name)
# N. B. Unlike MetaGraphDef, the GraphDef alone contains no information
# about collections, and hence, when we import a graph with Variables,
# those Variables are not added to global collections, and hence
# functions like tf.compat.v1.global_variables_initializers() will not
# contain their initialization ops.
output_tensors = tf.import_graph_def(
serialization_utils.unpack_graph_def(
computation_proto.tensorflow.graph_def),
input_map,
return_elements,
# N. B. It is very important not to return any names from the original
# computation_proto.tensorflow.graph_def, those names might or might not
# be valid in the current graph. Using a different scope makes the graph
# somewhat more readable, since _N style de-duplication of graph
# node names is less likely to be needed.
name='subcomputation')
output_map = {k: v for k, v in zip(return_elements, output_tensors)}
new_init_op_name = output_map.pop(orig_init_op_name, None)
return (new_init_op_name,
tensorflow_utils.assemble_result_from_graph(
type_spec.result, computation_proto.tensorflow.result,
output_map))
