async def create_selection(self, source, index=None, name=None):
py_typecheck.check_type(source, CachedValue)
py_typecheck.check_type(source.type_signature,
computation_types.NamedTupleType)
source_val = await source.target_future
if index is not None:
py_typecheck.check_none(name)
identifier_str = '{}[{}]'.format(source.identifier, index)
type_spec = source.type_signature[index]
else:
py_typecheck.check_not_none(name)
identifier_str = '{}.{}'.format(source.identifier, name)
type_spec = getattr(source.type_signature, name)
identifier = CachedValueIdentifier(identifier_str)
try:
cached_value = self._cache[identifier]
except KeyError:
target_future = asyncio.ensure_future(
self._target_executor.create_selection(
source_val, index=index, name=name))
cached_value = CachedValue(identifier, None, type_spec, target_future)
self._cache[identifier] = cached_value
try:
target_value = await cached_value.target_future
except Exception as e:
# TODO(b/145514490): This is a bit heavy handed, there maybe caches where
# only the current cache item needs to be invalidated; however this
# currently only occurs when an inner RemoteExecutor has the backend go
# down.
self._cache = {}
raise e
type_analysis.check_assignable_from(type_spec, target_value.type_signature)
return cached_value
async def create_call(self, comp, arg=None):
py_typecheck.check_type(comp, CachedValue)
py_typecheck.check_type(comp.type_signature, computation_types.FunctionType)
to_gather = [comp.target_future]
if arg is not None:
py_typecheck.check_type(arg, CachedValue)
type_analysis.check_assignable_from(comp.type_signature.parameter,
arg.type_signature)
to_gather.append(arg.target_future)
identifier_str = '{}({})'.format(comp.identifier, arg.identifier)
else:
identifier_str = '{}()'.format(comp.identifier)
gathered = await asyncio.gather(*to_gather)
type_spec = comp.type_signature.result
identifier = CachedValueIdentifier(identifier_str)
try:
cached_value = self._cache[identifier]
except KeyError:
target_future = asyncio.ensure_future(
self._target_executor.create_call(*gathered))
cached_value = CachedValue(identifier, None, type_spec, target_future)
self._cache[identifier] = cached_value
try:
target_value = await cached_value.target_future
except Exception as e:
# TODO(b/145514490): This is a bit heavy handed, there maybe caches where
# only the current cache item needs to be invalidated; however this
# currently only occurs when an inner RemoteExecutor has the backend go
# down.
self._cache = {}
raise e
type_analysis.check_assignable_from(type_spec, target_value.type_signature)
return cached_value
def deserialize_value(value_proto):
"""Deserializes a value (of any type) from `executor_pb2.Value`.
Args:
value_proto: An instance of `executor_pb2.Value`.
Returns:
A tuple `(value, type_spec)`, where `value` is a deserialized
representation of the transmitted value (e.g., Numpy array, or a
`pb.Computation` instance), and `type_spec` is an instance of
`tff.TensorType` that represents its type.
Raises:
TypeError: If the arguments are of the wrong types.
ValueError: If the value is malformed.
"""
py_typecheck.check_type(value_proto, executor_pb2.Value)
which_value = value_proto.WhichOneof('value')
if which_value == 'tensor':
return deserialize_tensor_value(value_proto)
elif which_value == 'computation':
return (value_proto.computation,
type_serialization.deserialize_type(
value_proto.computation.type))
elif which_value == 'tuple':
val_elems = []
type_elems = []
for e in value_proto.tuple.element:
name = e.name if e.name else None
e_val, e_type = deserialize_value(e.value)
val_elems.append((name, e_val))
type_elems.append((name, e_type) if name else e_type)
return (anonymous_tuple.AnonymousTuple(val_elems),
computation_types.NamedTupleType(type_elems))
elif which_value == 'sequence':
return deserialize_sequence_value(value_proto.sequence)
elif which_value == 'federated':
type_spec = type_serialization.deserialize_type(
computation_pb2.Type(federated=value_proto.federated.type))
value = []
for item in value_proto.federated.value:
item_value, item_type = deserialize_value(item)
type_analysis.check_assignable_from(type_spec.member, item_type)
value.append(item_value)
if type_spec.all_equal:
if len(value) == 1:
value = value[0]
else:
raise ValueError(
'Return an all_equal value with {} member consatituents.'.
format(len(value)))
return value, type_spec
else:
raise ValueError(
'Unable to deserialize a value of type {}.'.format(which_value))
def serialize_tensor_value(value, type_spec=None):
"""Serializes a tensor value into `executor_pb2.Value`.
Args:
value: A Numpy array or other object understood by `tf.make_tensor_proto`.
type_spec: An optional type spec, a `tff.TensorType` or something
convertible to it.
Returns:
A tuple `(value_proto, ret_type_spec)` in which `value_proto` is an instance
of `executor_pb2.Value` with the serialized content of `value`, and
`ret_type_spec` is the type of the serialized value. The `ret_type_spec` is
the same as the argument `type_spec` if that argument was not `None`. If
the argument was `None`, `ret_type_spec` is a type determined from `value`.
Raises:
TypeError: If the arguments are of the wrong types.
ValueError: If the value is malformed.
"""
if isinstance(value, tf.Tensor):
if type_spec is None:
type_spec = computation_types.TensorType(
dtype=tf.dtypes.as_dtype(value.dtype),
shape=tf.TensorShape(value.shape))
value = value.numpy()
if type_spec is not None:
type_spec = computation_types.to_type(type_spec)
py_typecheck.check_type(type_spec, computation_types.TensorType)
if isinstance(value, np.ndarray):
tensor_proto = tf.make_tensor_proto(value,
dtype=type_spec.dtype,
verify_shape=False)
type_analysis.check_assignable_from(
type_spec,
computation_types.TensorType(
dtype=tf.dtypes.as_dtype(tensor_proto.dtype),
shape=tf.TensorShape(tensor_proto.tensor_shape)))
else:
tensor_proto = tf.make_tensor_proto(value,
dtype=type_spec.dtype,
shape=type_spec.shape,
verify_shape=True)
else:
tensor_proto = tf.make_tensor_proto(value)
type_spec = computation_types.TensorType(
dtype=tf.dtypes.as_dtype(tensor_proto.dtype),
shape=tf.TensorShape(tensor_proto.tensor_shape))
any_pb = any_pb2.Any()
any_pb.Pack(tensor_proto)
return executor_pb2.Value(tensor=any_pb), type_spec
def _stamp_value_into_graph(value, type_signature, graph):
"""Stamps `value` in `graph` as an object of type `type_signature`.
Args:
value: An value to stamp.
type_signature: An instance of `computation_types.Type`.
graph: The graph to stamp in.
Returns:
A Python object made of tensors stamped into `graph`, `tf.data.Dataset`s,
and `anonymous_tuple.AnonymousTuple`s that structurally corresponds to the
value passed at input.
"""
py_typecheck.check_type(type_signature, computation_types.Type)
py_typecheck.check_type(graph, tf.Graph)
if value is None:
return None
if isinstance(type_signature, computation_types.TensorType):
if isinstance(value, np.ndarray):
value_type = computation_types.TensorType(
tf.dtypes.as_dtype(value.dtype), tf.TensorShape(value.shape))
type_analysis.check_assignable_from(type_signature, value_type)
with graph.as_default():
return tf.constant(value)
else:
with graph.as_default():
return tf.constant(value,
dtype=type_signature.dtype,
shape=type_signature.shape)
elif isinstance(type_signature, computation_types.NamedTupleType):
if isinstance(value, (list, dict)):
value = anonymous_tuple.from_container(value)
stamped_elements = []
named_type_signatures = anonymous_tuple.to_elements(type_signature)
for (name, type_signature), element in zip(named_type_signatures,
value):
stamped_element = _stamp_value_into_graph(element, type_signature,
graph)
stamped_elements.append((name, stamped_element))
return anonymous_tuple.AnonymousTuple(stamped_elements)
elif isinstance(type_signature, computation_types.SequenceType):
return tensorflow_utils.make_data_set_from_elements(
graph, value, type_signature.element)
else:
raise NotImplementedError(
'Unable to stamp a value of type {} in graph.'.format(
type_signature))
async def create_call(self, comp, arg=None):
py_typecheck.check_type(comp, CompositeValue)
if arg is not None:
py_typecheck.check_type(arg, CompositeValue)
py_typecheck.check_type(comp.type_signature,
computation_types.FunctionType)
param_type = comp.type_signature.parameter
type_analysis.check_assignable_from(param_type, arg.type_signature)
arg = CompositeValue(arg.internal_representation, param_type)
if isinstance(comp.internal_representation, pb.Computation):
which_computation = comp.internal_representation.WhichOneof(
'computation')
if which_computation == 'tensorflow':
call_args = [
self._parent_executor.create_value(
comp.internal_representation, comp.type_signature)
]
if arg is not None:
call_args.append(
executor_utils.delegate_entirely_to_executor(
arg.internal_representation, arg.type_signature,
self._parent_executor))
result = await self._parent_executor.create_call(
*(await asyncio.gather(*call_args)))
return CompositeValue(result, result.type_signature)
else:
raise ValueError(
'Directly calling computations of type {} is unsupported.'.
format(which_computation))
elif isinstance(comp.internal_representation,
intrinsic_defs.IntrinsicDef):
coro = getattr(
self, '_compute_intrinsic_{}'.format(
comp.internal_representation.uri), None)
if coro is not None:
return await coro(arg) # pylint: disable=not-callable
else:
raise NotImplementedError(
'Support for intrinsic "{}" has not been implemented yet.'.
format(comp.internal_representation.uri))
else:
raise ValueError(
'Calling objects of type {} is unsupported.'.format(
py_typecheck.type_string(type(
comp.internal_representation))))
async def create_tuple(self, elements):
if not isinstance(elements, anonymous_tuple.AnonymousTuple):
elements = anonymous_tuple.from_container(elements)
element_strings = []
element_kv_pairs = anonymous_tuple.to_elements(elements)
to_gather = []
type_elements = []
for k, v in element_kv_pairs:
py_typecheck.check_type(v, CachedValue)
to_gather.append(v.target_future)
if k is not None:
py_typecheck.check_type(k, str)
element_strings.append('{}={}'.format(k, v.identifier))
type_elements.append((k, v.type_signature))
else:
element_strings.append(str(v.identifier))
type_elements.append(v.type_signature)
type_spec = computation_types.NamedTupleType(type_elements)
gathered = await asyncio.gather(*to_gather)
identifier = CachedValueIdentifier('<{}>'.format(
','.join(element_strings)))
try:
cached_value = self._cache[identifier]
except KeyError:
target_future = asyncio.ensure_future(
self._target_executor.create_tuple(
anonymous_tuple.AnonymousTuple(
(k, v)
for (k, _), v in zip(element_kv_pairs, gathered))))
cached_value = CachedValue(identifier, None, type_spec,
target_future)
self._cache[identifier] = cached_value
try:
target_value = await cached_value.target_future
except Exception as e:
# TODO(b/145514490): This is a bit heavy handed, there maybe caches where
# only the current cache item needs to be invalidated; however this
# currently only occurs when an inner RemoteExecutor has the backend go
# down.
self._cache = {}
raise e
type_analysis.check_assignable_from(type_spec,
target_value.type_signature)
return cached_value
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 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 __init__(self, initialize, prepare, work, zero, accumulate, merge,
report, bitwidth, update):
"""Constructs a representation of a MapReduce-like iterative process.
Note: All the computations supplied here as arguments must be TensorFlow
computations, i.e., instances of `tff.Computation` constructed by the
`tff.tf_computation` decorator/wrapper.
Args:
initialize: The computation that produces the initial server state.
prepare: The computation that prepares the input for the clients.
work: The client-side work computation.
zero: The computation that produces the initial state for accumulators.
accumulate: The computation that adds a client update to an accumulator.
merge: The computation to use for merging pairs of accumulators.
report: The computation that produces the final server-side aggregate for
the top level accumulator (the global update).
bitwidth: The computation that produces the bitwidth for secure sum.
update: The computation that takes the global update and the server state
and produces the new server state, as well as server-side output.
Raises:
TypeError: If the Python or TFF types of the arguments are invalid or not
compatible with each other.
AssertionError: If the manner in which the given TensorFlow computations
are represented by TFF does not match what this code is expecting (this
is an internal error that requires code update).
"""
for label, comp in [
('initialize', initialize),
('prepare', prepare),
('work', work),
('zero', zero),
('accumulate', accumulate),
('merge', merge),
('report', report),
('bitwidth', bitwidth),
('update', update),
]:
py_typecheck.check_type(comp, computation_base.Computation, label)
# TODO(b/130633916): Remove private access once an appropriate API for it
# becomes available.
comp_proto = comp._computation_proto # pylint: disable=protected-access
if not isinstance(comp_proto, computation_pb2.Computation):
# Explicitly raised to force it to be done in non-debug mode as well.
raise AssertionError(
'Cannot find the embedded computation definition.')
which_comp = comp_proto.WhichOneof('computation')
if which_comp != 'tensorflow':
raise TypeError(
'Expected all computations supplied as arguments to '
'be plain TensorFlow, found {}.'.format(which_comp))
if prepare.type_signature.parameter != initialize.type_signature.result:
raise TypeError(
'The `prepare` computation expects an argument of type {}, '
'which does not match the result type {} of `initialize`.'.
format(prepare.type_signature.parameter,
initialize.type_signature.result))
if (not isinstance(work.type_signature.parameter,
computation_types.NamedTupleType)
or len(work.type_signature.parameter) != 2):
raise TypeError(
'The `work` computation expects an argument of type {} that is not '
'a two-tuple.'.format(work.type_signature.parameter))
if work.type_signature.parameter[1] != prepare.type_signature.result:
raise TypeError(
'The `work` computation expects an argument tuple with type {} as '
'the second element (the initial client state from the server), '
'which does not match the result type {} of `prepare`.'.format(
work.type_signature.parameter[1],
prepare.type_signature.result))
if (not isinstance(work.type_signature.result,
computation_types.NamedTupleType)
or len(work.type_signature.result) != 2):
raise TypeError(
'The `work` computation returns a result of type {} that is not a '
'two-tuple.'.format(work.type_signature.result))
py_typecheck.check_type(zero.type_signature,
computation_types.FunctionType)
py_typecheck.check_type(accumulate.type_signature,
computation_types.FunctionType)
py_typecheck.check_len(accumulate.type_signature.parameter, 2)
type_analysis.check_assignable_from(
accumulate.type_signature.parameter[0], zero.type_signature.result)
if (accumulate.type_signature.parameter[1] !=
work.type_signature.result[0][0]):
raise TypeError(
'The `accumulate` computation expects a second argument of type {}, '
'which does not match the expected {} as implied by the type '
'signature of `work`.'.format(
accumulate.type_signature.parameter[1],
work.type_signature.result[0][0]))
type_analysis.check_assignable_from(
accumulate.type_signature.parameter[0],
accumulate.type_signature.result)
py_typecheck.check_type(merge.type_signature,
computation_types.FunctionType)
py_typecheck.check_len(merge.type_signature.parameter, 2)
type_analysis.check_assignable_from(merge.type_signature.parameter[0],
accumulate.type_signature.result)
type_analysis.check_assignable_from(merge.type_signature.parameter[1],
accumulate.type_signature.result)
type_analysis.check_assignable_from(merge.type_signature.parameter[0],
merge.type_signature.result)
py_typecheck.check_type(report.type_signature,
computation_types.FunctionType)
type_analysis.check_assignable_from(report.type_signature.parameter,
merge.type_signature.result)
py_typecheck.check_type(bitwidth.type_signature,
computation_types.FunctionType)
expected_update_parameter_type = computation_types.to_type([
initialize.type_signature.result,
[report.type_signature.result, work.type_signature.result[0][1]],
])
if update.type_signature.parameter != expected_update_parameter_type:
raise TypeError(
'The `update` computation expects an argument of type {}, '
'which does not match the expected {} as implied by the type '
'signatures of `initialize`, `report`, and `work`.'.format(
update.type_signature.parameter,
expected_update_parameter_type))
if (not isinstance(update.type_signature.result,
computation_types.NamedTupleType)
or len(update.type_signature.result) != 2):
raise TypeError(
'The `update` computation returns a result of type {} that is not '
'a two-tuple.'.format(update.type_signature.result))
if update.type_signature.result[0] != initialize.type_signature.result:
raise TypeError(
'The `update` computation returns a result tuple with type {} as '
'the first element (the updated state of the server), which does '
'not match the result type {} of `initialize`.'.format(
update.type_signature.result[0],
initialize.type_signature.result))
self._initialize = initialize
self._prepare = prepare
self._work = work
self._zero = zero
self._accumulate = accumulate
self._merge = merge
self._report = report
self._bitwidth = bitwidth
self._update = update
async def create_call(self, comp, arg=None):
"""A coroutine that creates a call to `comp` with optional argument `arg`.
Args:
comp: The computation to invoke.
arg: An optional argument of the call, or `None` if no argument was
supplied.
Returns:
An instance of `FederatingExecutorValue` that represents the constructed
call.
Raises:
TypeError: If `comp` or `arg` are not embedded in the executor, before
calling `create_call` or if the `type_signature` of `arg` do not match
the expected `type_signature` of the parameter to `comp`.
ValueError: If `comp` is not a functional kind recognized by
`FederatingExecutor` or if `comp` is a lambda with an argument.
NotImplementedError: If `comp` is an intrinsic and it has not been
implemented by the `FederatingExecutor`.
"""
py_typecheck.check_type(comp, FederatingExecutorValue)
if arg is not None:
py_typecheck.check_type(arg, FederatingExecutorValue)
py_typecheck.check_type(comp.type_signature,
computation_types.FunctionType)
param_type = comp.type_signature.parameter
type_analysis.check_assignable_from(param_type, arg.type_signature)
arg = FederatingExecutorValue(arg.internal_representation,
param_type)
if isinstance(comp.internal_representation, pb.Computation):
which_computation = comp.internal_representation.WhichOneof(
'computation')
if which_computation == 'tensorflow':
child = self._target_executors[None][0]
embedded_comp = await child.create_value(
comp.internal_representation, comp.type_signature)
if arg is not None:
embedded_arg = await executor_utils.delegate_entirely_to_executor(
arg.internal_representation, arg.type_signature, child)
else:
embedded_arg = None
result = await child.create_call(embedded_comp, embedded_arg)
return FederatingExecutorValue(result, result.type_signature)
else:
raise ValueError(
'Directly calling computations of type {} is unsupported.'.
format(which_computation))
elif isinstance(comp.internal_representation,
intrinsic_defs.IntrinsicDef):
coro = getattr(
self, '_compute_intrinsic_{}'.format(
comp.internal_representation.uri), None)
if coro is not None:
return await coro(arg) # pylint: disable=not-callable
else:
raise NotImplementedError(
'Support for intrinsic "{}" has not been implemented yet.'.
format(comp.internal_representation.uri))
else:
raise ValueError(
'Calling objects of type {} is unsupported.'.format(
py_typecheck.type_string(type(
comp.internal_representation))))
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_analysis.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):
deserialized_type = type_serialization.deserialize_type(value.type)
if type_spec is None:
type_spec = deserialized_type
else:
type_analysis.check_assignable_from(type_spec,
deserialized_type)
which_computation = value.WhichOneof('computation')
if which_computation in ['lambda', 'tensorflow']:
return FederatingExecutorValue(value, type_spec)
elif which_computation == 'intrinsic':
intrinsic_def = intrinsic_defs.uri_to_intrinsic_def(
value.intrinsic.uri)
if intrinsic_def is None:
raise ValueError(
'Encountered an unrecognized intrinsic "{}".'.format(
value.intrinsic.uri))
return await self.create_value(intrinsic_def, type_spec)
else:
raise ValueError(
'Unsupported computation building block of type "{}".'.
format(which_computation))
elif isinstance(type_spec, computation_types.FederatedType):
self._check_value_compatible_with_placement(
value, type_spec.placement, type_spec.all_equal)
children = self._target_executors[type_spec.placement]
if type_spec.all_equal:
value = [value for _ in children]
results = await asyncio.gather(*[
c.create_value(v, type_spec.member)
for v, c in zip(value, children)
])
return FederatingExecutorValue(results, type_spec)
else:
child = self._target_executors[None][0]
return FederatingExecutorValue(
await child.create_value(value, type_spec), type_spec)
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_analysis.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):
deserialized_type = type_serialization.deserialize_type(value.type)
if type_spec is None:
type_spec = deserialized_type
else:
type_analysis.check_assignable_from(type_spec,
deserialized_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)
