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))