def __init__(self, type_spec):
"""Constructs a computation building block with the given TFF type.
Args:
type_spec: An instance of types.Type, or something convertible to it via
types.to_type().
"""
type_signature = computation_types.to_type(type_spec)
type_utils.check_well_formed(type_signature)
self._type_signature = type_signature
def test_extra_well_formed_check_nested_types(self):
nest_federated = computation_types.FederatedType(
computation_types.FederatedType(tf.int32, placements.CLIENTS),
placements.CLIENTS)
tuple_federated_nest = computation_types.NamedTupleType(
[nest_federated])
with self.assertRaisesRegex(TypeError,
r'A {int32}@CLIENTS has been encountered'):
type_utils.check_well_formed(tuple_federated_nest)
federated_inner = computation_types.FederatedType(
tf.int32, placements.CLIENTS)
tuple_on_federated = computation_types.NamedTupleType(
[federated_inner])
federated_outer = computation_types.FederatedType(
tuple_on_federated, placements.CLIENTS)
with self.assertRaisesRegex(TypeError,
r'A {int32}@CLIENTS has been encountered'):
type_utils.check_well_formed(federated_outer)
multiple_nest = computation_types.NamedTupleType(
[computation_types.NamedTupleType([federated_outer])])
with self.assertRaisesRegex(TypeError,
r'A {int32}@CLIENTS has been encountered'):
type_utils.check_well_formed(multiple_nest)
sequence_of_federated = computation_types.SequenceType(federated_inner)
with self.assertRaisesRegex(TypeError,
r'A {int32}@CLIENTS has been encountered'):
type_utils.check_well_formed(sequence_of_federated)
def __init__(self, computation_proto, context_stack, annotated_type=None):
"""Constructs a new instance of ComputationImpl from the computation_proto.
Args:
computation_proto: The protocol buffer that represents the computation, an
instance of pb.Computation.
context_stack: The context stack to use.
annotated_type: Optional, type information with additional annotations
that replaces the information in `computation_proto.type`.
Raises:
TypeError: if `annotated_type` is not `None` and is not compatible with
`computation_proto.type`.
"""
py_typecheck.check_type(computation_proto, pb.Computation)
py_typecheck.check_type(context_stack, context_stack_base.ContextStack)
type_spec = type_serialization.deserialize_type(computation_proto.type)
py_typecheck.check_type(type_spec, computation_types.Type)
if annotated_type is not None:
py_typecheck.check_type(annotated_type, computation_types.Type)
# Extra information is encoded in a NamedTupleTypeWithPyContainerType
# subclass which does not override __eq__. The two type specs should still
# compare as equal.
if type_spec != annotated_type:
raise TypeError(
'annotated_type not compatible with computation_proto.type\n'
'computation_proto.type: {!s}\n'
'annotated_type: {!s}'.format(type_spec, annotated_type)
)
type_spec = annotated_type
type_utils.check_well_formed(type_spec)
# We may need to modify the type signature to reflect the fact that in the
# underlying framework for composing computations, there is no concept of
# no-argument lambdas, but in Python, every computation needs to look like
# a function that needs to be invoked.
if not isinstance(type_spec, computation_types.FunctionType):
type_spec = computation_types.FunctionType(None, type_spec)
super(ComputationImpl, self).__init__(type_spec, context_stack)
self._computation_proto = computation_proto
def __init__(self, computation_proto, context_stack, annotated_type=None):
"""Constructs a new instance of ComputationImpl from the computation_proto.
Args:
computation_proto: The protocol buffer that represents the computation, an
instance of pb.Computation.
context_stack: The context stack to use.
annotated_type: Optional, type information with additional annotations
that replaces the information in `computation_proto.type`.
Raises:
TypeError: if `annotated_type` is not `None` and is not compatible with
`computation_proto.type`.
"""
py_typecheck.check_type(computation_proto, pb.Computation)
py_typecheck.check_type(context_stack, context_stack_base.ContextStack)
type_spec = type_serialization.deserialize_type(computation_proto.type)
py_typecheck.check_type(type_spec, computation_types.Type)
if annotated_type is not None:
py_typecheck.check_type(annotated_type, computation_types.Type)
# Extra information is encoded in a NamedTupleTypeWithPyContainerType
# subclass which does not override __eq__. The two type specs should still
# compare as equal.
if type_spec != annotated_type:
raise TypeError(
'annotated_type not compatible with computation_proto.type\n'
'computation_proto.type: {!s}\n'
'annotated_type: {!s}'.format(type_spec, annotated_type))
type_spec = annotated_type
type_utils.check_well_formed(type_spec)
if not isinstance(type_spec, computation_types.FunctionType):
raise TypeError(
'{} is not a functional type, from proto: {}'.format(
str(type_spec), str(computation_proto)))
super().__init__(type_spec, context_stack)
self._computation_proto = computation_proto
def __init__(self, computation_proto, context_stack):
"""Constructs a new instance of ComputationImpl from the computation_proto.
Args:
computation_proto: The protocol buffer that represents the computation, an
instance of pb.Computation.
context_stack: The context stack to use.
"""
py_typecheck.check_type(computation_proto, pb.Computation)
py_typecheck.check_type(context_stack, context_stack_base.ContextStack)
type_spec = type_serialization.deserialize_type(computation_proto.type)
py_typecheck.check_type(type_spec, computation_types.Type)
type_utils.check_well_formed(type_spec)
# We may need to modify the type signature to reflect the fact that in the
# underlying framework for composing computations, there is no concept of
# no-argument lambdas, but in Python, every computation needs to look like
# a function that needs to be invoked.
if not isinstance(type_spec, computation_types.FunctionType):
type_spec = computation_types.FunctionType(None, type_spec)
super(ComputationImpl, self).__init__(type_spec, context_stack)
self._computation_proto = computation_proto
def test_well_formed_check_fails_bad_types(self):
nest_federated = computation_types.FederatedType(
computation_types.FederatedType(tf.int32, placements.CLIENTS),
placements.CLIENTS)
with self.assertRaisesRegex(TypeError,
'A {int32}@CLIENTS has been encountered'):
type_utils.check_well_formed(nest_federated)
sequence_in_sequence = computation_types.SequenceType(
computation_types.SequenceType([tf.int32]))
with self.assertRaisesRegex(TypeError, r'A <int32>\* has been encountered'):
type_utils.check_well_formed(sequence_in_sequence)
federated_fn = computation_types.FederatedType(
computation_types.FunctionType(tf.int32, tf.int32), placements.CLIENTS)
with self.assertRaisesRegex(TypeError,
r'A \(int32 -> int32\) has been encountered'):
type_utils.check_well_formed(federated_fn)
tuple_federated_fn = computation_types.NamedTupleType([federated_fn])
with self.assertRaisesRegex(TypeError,
r'A \(int32 -> int32\) has been encountered'):
type_utils.check_well_formed(tuple_federated_fn)
def test_well_formed_check_succeeds_good_types(self):
federated = computation_types.FederatedType(tf.int32, placements.CLIENTS)
self.assertTrue(type_utils.check_well_formed(federated))
tensor = computation_types.TensorType(tf.int32)
self.assertTrue(type_utils.check_well_formed(tensor))
namedtuple = computation_types.NamedTupleType(
[tf.int32,
computation_types.NamedTupleType([tf.int32, tf.int32])])
self.assertTrue(type_utils.check_well_formed(namedtuple))
sequence = computation_types.SequenceType(tf.int32)
self.assertTrue(type_utils.check_well_formed(sequence))
fn = computation_types.FunctionType(tf.int32, tf.int32)
self.assertTrue(type_utils.check_well_formed(fn))
abstract = computation_types.AbstractType('T')
self.assertTrue(type_utils.check_well_formed(abstract))
placement = computation_types.PlacementType()
self.assertTrue(type_utils.check_well_formed(placement))
def to_value(arg, type_spec, context_stack):
"""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: A type specifier that 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.
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_utils.check_well_formed(type_spec)
if isinstance(arg, ValueImpl):
result = arg
elif isinstance(arg, computation_building_blocks.ComputationBuildingBlock):
result = ValueImpl(arg, context_stack)
elif isinstance(arg, placement_literals.PlacementLiteral):
result = ValueImpl(computation_building_blocks.Placement(arg),
context_stack)
elif isinstance(arg, computation_base.Computation):
result = ValueImpl(
computation_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(
computation_building_blocks.Tuple([
(k, ValueImpl.get_comp(to_value(v, None, context_stack)))
for k, v in anonymous_tuple.to_elements(arg)
]), context_stack)
elif py_typecheck.is_named_tuple(arg):
result = to_value(arg._asdict(), None, context_stack)
elif isinstance(arg, dict):
if isinstance(arg, collections.OrderedDict):
items = six.iteritems(arg)
else:
items = sorted(six.iteritems(arg))
value = computation_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(
computation_building_blocks.Tuple([
ValueImpl.get_comp(to_value(x, None, context_stack))
for x in arg
]), context_stack)
elif isinstance(arg, dtype_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_utils.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(str(result.type_signature),
str(type_spec)))
return result
def test_check_abstract_types_are_bound_valid_cases(self, type_spec):
type_utils.check_well_formed(type_spec)
type_utils.check_all_abstract_types_are_bound(type_spec)
def test_types_are_well_formed(self, name):
type_utils.check_well_formed(
getattr(intrinsic_defs, name).type_signature)
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_utils.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_utils.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_utils.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
