def serialize_value(
value: Any,
type_spec: Optional[computation_types.Type] = None,
*,
value_proto: Optional[serialization_bindings.Value] = None
) -> computation_types.Type:
"""Serializes a value into `serialization_bindings.Value`.
We use a switch/function pattern in the body here (and in `deserialize_value`
below in order to persist more information in traces and profiling.
Args:
value: A value to be serialized.
type_spec: Optional type spec, a `tff.Type` or something convertible to it.
value_proto: The protobuf instance to serialize into. `value_proto` will
first be cleared before serializing.
Returns:
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.
"""
if value_proto is None:
value_proto = serialization_bindings.Value()
type_spec = computation_types.to_type(type_spec)
if isinstance(value, computation_pb2.Computation):
return _serialize_computation(value, type_spec)
elif isinstance(value, computation_impl.ComputationImpl):
return _serialize_computation(
computation_impl.ComputationImpl.get_proto(value),
executor_utils.reconcile_value_with_type_spec(value, type_spec))
elif type_spec is None:
raise TypeError(
'A type hint is required when serializing a value which '
'is not a TFF computation. Asked to serialized value {v} '
' of type {t} with None type spec.'.format(v=value, t=type(value)))
elif type_spec.is_tensor():
return _serialize_tensor_value(value_proto, value, type_spec)
elif type_spec.is_sequence():
return _serialize_sequence_value(value_proto, value, type_spec)
elif type_spec.is_struct():
return _serialize_struct_type(value_proto, value, type_spec)
elif type_spec.is_federated():
proto, type_spec = _serialize_federated_value(value_proto, value,
type_spec)
return 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'))
async def create_value(self, value, type_spec=None):
type_spec = computation_types.to_type(type_spec)
if isinstance(value, computation_impl.ComputationImpl):
return await self.create_value(
computation_impl.ComputationImpl.get_proto(value),
executor_utils.reconcile_value_with_type_spec(
value, type_spec))
py_typecheck.check_type(type_spec, computation_types.Type)
hashable_key = _get_hashable_key(value, type_spec)
try:
identifier = self._cache.get(hashable_key)
except TypeError as err:
raise RuntimeError(
'Failed to perform a hash table lookup with a value of Python '
'type {} and TFF type {}, and payload {}: {}'.format(
py_typecheck.type_string(type(value)), type_spec, value,
err))
if isinstance(identifier, CachedValueIdentifier):
cached_value = self._cache.get(identifier)
# If may be that the same payload appeared with a mismatching type spec,
# which may be a legitimate use case if (as it happens) the payload alone
# does not uniquely determine the type, so we simply opt not to reuse the
# cache value and fallback on the regular behavior.
if (cached_value is not None and type_spec is not None and
not cached_value.type_signature.is_equivalent_to(type_spec)
):
identifier = None
else:
identifier = None
if identifier is None:
self._num_values_created = self._num_values_created + 1
identifier = CachedValueIdentifier(str(self._num_values_created))
self._cache[hashable_key] = identifier
target_future = asyncio.ensure_future(
self._target_executor.create_value(value, type_spec))
cached_value = None
if cached_value is None:
cached_value = CachedValue(identifier, hashable_key, type_spec,
target_future)
self._cache[identifier] = cached_value
try:
await cached_value.target_future
except Exception:
# Invalidate the entire cache in the inner executor had an exception.
# 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
# No type check is necessary here; we have either checked
# `is_equivalent_to` or just constructed `target_value`
# explicitly with `type_spec`.
return cached_value
def to_representation_for_type(value, type_spec=None, backend=None):
"""Verifies or converts the `value` to executor payload matching `type_spec`.
The following kinds of `value` are supported:
* Computations, `pb.Computation` or `computation_impl.ConcreteComputation`.
These are compiled and converted into `runtime.ComputationCallable`.
* Numpy arrays and scalars, or Python scalars that are converted to Numpy.
Args:
value: The raw representation of a value to compare against `type_spec` and
potentially to be converted.
type_spec: An instance of `tff.Type`. Can be `None` for values that derive
from `typed_object.TypedObject`.
backend: Optional information about the backend, only required for
computations. Must be `None` or an instance of `backend_info.BackendInfo`.
Returns:
Either `value` itself, or a modified version of it.
Raises:
TypeError: If the `value` is not compatible with `type_spec`.
ValueError: If the arguments are incorrect (e.g., missing `backend` for a
computation-typed `value`).
"""
type_spec = executor_utils.reconcile_value_with_type_spec(value, type_spec)
if backend is not None:
py_typecheck.check_type(backend, backend_info.BackendInfo)
if isinstance(value, computation_impl.ConcreteComputation):
return to_representation_for_type(
computation_impl.ConcreteComputation.get_proto(value),
type_spec=type_spec,
backend=backend)
elif isinstance(value, pb.Computation):
if backend is None:
raise ValueError('Missing backend info for a computation.')
module = compiler.import_tensorflow_computation(value)
return runtime.ComputationCallable(module, backend)
elif isinstance(type_spec, computation_types.TensorType):
type_spec.shape.assert_is_fully_defined()
type_analysis.check_type(value, type_spec)
if type_spec.shape.rank == 0:
return np.dtype(type_spec.dtype.as_numpy_dtype).type(value)
elif type_spec.shape.rank > 0:
return np.array(value, dtype=type_spec.dtype.as_numpy_dtype)
else:
raise TypeError('Unsupported tensor shape {}.'.format(
type_spec.shape))
else:
raise TypeError('Unexpected type {}.'.format(type_spec))
def to_representation_for_type(value, type_spec, backend=None):
"""Verifies or converts the `value` to executor payload matching `type_spec`.
The following kinds of `value` are supported:
* Computations, either `pb.Computation` or `computation_impl.ComputationImpl`.
* Numpy arrays and scalars, or Python scalars that are converted to Numpy.
* Nested structures of the above.
Args:
value: The raw representation of a value to compare against `type_spec` and
potentially to be converted.
type_spec: An instance of `tff.Type`. Can be `None` for values that derive
from `typed_object.TypedObject`.
backend: The backend to use; an instance of `xla_client.Client`. Only used
for functional types. Can be `None` if unused.
Returns:
Either `value` itself, or a modified version of it.
Raises:
TypeError: If the `value` is not compatible with `type_spec`.
ValueError: If the arguments are incorrect.
"""
if backend is not None:
py_typecheck.check_type(backend, xla_client.Client)
if type_spec is not None:
type_spec = computation_types.to_type(type_spec)
type_spec = executor_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), type_spec,
backend)
if isinstance(value, pb.Computation):
comp_type = type_serialization.deserialize_type(value.type)
if type_spec is not None:
comp_type.check_equivalent_to(type_spec)
return runtime.ComputationCallable(value, comp_type, backend)
if isinstance(type_spec, computation_types.StructType):
return structure.map_structure(
lambda v, t: to_representation_for_type(v, t, backend),
structure.from_container(value, recursive=True), type_spec)
if isinstance(type_spec, computation_types.TensorType):
return runtime.normalize_tensor_representation(value, type_spec)
raise TypeError('Unexpected type {}.'.format(type_spec))
async def create_value(self, value, type_spec=None):
type_spec = computation_types.to_type(type_spec)
if isinstance(value, computation_impl.ComputationImpl):
return await self.create_value(
computation_impl.ComputationImpl.get_proto(value),
executor_utils.reconcile_value_with_type_spec(value, type_spec))
elif isinstance(value, pb.Computation):
return await self._evaluate(value)
elif type_spec is not None and type_spec.is_struct():
v_el = structure.to_elements(structure.from_container(value))
vals = await asyncio.gather(
*[self.create_value(val, t) for (_, val), t in zip(v_el, type_spec)])
return ReferenceResolvingExecutorValue(
structure.Struct((name, val) for (name, _), val in zip(v_el, vals)))
else:
return ReferenceResolvingExecutorValue(await
self._target_executor.create_value(
value, type_spec))
async def create_value(self, value, type_spec=None):
if isinstance(value, computation_impl.ComputationImpl):
return await self.create_value(
computation_impl.ComputationImpl.get_proto(value),
executor_utils.reconcile_value_with_type_spec(
value, type_spec))
elif isinstance(value, pb.Computation):
return await self.create_value(
building_blocks.ComputationBuildingBlock.from_proto(value),
type_spec)
elif isinstance(value, building_blocks.ComputationBuildingBlock):
value = self._transformation_fn(value)
py_typecheck.check_type(value,
building_blocks.ComputationBuildingBlock)
return await self._target_executor.create_value(
value.proto, type_spec)
else:
return await self._target_executor.create_value(value, type_spec)
def test_reconcile_value_with_type_spec_raises_type_error(
self, value, type_spec):
with self.assertRaises(TypeError):
executor_utils.reconcile_value_with_type_spec(value, type_spec)
def test_reconcile_value_with_type_spec_returns_type(self, value, type_spec,
expected_type):
actual_type = executor_utils.reconcile_value_with_type_spec(
value, type_spec)
self.assertEqual(actual_type, expected_type)
def to_representation_for_type(
value: Any,
tf_function_cache: MutableMapping[str, Any],
type_spec: Optional[computation_types.Type] = None,
device: Optional[tf.config.LogicalDevice] = None) -> Any:
"""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: An optional `tf.config.LogicalDevice` 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 = executor_utils.reconcile_value_with_type_spec(value, type_spec)
if isinstance(value, computation_impl.ConcreteComputation):
return to_representation_for_type(
computation_impl.ConcreteComputation.get_proto(value),
tf_function_cache, type_spec, device)
elif isinstance(value, pb.Computation):
computation_oneof = value.WhichOneof('computation')
if computation_oneof != 'tensorflow':
raise ValueError('Eager TF Executor can only execute computations of '
'TensorFlow flavor; encountered a computation of type '
f'{computation_oneof}')
return _to_computation_internal_rep(
value=value,
tf_function_cache=tf_function_cache,
type_spec=type_spec,
device=device)
elif type_spec.is_struct():
return _to_struct_internal_rep(
value=value,
tf_function_cache=tf_function_cache,
type_spec=type_spec,
device=device)
elif device is not None:
py_typecheck.check_type(device, tf.config.LogicalDevice)
with tf.device(device.name):
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 type_spec.is_tensor():
return _to_tensor_internal_rep(value=value, type_spec=type_spec)
elif type_spec.is_sequence():
return _to_sequence_internal_rep(value=value, type_spec=type_spec)
else:
raise TypeError(
f'Unexpected type {type_spec} for value of type {type(value)}: {value}')
async def create_value(
self,
value: Any,
type_spec: Any = None) -> executor_value_base.ExecutorValue:
"""Creates an embedded value from the given `value` and `type_spec`.
The kinds of supported `value`s are:
* An instance of `intrinsic_defs.IntrinsicDef`.
* An instance of `placements.PlacementLiteral`.
* An instance of `pb.Computation` if of one of the following kinds:
intrinsic, lambda, tensorflow, xla, or data.
* A Python `list` if `type_spec` is a federated type.
Note: The `value` must be a list even if it is of an `all_equal` type or
if there is only a single participant associated with the given placement.
* A Python value if `type_spec` is a non-functional, non-federated type.
Args:
value: An object to embed in the executor, one of the supported types
defined by above.
type_spec: An optional type convertible to instance of `tff.Type` via
`tff.to_type`, the type of `value`.
Returns:
An instance of `executor_value_base.ExecutorValue` representing a value
embedded in the `FederatingExecutor` using a particular
`FederatingStrategy`.
Raises:
TypeError: If the `value` and `type_spec` do not match.
ValueError: If `value` is not a kind supported by the
`FederatingExecutor`.
"""
type_spec = computation_types.to_type(type_spec)
if isinstance(value, intrinsic_defs.IntrinsicDef):
type_analysis.check_concrete_instance_of(type_spec, value.type_signature)
return self._strategy.ingest_value(value, type_spec)
elif isinstance(value, placements.PlacementLiteral):
if type_spec is None:
type_spec = computation_types.PlacementType()
type_spec.check_placement()
return self._strategy.ingest_value(value, type_spec)
elif isinstance(value, computation_impl.ConcreteComputation):
return await self.create_value(
computation_impl.ConcreteComputation.get_proto(value),
executor_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_spec.check_assignable_from(deserialized_type)
which_computation = value.WhichOneof('computation')
if which_computation in ['lambda', 'tensorflow', 'xla', 'data']:
return self._strategy.ingest_value(value, type_spec)
elif which_computation == 'intrinsic':
if value.intrinsic.uri in FederatingExecutor._FORWARDED_INTRINSICS:
return self._strategy.ingest_value(value, type_spec)
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 type_spec is not None and type_spec.is_federated():
return await self._strategy.compute_federated_value(value, type_spec)
else:
result = await self._unplaced_executor.create_value(value, type_spec)
return self._strategy.ingest_value(result, type_spec)
