def capture_result_from_graph(result, graph): """Captures a result stamped into a tf.Graph as a type signature and binding. Args: result: The result to capture, a Python object that is composed of tensors, possibly nested within Python structures such as dictionaries, lists, tuples, or named tuples. graph: The instance of tf.Graph to use. Returns: A tuple (type_spec, binding), where 'type_spec' is an instance of computation_types.Type that describes the type of the result, and 'binding' is an instance of TensorFlow.Binding that indicates how parts of the result type relate to the tensors and ops that appear in the result. Raises: TypeError: If the argument or any of its parts are of an uexpected type. """ def _get_bindings_for_elements(name_value_pairs, graph, type_fn): """Build `(type_spec, binding)` tuple for name value pairs.""" element_name_type_binding_triples = [ ((k,) + capture_result_from_graph(v, graph)) for k, v in name_value_pairs ] type_spec = type_fn([((e[0], e[1]) if e[0] else e[1]) for e in element_name_type_binding_triples]) binding = pb.TensorFlow.Binding( tuple=pb.TensorFlow.NamedTupleBinding( element=[e[2] for e in element_name_type_binding_triples])) return type_spec, binding # TODO(b/113112885): The emerging extensions for serializing SavedModels may # end up introducing similar concepts of bindings, etc., we should look here # into the possibility of reusing some of that code when it's available. if isinstance(result, dtype_utils.TENSOR_REPRESENTATION_TYPES): with graph.as_default(): result = tf.constant(result) if tf.is_tensor(result): if hasattr(result, 'read_value'): # We have a tf.Variable-like result, get a proper tensor to fetch. with graph.as_default(): result = result.read_value() return (computation_types.TensorType(result.dtype.base_dtype, result.shape), pb.TensorFlow.Binding( tensor=pb.TensorFlow.TensorBinding(tensor_name=result.name))) elif py_typecheck.is_named_tuple(result): # Special handling needed for collections.namedtuples since they do not have # anything in the way of a shared base class. Note we don't want to rely on # the fact that collections.namedtuples inherit from 'tuple' because we'd be # failing to retain the information about naming of tuple members. # pylint: disable=protected-access name_value_pairs = six.iteritems(result._asdict()) # pylint: enable=protected-access return _get_bindings_for_elements( name_value_pairs, graph, functools.partial( computation_types.NamedTupleTypeWithPyContainerType, container_type=type(result))) elif py_typecheck.is_attrs(result): name_value_pairs = attr.asdict( result, dict_factory=collections.OrderedDict, recurse=False) return _get_bindings_for_elements( six.iteritems(name_value_pairs), graph, functools.partial( computation_types.NamedTupleTypeWithPyContainerType, container_type=type(result))) elif isinstance(result, anonymous_tuple.AnonymousTuple): return _get_bindings_for_elements( anonymous_tuple.to_elements(result), graph, computation_types.NamedTupleType) elif isinstance(result, collections.Mapping): if isinstance(result, collections.OrderedDict): name_value_pairs = six.iteritems(result) else: name_value_pairs = sorted(six.iteritems(result)) return _get_bindings_for_elements( name_value_pairs, graph, functools.partial( computation_types.NamedTupleTypeWithPyContainerType, container_type=type(result))) elif isinstance(result, (list, tuple)): element_type_binding_pairs = [ capture_result_from_graph(e, graph) for e in result ] return (computation_types.NamedTupleTypeWithPyContainerType( [e[0] for e in element_type_binding_pairs], type(result)), pb.TensorFlow.Binding( tuple=pb.TensorFlow.NamedTupleBinding( element=[e[1] for e in element_type_binding_pairs]))) elif isinstance(result, (tf.compat.v1.data.Dataset, tf.compat.v2.data.Dataset)): variant_tensor = tf.data.experimental.to_variant(result) # TODO(b/130032140): Switch to TF2.0 way of doing it while cleaning up the # legacy structures all over the code base and replacing them with the new # tf.data.experimenta.Structure variants. element_type = type_utils.tf_dtypes_and_shapes_to_type( tf.compat.v1.data.get_output_types(result), tf.compat.v1.data.get_output_shapes(result)) return (computation_types.SequenceType(element_type), pb.TensorFlow.Binding( sequence=pb.TensorFlow.SequenceBinding( variant_tensor_name=variant_tensor.name))) elif isinstance(result, OneShotDataset): # TODO(b/129956296): Eventually delete this deprecated code path. element_type = type_utils.tf_dtypes_and_shapes_to_type( tf.compat.v1.data.get_output_types(result), tf.compat.v1.data.get_output_shapes(result)) handle_name = result.make_one_shot_iterator().string_handle().name return (computation_types.SequenceType(element_type), pb.TensorFlow.Binding( sequence=pb.TensorFlow.SequenceBinding( iterator_string_handle_name=handle_name))) else: raise TypeError('Cannot capture a result of an unsupported type {}.'.format( py_typecheck.type_string(type(result))))
def __repr__(self): return 'Tuple([{}])'.format(', '.join( '({}, {})'.format('\'{}\''.format(e[0]) if e[0] is not None else 'None', repr(e[1])) for e in anonymous_tuple.to_elements(self)))
def infer_type(arg): """Infers the TFF type of the argument (a `computation_types.Type` instance). WARNING: This function is only partially implemented. The kinds of arguments that are currently correctly recognized: - tensors, variables, and data sets, - things that are convertible to tensors (including numpy arrays, builtin types, as well as lists and tuples of any of the above, etc.), - nested lists, tuples, namedtuples, anonymous tuples, dict, and OrderedDicts. Args: arg: The argument, the TFF type of which to infer. Returns: Either an instance of `computation_types.Type`, or `None` if the argument is `None`. """ # TODO(b/113112885): Implement the remaining cases here on the need basis. if arg is None: return None elif isinstance(arg, typed_object.TypedObject): return arg.type_signature elif tf.contrib.framework.is_tensor(arg): return computation_types.TensorType(arg.dtype.base_dtype, arg.shape) elif isinstance(arg, tf.data.Dataset): return computation_types.SequenceType( tf_dtypes_and_shapes_to_type(arg.output_types, arg.output_shapes)) elif isinstance(arg, anonymous_tuple.AnonymousTuple): return computation_types.NamedTupleType([ (k, infer_type(v)) if k else infer_type(v) for k, v in anonymous_tuple.to_elements(arg) ]) elif py_typecheck.is_named_tuple(arg): # Special handling needed for collections.namedtuple. return infer_type(arg._asdict()) elif isinstance(arg, dict): if isinstance(arg, collections.OrderedDict): items = six.iteritems(arg) else: items = sorted(six.iteritems(arg)) return computation_types.NamedTupleType([(k, infer_type(v)) for k, v in items]) elif isinstance(arg, (tuple, list)): return computation_types.NamedTupleType([infer_type(e) for e in arg]) elif isinstance(arg, six.string_types): return computation_types.TensorType(tf.string) elif isinstance(arg, (np.generic, np.ndarray)): return computation_types.TensorType(tf.as_dtype(arg.dtype), arg.shape) else: dtype = { bool: tf.bool, int: tf.int32, float: tf.float32 }.get(type(arg)) if dtype: return computation_types.TensorType(dtype) else: # Now fall back onto the heavier-weight processing, as all else failed. # Use make_tensor_proto() to make sure to handle it consistently with # how TensorFlow is handling values (e.g., recognizing int as int32, as # opposed to int64 as in NumPy). try: # TODO(b/113112885): Find something more lightweight we could use here. tensor_proto = tf.make_tensor_proto(arg) return computation_types.TensorType( tf.DType(tensor_proto.dtype), tf.TensorShape(tensor_proto.tensor_shape)) except TypeError as err: raise TypeError( 'Could not infer the TFF type of {}: {}.'.format( py_typecheck.type_string(type(arg)), str(err)))
def pack_args_into_anonymous_tuple( args: Sequence[Any], kwargs: Mapping[str, Any], type_spec=None, context: Optional[context_base.Context] = None ) -> anonymous_tuple.AnonymousTuple: """Packs positional and keyword arguments into an anonymous tuple. If 'type_spec' is not None, it must be a tuple type or something that's convertible to it by computation_types.to_type(). The assignment of arguments to fields of the tuple follows the same rule as during function calls. If 'type_spec' is None, the positional arguments precede any of the keyword arguments, and the ordering of the keyword arguments matches the ordering in which they appear in kwargs. If the latter is an OrderedDict, the ordering will be preserved. On the other hand, if the latter is an ordinary unordered dict, the ordering is arbitrary. Args: args: Positional arguments. kwargs: Keyword arguments. type_spec: The optional type specification (either an instance of computation_types.NamedTupleType or something convertible to it), or None if there's no type. Used to drive the arrangements of args into fields of the constructed anonymous tuple, as noted in the description. context: The optional context (an instance of `context_base.Context`) in which the arguments are being packed. Required if and only if the `type_spec` is not `None`. Returns: An anoymous tuple containing all the arguments. Raises: TypeError: if the arguments are of the wrong computation_types. """ type_spec = computation_types.to_type(type_spec) if not type_spec: return anonymous_tuple.AnonymousTuple([(None, arg) for arg in args] + list(kwargs.items())) else: py_typecheck.check_type(type_spec, computation_types.NamedTupleType) py_typecheck.check_type(context, context_base.Context) context = context # type: context_base.Context if not is_argument_tuple(type_spec): # pylint: disable=attribute-error raise TypeError( 'Parameter type {} does not have a structure of an argument tuple, ' 'and cannot be populated from multiple positional and keyword ' 'arguments'.format(type_spec)) else: result_elements = [] positions_used = set() keywords_used = set() for index, (name, elem_type) in enumerate( anonymous_tuple.to_elements(type_spec)): if index < len(args): if name is not None and name in kwargs: raise TypeError( 'Argument {} specified twice.'.format(name)) else: arg_value = args[index] result_elements.append( (name, context.ingest(arg_value, elem_type))) positions_used.add(index) elif name is not None and name in kwargs: arg_value = kwargs[name] result_elements.append( (name, context.ingest(arg_value, elem_type))) keywords_used.add(name) elif name: raise TypeError( 'Argument named {} is missing.'.format(name)) else: raise TypeError( 'Argument at position {} is missing.'.format(index)) positions_missing = set(range( len(args))).difference(positions_used) if positions_missing: raise TypeError('Positional arguments at {} not used.'.format( positions_missing)) keywords_missing = set(kwargs.keys()).difference(keywords_used) if keywords_missing: raise TypeError('Keyword arguments at {} not used.'.format( keywords_missing)) return anonymous_tuple.AnonymousTuple(result_elements)
async def _compute_tuple(anon_tuple): elements = anonymous_tuple.to_elements(anon_tuple) keys = [k for k, _ in elements] vals = await asyncio.gather(*[_compute_element(v) for _, v in elements]) return anonymous_tuple.AnonymousTuple(zip(keys, vals))
def _get_index_from_name(selection_name, tuple_type_signature): named_type_signatures = anonymous_tuple.to_elements( tuple_type_signature) return [x[0] for x in named_type_signatures].index(selection_name)
def test_elements(self): self.assertEqual( repr( anonymous_tuple.to_elements( computation_types.NamedTupleType([tf.int32, ('a', tf.bool)]))), '[(None, TensorType(tf.int32)), (\'a\', TensorType(tf.bool))]')
def assemble_result_from_graph(type_spec, binding, output_map): """Assembles a result stamped into a `tf.Graph` given type signature/binding. This method does roughly the opposite of `capture_result_from_graph`, in that whereas `capture_result_from_graph` starts with a single structured object made up of tensors and computes its type and bindings, this method starts with the type/bindings and constructs a structured object made up of tensors. Args: type_spec: The type signature of the result to assemble, an instance of `types.Type` or something convertible to it. binding: The binding that relates the type signature to names of tensors in the graph, an instance of `pb.TensorFlow.Binding`. output_map: The mapping from tensor names that appear in the binding to actual stamped tensors (possibly renamed during import). Returns: The assembled result, a Python object that is composed of tensors, possibly nested within Python structures such as anonymous tuples. Raises: TypeError: If the argument or any of its parts are of an uexpected type. ValueError: If the arguments are invalid or inconsistent witch other, e.g., the type and binding don't match, or the tensor is not found in the map. """ type_spec = computation_types.to_type(type_spec) py_typecheck.check_type(type_spec, computation_types.Type) py_typecheck.check_type(binding, pb.TensorFlow.Binding) py_typecheck.check_type(output_map, dict) for k, v in six.iteritems(output_map): py_typecheck.check_type(k, six.string_types) if not tf.contrib.framework.is_tensor(v): raise TypeError( 'Element with key {} in the output map is {}, not a tensor.'. format(k, py_typecheck.type_string(type(v)))) binding_oneof = binding.WhichOneof('binding') if isinstance(type_spec, computation_types.TensorType): if binding_oneof != 'tensor': raise ValueError( 'Expected a tensor binding, found {}.'.format(binding_oneof)) elif binding.tensor.tensor_name not in output_map: raise ValueError( 'Tensor named {} not found in the output map.'.format( binding.tensor.tensor_name)) else: return output_map[binding.tensor.tensor_name] elif isinstance(type_spec, computation_types.NamedTupleType): if binding_oneof != 'tuple': raise ValueError( 'Expected a tuple binding, found {}.'.format(binding_oneof)) else: type_elements = anonymous_tuple.to_elements(type_spec) if len(binding.tuple.element) != len(type_elements): raise ValueError( 'Mismatching tuple sizes in type ({}) and binding ({}).'. format(len(type_elements), len(binding.tuple.element))) result_elements = [] for (element_name, element_type), element_binding in zip(type_elements, binding.tuple.element): element_object = assemble_result_from_graph( element_type, element_binding, output_map) result_elements.append((element_name, element_object)) return anonymous_tuple.AnonymousTuple(result_elements) elif isinstance(type_spec, computation_types.SequenceType): if binding_oneof != 'sequence': raise ValueError( 'Expected a sequence binding, found {}.'.format(binding_oneof)) else: handle = output_map[binding.sequence.iterator_string_handle_name] return make_dataset_from_string_handle(handle, type_spec.element) else: raise ValueError('Unsupported type \'{}\'.'.format(str(type_spec)))
def stamp_parameter_in_graph(parameter_name, parameter_type, graph): """Stamps a parameter of a given type in the given tf.Graph instance. Tensors are stamped as placeholders, sequences are stamped as data sets constructed from string tensor handles, and named tuples are stamped by independently stamping their elements. Args: parameter_name: The suggested (string) name of the parameter to use in determining the names of the graph components to construct. The names that will actually appear in the graph are not guaranteed to be based on this suggested name, and may vary, e.g., due to existing naming conflicts, but a best-effort attempt will be made to make them similar for ease of debugging. parameter_type: The type of the parameter to stamp. Must be either an instance of computation_types.Type (or convertible to it), or None. graph: The instance of tf.Graph to stamp in. Returns: A tuple (val, binding), where 'val' is a Python object (such as a dataset, a placeholder, or a dictionary that represents a named tuple) that represents the stamped parameter for use in the body of a Python function that consumes this parameter, and the 'binding' is an instance of TensorFlow.Binding that indicates how parts of the type signature relate to the tensors and ops stamped into the graph. Raises: TypeError: If the arguments are of the wrong computation_types. ValueError: If the parameter type cannot be stamped in a TensorFlow graph. """ py_typecheck.check_type(parameter_name, six.string_types) py_typecheck.check_type(graph, tf.Graph) if parameter_type is None: return (None, None) parameter_type = computation_types.to_type(parameter_type) if isinstance(parameter_type, computation_types.TensorType): with graph.as_default(): placeholder = tf.placeholder(dtype=parameter_type.dtype, shape=parameter_type.shape, name=parameter_name) binding = pb.TensorFlow.Binding(tensor=pb.TensorFlow.TensorBinding( tensor_name=placeholder.name)) return (placeholder, binding) elif isinstance(parameter_type, computation_types.NamedTupleType): element_name_value_pairs = [] element_bindings = [] for e in anonymous_tuple.to_elements(parameter_type): e_val, e_binding = stamp_parameter_in_graph( '{}_{}'.format(parameter_name, e[0]), e[1], graph) element_name_value_pairs.append((e[0], e_val)) element_bindings.append(e_binding) return (anonymous_tuple.AnonymousTuple(element_name_value_pairs), pb.TensorFlow.Binding(tuple=pb.TensorFlow.NamedTupleBinding( element=element_bindings))) elif isinstance(parameter_type, computation_types.SequenceType): with graph.as_default(): handle = tf.placeholder(tf.string, shape=[]) ds = make_dataset_from_string_handle(handle, parameter_type.element) return (ds, pb.TensorFlow.Binding(sequence=pb.TensorFlow.SequenceBinding( iterator_string_handle_name=handle.name))) else: raise ValueError( 'Parameter type component {} cannot be stamped into a TensorFlow ' 'graph.'.format(repr(parameter_type)))
def _concretize_abstract_types(abstract_type_spec, concrete_type_spec): """Recursive helper function to construct concrete type spec.""" if isinstance(abstract_type_spec, computation_types.AbstractType): bound_type = bound_abstract_types.get(str( abstract_type_spec.label)) if bound_type: return bound_type else: bound_abstract_types[str( abstract_type_spec.label)] = concrete_type_spec return concrete_type_spec elif isinstance(abstract_type_spec, computation_types.TensorType): return abstract_type_spec elif isinstance(abstract_type_spec, computation_types.NamedTupleType): if not isinstance(concrete_type_spec, computation_types.NamedTupleType): raise TypeError(type_error_string) abstract_elements = anonymous_tuple.to_elements(abstract_type_spec) concrete_elements = anonymous_tuple.to_elements(concrete_type_spec) if len(abstract_elements) != len(concrete_elements): raise TypeError(type_error_string) concretized_tuple_elements = [] for k in range(len(abstract_elements)): if abstract_elements[k][0] != concrete_elements[k][0]: raise TypeError(type_error_string) concretized_tuple_elements.append( (abstract_elements[k][0], _concretize_abstract_types(abstract_elements[k][1], concrete_elements[k][1]))) return computation_types.NamedTupleType(concretized_tuple_elements) elif isinstance(abstract_type_spec, computation_types.SequenceType): if not isinstance(concrete_type_spec, computation_types.SequenceType): raise TypeError(type_error_string) return computation_types.SequenceType( _concretize_abstract_types(abstract_type_spec.element, concrete_type_spec.element)) elif isinstance(abstract_type_spec, computation_types.FunctionType): if not isinstance(concrete_type_spec, computation_types.FunctionType): raise TypeError(type_error_string) concretized_param = _concretize_abstract_types( abstract_type_spec.parameter, concrete_type_spec.parameter) concretized_result = _concretize_abstract_types( abstract_type_spec.result, concrete_type_spec.result) return computation_types.FunctionType(concretized_param, concretized_result) elif isinstance(abstract_type_spec, computation_types.PlacementType): if not isinstance(concrete_type_spec, computation_types.PlacementType): raise TypeError(type_error_string) return abstract_type_spec elif isinstance(abstract_type_spec, computation_types.FederatedType): if not isinstance(concrete_type_spec, computation_types.FederatedType): raise TypeError(type_error_string) new_member = _concretize_abstract_types(abstract_type_spec.member, concrete_type_spec.member) return computation_types.FederatedType( new_member, abstract_type_spec.placement, abstract_type_spec.all_equal) elif abstract_type_spec is None: if concrete_type_spec is not None: raise TypeError(type_error_string) return None else: raise TypeError( 'Unexpected abstract typespec {}.'.format(abstract_type_spec))
def capture_result_from_graph(result, graph): """Captures a result stamped into a tf.Graph as a type signature and binding. Args: result: The result to capture, a Python object that is composed of tensors, possibly nested within Python structures such as dictionaries, lists, tuples, or named tuples. graph: The instance of tf.Graph to use. Returns: A tuple (type_spec, binding), where 'type_spec' is an instance of computation_types.Type that describes the type of the result, and 'binding' is an instance of TensorFlow.Binding that indicates how parts of the result type relate to the tensors and ops that appear in the result. Raises: TypeError: If the argument or any of its parts are of an uexpected type. """ # TODO(b/113112885): The emerging extensions for serializing SavedModels may # end up introducing similar concepts of bindings, etc., we should look here # into the possibility of reusing some of that code when it's available. if isinstance(result, dtype_utils.TENSOR_REPRESENTATION_TYPES): with graph.as_default(): result = tf.constant(result) if tf.contrib.framework.is_tensor(result): if hasattr(result, 'read_value'): # We have a tf.Variable-like result, get a proper tensor to fetch. with graph.as_default(): result = result.read_value() return (computation_types.TensorType(result.dtype.base_dtype, result.shape), pb.TensorFlow.Binding(tensor=pb.TensorFlow.TensorBinding( tensor_name=result.name))) elif py_typecheck.is_named_tuple(result): # Special handling needed for collections.namedtuples since they do not have # anything in the way of a shared base class. Note we don't want to rely on # the fact that collections.namedtuples inherit from 'tuple' because we'd be # failing to retain the information about naming of tuple members. # pylint: disable=protected-access return capture_result_from_graph(result._asdict(), graph) # pylint: enable=protected-access elif isinstance(result, (dict, anonymous_tuple.AnonymousTuple)): if isinstance(result, anonymous_tuple.AnonymousTuple): name_value_pairs = anonymous_tuple.to_elements(result) elif isinstance(result, collections.OrderedDict): name_value_pairs = six.iteritems(result) else: name_value_pairs = sorted(six.iteritems(result)) element_name_type_binding_triples = [ ((k, ) + capture_result_from_graph(v, graph)) for k, v in name_value_pairs ] return (computation_types.NamedTupleType([ ((e[0], e[1]) if e[0] else e[1]) for e in element_name_type_binding_triples ]), pb.TensorFlow.Binding(tuple=pb.TensorFlow.NamedTupleBinding( element=[e[2] for e in element_name_type_binding_triples]))) elif isinstance(result, (list, tuple)): element_type_binding_pairs = [ capture_result_from_graph(e, graph) for e in result ] return (computation_types.NamedTupleType( [e[0] for e in element_type_binding_pairs]), pb.TensorFlow.Binding(tuple=pb.TensorFlow.NamedTupleBinding( element=[e[1] for e in element_type_binding_pairs]))) elif isinstance(result, DATASET_REPRESENTATION_TYPES): element_type = type_utils.tf_dtypes_and_shapes_to_type( result.output_types, result.output_shapes) handle_name = result.make_one_shot_iterator().string_handle().name return (computation_types.SequenceType(element_type), pb.TensorFlow.Binding(sequence=pb.TensorFlow.SequenceBinding( iterator_string_handle_name=handle_name))) else: raise TypeError( 'Cannot capture a result of an unsupported type {}.'.format( py_typecheck.type_string(type(result))))
def to_representation_for_type(value, type_spec, callable_handler=None): """Verifies or converts the `value` representation to match `type_spec`. This method first tries to determine whether `value` is a valid representation of TFF type `type_spec`. If so, it is returned unchanged. If not, but if it can be converted into a valid representation, it is converted to such, and the valid representation is returned. If no conversion to a valid representation is possible, TypeError is raised. The accepted forms of `value` for various TFF types are as follows: * For TFF tensor types listed in `tensorflow_utils.TENSOR_REPRESENTATION_TYPES`. * For TFF named tuple types, instances of `anonymous_tuple.AnonymousTuple`. * For TFF sequences, Python lists. * For TFF functional types, Python callables that accept a single argument that is an instance of `ComputedValue` (if the function has a parameter) or `None` (otherwise), and return a `ComputedValue` instance as a result. This function only verifies that `value` is a callable. * For TFF abstract types, there is no valid representation. The reference executor requires all types in an executable computation to be concrete. * For TFF placement types, the valid representations are the placement literals (currently only `tff.SERVER` and `tff.CLIENTS`). * For TFF federated types with `all_equal` set to `True`, the representation is the same as the representation of the member constituent (thus, e.g., a valid representation of `int32@SERVER` is the same as that of `int32`). For those types that have `all_equal_` set to `False`, the representation is a Python list of member constituents. NOTE: This function does not attempt at validating that the sizes of lists that represent federated values match the corresponding placemenets. The cardinality analysis is a separate step, handled by the reference executor at a different point. As long as values can be packed into a Python list, they are accepted as they are. Args: value: The raw representation of a value to compare against `type_spec` and potentially to be converted into a canonical form for the given TFF type. type_spec: The TFF type, an instance of `tff.Type` or something convertible to it that determines what the valid representation should be. callable_handler: The function to invoke to handle TFF functional types. If this is `None`, functional types are not supported. The function must accept `value` and `type_spec` as arguments and return the converted valid representation, just as `to_representation_for_type`. Returns: Either `value` itself, or the `value` converted into a valid representation for `type_spec`. Raises: TypeError: If `value` is not a valid representation for given `type_spec`. NotImplementedError: If verification for `type_spec` is not supported. """ type_spec = computation_types.to_type(type_spec) py_typecheck.check_type(type_spec, computation_types.Type) if callable_handler is not None: py_typecheck.check_callable(callable_handler) # NOTE: We do not simply call `type_utils.infer_type()` on `value`, as the # representations of values in the reference executor are only a subset of # the Python types recognized by that helper function. if isinstance(type_spec, computation_types.TensorType): if tf.executing_eagerly() and isinstance(value, (tf.Tensor, tf.Variable)): value = value.numpy() py_typecheck.check_type(value, tensorflow_utils.TENSOR_REPRESENTATION_TYPES) inferred_type_spec = type_utils.infer_type(value) if not type_utils.is_assignable_from(type_spec, inferred_type_spec): raise TypeError( 'The tensor type {} of the value representation does not match ' 'the type spec {}.'.format(inferred_type_spec, type_spec)) return value elif isinstance(type_spec, computation_types.NamedTupleType): type_spec_elements = anonymous_tuple.to_elements(type_spec) # Special-casing unodered dictionaries to allow their elements to be fed in # the order in which they're defined in the named tuple type. if (isinstance(value, dict) and (set(value.keys()) == set(k for k, _ in type_spec_elements))): value = collections.OrderedDict([ (k, value[k]) for k, _ in type_spec_elements ]) value = anonymous_tuple.from_container(value) value_elements = anonymous_tuple.to_elements(value) if len(value_elements) != len(type_spec_elements): raise TypeError( 'The number of elements {} in the value tuple {} does not match the ' 'number of elements {} in the type spec {}.'.format( len(value_elements), value, len(type_spec_elements), type_spec)) result_elements = [] for index, (type_elem_name, type_elem) in enumerate(type_spec_elements): value_elem_name, value_elem = value_elements[index] if value_elem_name not in [type_elem_name, None]: raise TypeError( 'Found element named `{}` where `{}` was expected at position {} ' 'in the value tuple. Value: {}. Type: {}'.format( value_elem_name, type_elem_name, index, value, type_spec)) converted_value_elem = to_representation_for_type(value_elem, type_elem, callable_handler) result_elements.append((type_elem_name, converted_value_elem)) return anonymous_tuple.AnonymousTuple(result_elements) elif isinstance(type_spec, computation_types.SequenceType): if isinstance(value, tf.data.Dataset): inferred_type_spec = computation_types.SequenceType( computation_types.to_type(tf.data.experimental.get_structure(value))) if not type_utils.is_assignable_from(type_spec, inferred_type_spec): raise TypeError( 'Value of type {!s} not assignable to expected type {!s}'.format( inferred_type_spec, type_spec)) if tf.executing_eagerly(): return [ to_representation_for_type(v, type_spec.element, callable_handler) for v in value ] else: raise ValueError( 'Processing `tf.data.Datasets` outside of eager mode is not ' 'currently supported.') return [ to_representation_for_type(v, type_spec.element, callable_handler) for v in value ] elif isinstance(type_spec, computation_types.FunctionType): if callable_handler is not None: return callable_handler(value, type_spec) else: raise TypeError( 'Values that are callables have been explicitly disallowed ' 'in this context. If you would like to supply here a function ' 'as a parameter, please construct a computation that contains ' 'this call.') elif isinstance(type_spec, computation_types.AbstractType): raise TypeError( 'Abstract types are not supported by the reference executor.') elif isinstance(type_spec, computation_types.PlacementType): py_typecheck.check_type(value, placement_literals.PlacementLiteral) return value elif isinstance(type_spec, computation_types.FederatedType): if type_spec.all_equal: return to_representation_for_type(value, type_spec.member, callable_handler) elif type_spec.placement is not placements.CLIENTS: raise TypeError( 'Unable to determine a valid value representation for a federated ' 'type with non-equal members placed at {}.'.format( type_spec.placement)) elif not isinstance(value, (list, tuple)): raise ValueError('Please pass a list or tuple to any function that' ' expects a federated type placed at {};' ' you passed {}'.format(type_spec.placement, value)) else: return [ to_representation_for_type(v, type_spec.member, callable_handler) for v in value ] else: raise NotImplementedError( 'Unable to determine valid value representation for {} for what ' 'is currently an unsupported TFF type {}.'.format(value, type_spec))
def fit_argument(arg, type_spec, context): """Fits the given argument `arg` to match the given parameter `type_spec`. Args: arg: The argument to fit, an instance of `ComputedValue`. type_spec: The type of the parameter to fit to, an instance of `tff.Type` or something convertible to it. context: The context in which to perform the fitting, either an instance of `ComputationContext`, or `None` if unspecified. Returns: An instance of `ComputationValue` with the payload from `arg`, but matching the `type_spec` in the given context. Raises: TypeError: If the types mismatch. ValueError: If the value is invalid or does not fit the requested type. """ py_typecheck.check_type(arg, ComputedValue) type_spec = computation_types.to_type(type_spec) py_typecheck.check_type(type_spec, computation_types.Type) if context is not None: py_typecheck.check_type(context, ComputationContext) type_utils.check_assignable_from(type_spec, arg.type_signature) if arg.type_signature == type_spec: return arg elif isinstance(type_spec, computation_types.NamedTupleType): py_typecheck.check_type(arg.value, anonymous_tuple.AnonymousTuple) result_elements = [] for idx, (elem_name, elem_type) in enumerate(anonymous_tuple.to_elements(type_spec)): elem_val = ComputedValue(arg.value[idx], arg.type_signature[idx]) if elem_val != elem_type: elem_val = fit_argument(elem_val, elem_type, context) result_elements.append((elem_name, elem_val.value)) return ComputedValue( anonymous_tuple.AnonymousTuple(result_elements), type_spec) elif isinstance(type_spec, computation_types.FederatedType): type_utils.check_federated_type( arg.type_signature, placement=type_spec.placement) if arg.type_signature.all_equal: member_val = ComputedValue(arg.value, arg.type_signature.member) if type_spec.member != arg.type_signature.member: member_val = fit_argument(member_val, type_spec.member, context) if type_spec.all_equal: return ComputedValue(member_val.value, type_spec) else: cardinality = context.get_cardinality(type_spec.placement) return ComputedValue([member_val.value for _ in range(cardinality)], type_spec) elif type_spec.all_equal: raise TypeError('Cannot fit a non all-equal {} into all-equal {}.'.format( arg.type_signature, type_spec)) else: py_typecheck.check_type(arg.value, list) def _fit_member_val(x): x_val = ComputedValue(x, arg.type_signature.member) return fit_argument(x_val, type_spec.member, context).value return ComputedValue([_fit_member_val(x) for x in arg.value], type_spec) else: # TODO(b/113123634): Possibly add more conversions, e.g., for tensor types. return arg
def _transform(comp): """Returns a new transformed computation or `comp`.""" if not _should_transform(comp): return comp def _get_comps(comp): """Constructs a 2 dimentional Python list of computations. Args: comp: A `computation_building_blocks.Tuple` containing `n` called intrinsics with `m` arguments. Returns: A 2 dimentional Python list of computations. """ first_call = comp[0] comps = [[] for _ in range(len(first_call.argument))] for _, call in anonymous_tuple.to_elements(comp): for index, arg in enumerate(call.argument): comps[index].append(arg) return comps def _create_block_to_calls(call_names, comps): r"""Constructs a transformed block computation from `comps`. Given the "original" computation containing `n` called intrinsics with `m` arguments, this function constructs the following computation: Block / \ fn=Tuple Lambda(arg) | \ [Comp(f1), Comp(f2), ...] Tuple | [Call, Call, ...] / \ / \ Sel(0) Sel(0) Sel(1) Sel(1) / / / / Ref(fn) Ref(arg) Ref(fn) Ref(arg) with one `computation_building_blocks.Call` for each `n`. This computation represents one of `m` arguments that should be passed to the call of the "transformed" computation. Args: call_names: a Python list of names. comps: a Python list of computations. Returns: A `computation_building_blocks.Block`. """ functions = computation_building_blocks.Tuple( zip(call_names, comps)) fn = computation_building_blocks.Reference( 'fn', functions.type_signature) arg_type = [element.type_signature.parameter for element in comps] arg = computation_building_blocks.Reference('arg', arg_type) elements = [] for index, name in enumerate(call_names): sel_fn = computation_building_blocks.Selection(fn, index=index) sel_arg = computation_building_blocks.Selection(arg, index=index) call = computation_building_blocks.Call(sel_fn, sel_arg) elements.append((name, call)) calls = computation_building_blocks.Tuple(elements) lam = computation_building_blocks.Lambda(arg.name, arg.type_signature, calls) return computation_building_blocks.Block([('fn', functions)], lam) def _create_transformed_args_from_comps(call_names, elements): """Constructs a Python list of transformed computations. Given the "original" computation containing `n` called intrinsics with `m` arguments, this function constructs the following Python list of computations: [Block, Tuple, ...] with one `computation_building_blocks.Block` for each functional computation in `m` and one `computation_building_blocks.Tuple` for each non-functional computation in `m`. This list of computations represent the arguments that should be passed to the `computation_building_blocks.Call` of the "transformed" computation. Args: call_names: a Python list of names. elements: A 2 dimentional Python list of computations. Returns: A Python list of computations. """ args = [] for comps in elements: if isinstance(comps[0].type_signature, computation_types.FunctionType): arg = _create_block_to_calls(call_names, comps) else: arg = computation_building_blocks.Tuple( zip(call_names, comps)) args.append(arg) return args elements = anonymous_tuple.to_elements(comp) call_names = [name for name, _ in elements] comps = _get_comps(comp) args = _create_transformed_args_from_comps(call_names, comps) arg = computation_building_blocks.Tuple(args) parameter_type = computation_types.to_type(arg.type_signature) result_type = [(name, call.type_signature) for name, call in elements] intrinsic_type = computation_types.FunctionType( parameter_type, result_type) intrinsic = computation_building_blocks.Intrinsic( comp[0].function.uri, intrinsic_type) return computation_building_blocks.Call(intrinsic, arg)
def test_construction_from_ordereddict(self): v = collections.OrderedDict(a=1, b=2, c=3) x = anonymous_tuple.AnonymousTuple(v.items()) self.assertSequenceEqual(anonymous_tuple.to_elements(x), list(v.items()))
def append_to_list_structure_for_element_type_spec(structure, value, type_spec): """Adds an element `value` to a nested `structure` of lists for `type_spec`. This function appends tensor-level constituents of an element `value` to the lists created by `make_empty_list_structure_for_element_type_spec`. The nested structure of `value` must match that created by the above function, and consistent with `type_spec`. Args: structure: Output of `make_empty_list_structure_for_element_type_spec`. value: A value (Python object) that a hierarchical structure of dictionary, list, and other containers holding tensor-like items that matches the hierarchy of `type_spec`. type_spec: An instance of `tff.Type` or something convertible to it, as in `make_empty_list_structure_for_element_type_spec`. Raises: TypeError: If the `type_spec` is not of a form described above, or the value is not of a type compatible with `type_spec`. """ if value is None: return type_spec = computation_types.to_type(type_spec) py_typecheck.check_type(type_spec, computation_types.Type) # TODO(b/113116813): This could be made more efficient, but for now we won't # need to worry about it as this is an odd corner case. if isinstance(value, anonymous_tuple.AnonymousTuple): value = collections.OrderedDict(anonymous_tuple.to_elements(value)) if isinstance(type_spec, computation_types.TensorType): py_typecheck.check_type(structure, list) structure.append(value) elif isinstance(type_spec, computation_types.NamedTupleType): elements = anonymous_tuple.to_elements(type_spec) if isinstance(structure, collections.OrderedDict): if py_typecheck.is_named_tuple(value): value = value._asdict() if isinstance(value, dict): if set(value.keys()) != set(k for k, _ in elements): raise TypeError('Value {} does not match type {}.'.format( str(value), str(type_spec))) for elem_name, elem_type in elements: append_to_list_structure_for_element_type_spec( structure[elem_name], value[elem_name], elem_type) elif isinstance(value, (list, tuple)): if len(value) != len(elements): raise TypeError('Value {} does not match type {}.'.format( str(value), str(type_spec))) for idx, (elem_name, elem_type) in enumerate(elements): append_to_list_structure_for_element_type_spec( structure[elem_name], value[idx], elem_type) else: raise TypeError( 'Unexpected type of value {} for TFF type {}.'.format( py_typecheck.type_string(type(value)), str(type_spec))) elif isinstance(structure, tuple): py_typecheck.check_type(value, (list, tuple)) if len(value) != len(elements): raise TypeError('Value {} does not match type {}.'.format( str(value), str(type_spec))) for idx, (_, elem_type) in enumerate(elements): append_to_list_structure_for_element_type_spec( structure[idx], value[idx], elem_type) else: raise TypeError( 'Invalid nested structure, unexpected container type {}.'. format(py_typecheck.type_string(type(structure)))) else: raise TypeError( 'Expected a tensor or named tuple type, found {}.'.format( str(type_spec)))
def test_construction_from_generator_expression(self): x = anonymous_tuple.AnonymousTuple( (name, i) for i, name in enumerate(('a', 'b', None))) self.assertSequenceEqual( anonymous_tuple.to_elements(x), [('a', 0), ('b', 1), (None, 2)])
def flatten_first_index(apply_fn, type_to_add, context_stack): """Returns a value `(arg -> APPEND(apply_fn(arg[0]), arg[1]))`. In the above, `APPEND(a,b)` refers to appending element b to tuple a. Constructs a Value of a TFF functional type that: 1. Takes as argument a 2-element tuple `(x, y)` of TFF type `[apply_fn.type_signature.parameter, type_to_add]`. 2. Transforms the 1st element `x` of this 2-tuple by applying `apply_fn`, producing a result `z` that must be a TFF tuple (e.g, as a result of flattening `x`). 3. Leaves the 2nd element `y` of the argument 2-tuple unchanged. 4. Returns the result of appending the unchanged `y` at the end of the tuple `z` returned by `apply_fn`. Args: apply_fn: TFF `Value` of type_signature `FunctionType`, a function taking TFF `Value`s to `Value`s of type `NamedTupleType`. type_to_add: 2-tuple specifying name and TFF type of arg[1]. Name can be `None` or `string`. context_stack: The context stack to use, as in `impl.value_impl.to_value`. Returns: TFF `Value` of `FunctionType`, taking 2-tuples to N-tuples, which calls `apply_fn` on the first index of its argument, appends the second index to the resulting (N-1)-tuple, then returns the N-tuple thus created. """ py_typecheck.check_type(apply_fn, value_base.Value) py_typecheck.check_type(apply_fn.type_signature, computation_types.FunctionType) py_typecheck.check_type(apply_fn.type_signature.result, computation_types.NamedTupleType) py_typecheck.check_type(type_to_add, tuple) if len(type_to_add) != 2: raise ValueError( 'Please pass a 2-tuple as type_to_add to ' 'flatten_first_index, with first index name or None ' 'and second index instance of `computation_types.Type` ' 'or something convertible to one by ' '`computationtypes.to_type`.') prev_param_type = apply_fn.type_signature.parameter inputs = value_impl.to_value( computation_building_blocks.Reference( 'inputs', computation_types.NamedTupleType([prev_param_type, type_to_add])), None, context_stack) intermediate = apply_fn(inputs[0]) full_type_spec = anonymous_tuple.to_elements( apply_fn.type_signature.result) + [type_to_add] named_values = [(full_type_spec[k][0], intermediate[k]) for k in range(len(intermediate)) ] + [(full_type_spec[-1][0], inputs[1])] new_elements = value_impl.to_value( anonymous_tuple.AnonymousTuple(named_values), type_spec=full_type_spec, context_stack=context_stack) return value_impl.to_value( computation_building_blocks.Lambda( 'inputs', inputs.type_signature, value_impl.ValueImpl.get_comp(new_elements)), None, context_stack)
def transform_postorder(comp, transform): """Traverses `comp` recursively postorder and replaces its constituents. For each element of `comp` viewed as an expression tree, the transformation `transform` is applied first to building blocks it is parameterized by, then the element itself. The transformation `transform` should act as an identity function on the kinds of elements (computation building blocks) it does not care to transform. This corresponds to a post-order traversal of the expression tree, i.e., parameters are alwaysd transformed left-to-right (in the order in which they are listed in building block constructors), then the parent is visited and transformed with the already-visited, and possibly transformed arguments in place. NOTE: In particular, in `Call(f,x)`, both `f` and `x` are arguments to `Call`. Therefore, `f` is transformed into `f'`, next `x` into `x'` and finally, `Call(f',x')` is transformed at the end. Args: comp: The computation to traverse and transform bottom-up. transform: The transformation to apply locally to each building block in `comp`. It is a Python function that accepts a building block at input, and should return either the same, or transformed building block as output. Both the input and output of `transform` are instances of `ComputationBuildingBlock`. Returns: The result of applying `transform` to parts of `comp` in a bottom-up fashion. Raises: TypeError: If the arguments are of the wrong computation_types. NotImplementedError: If the argument is a kind of computation building block that is currently not recognized. """ py_typecheck.check_type( comp, computation_building_blocks.ComputationBuildingBlock) if isinstance(comp, (computation_building_blocks.CompiledComputation, computation_building_blocks.Data, computation_building_blocks.Intrinsic, computation_building_blocks.Placement, computation_building_blocks.Reference)): return transform(comp) elif isinstance(comp, computation_building_blocks.Selection): return transform( computation_building_blocks.Selection( transform_postorder(comp.source, transform), comp.name, comp.index)) elif isinstance(comp, computation_building_blocks.Tuple): return transform( computation_building_blocks.Tuple([ (k, transform_postorder(v, transform)) for k, v in anonymous_tuple.to_elements(comp) ])) elif isinstance(comp, computation_building_blocks.Call): transformed_transform = transform_postorder(comp.function, transform) if comp.argument is not None: transformed_arg = transform_postorder(comp.argument, transform) else: transformed_arg = None return transform( computation_building_blocks.Call(transformed_transform, transformed_arg)) elif isinstance(comp, computation_building_blocks.Lambda): transformed_result = transform_postorder(comp.result, transform) return transform( computation_building_blocks.Lambda(comp.parameter_name, comp.parameter_type, transformed_result)) elif isinstance(comp, computation_building_blocks.Block): return transform( computation_building_blocks.Block( [(k, transform_postorder(v, transform)) for k, v in comp.locals], transform_postorder(comp.result, transform))) else: raise NotImplementedError( 'Unrecognized computation building block: {}'.format(str(comp)))
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, tensorflow_utils.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(tf.data.experimental.get_structure(value))) if not type_utils.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_utils.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'))
async def create_value(self, value, type_spec=None): type_spec = computation_types.to_type(type_spec) py_typecheck.check_type(type_spec, computation_types.Type) if isinstance(value, intrinsic_defs.IntrinsicDef): if not type_utils.is_concrete_instance_of(type_spec, value.type_signature): raise TypeError( 'Incompatible type {} used with intrinsic {}.'.format( type_spec, value.uri)) else: return CompositeValue(value, type_spec) elif isinstance(value, pb.Computation): which_computation = value.WhichOneof('computation') if which_computation in ['tensorflow', 'lambda']: return CompositeValue(value, type_spec) 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) else: raise NotImplementedError( 'Unimplemented computation type {}.'.format( which_computation)) elif isinstance(type_spec, computation_types.NamedTupleType): v_el = anonymous_tuple.to_elements( anonymous_tuple.from_container(value)) t_el = anonymous_tuple.to_elements(type_spec) items = await asyncio.gather(*[ self.create_value(v, t) for (_, v), (_, t) in zip(v_el, t_el) ]) return self.create_tuple( anonymous_tuple.AnonymousTuple([ (k, i) for (k, _), i in zip(t_el, items) ])) elif isinstance(type_spec, computation_types.FederatedType): if type_spec.placement == placement_literals.SERVER: if type_spec.all_equal: return CompositeValue( await self._parent_executor.create_value( value, type_spec.member), type_spec) else: raise ValueError( 'A non-all_equal value on the server is unexpected.') elif type_spec.placement == placement_literals.CLIENTS: if type_spec.all_equal: return CompositeValue( await asyncio.gather(*[ c.create_value(value, type_spec) for c in self._child_executors ]), type_spec) else: py_typecheck.check_type(value, list) if self._cardinalities is None: self._cardinalities = asyncio.ensure_future( self._get_cardinalities()) cardinalities = await self._cardinalities py_typecheck.check_len(cardinalities, len(self._child_executors)) count = sum(cardinalities) py_typecheck.check_len(value, count) result = [] offset = 0 for c, n in zip(self._child_executors, cardinalities): new_offset = offset + n result.append( c.create_value(value[offset:new_offset], type_spec)) offset = new_offset return CompositeValue(await asyncio.gather(*result), type_spec) else: raise ValueError('Unexpected placement {}.'.format( type_spec.placement)) else: return CompositeValue( await self._parent_executor.create_value(value, type_spec), type_spec)
def transform_postorder(comp, transform): """Traverses `comp` recursively postorder and replaces its constituents. For each element of `comp` viewed as an expression tree, the transformation `transform` is applied first to building blocks it is parameterized by, then the element itself. The transformation `transform` should act as an identity function on the kinds of elements (computation building blocks) it does not care to transform. This corresponds to a post-order traversal of the expression tree, i.e., parameters are alwaysd transformed left-to-right (in the order in which they are listed in building block constructors), then the parent is visited and transformed with the already-visited, and possibly transformed arguments in place. NOTE: In particular, in `Call(f,x)`, both `f` and `x` are arguments to `Call`. Therefore, `f` is transformed into `f'`, next `x` into `x'` and finally, `Call(f',x')` is transformed at the end. Args: comp: A `computation_building_block.ComputationBuildingBlock` to traverse and transform bottom-up. transform: The transformation to apply locally to each building block in `comp`. It is a Python function that accepts a building block at input, and should return a (building block, bool) tuple as output, where the building block is a `computation_building_block.ComputationBuildingBlock` representing either the original building block or a transformed building block and the bool is a flag indicating if the building block was modified as. Returns: The result of applying `transform` to parts of `comp` in a bottom-up fashion, along with a Boolean with the value `True` if `comp` was transformed and `False` if it was not. Raises: TypeError: If the arguments are of the wrong computation_types. NotImplementedError: If the argument is a kind of computation building block that is currently not recognized. """ py_typecheck.check_type( comp, computation_building_blocks.ComputationBuildingBlock) if isinstance(comp, (computation_building_blocks.CompiledComputation, computation_building_blocks.Data, computation_building_blocks.Intrinsic, computation_building_blocks.Placement, computation_building_blocks.Reference)): return transform(comp) elif isinstance(comp, computation_building_blocks.Selection): source, source_modified = transform_postorder(comp.source, transform) if source_modified: comp = computation_building_blocks.Selection( source, comp.name, comp.index) comp, comp_modified = transform(comp) return comp, comp_modified or source_modified elif isinstance(comp, computation_building_blocks.Tuple): elements = [] elements_modified = False for key, value in anonymous_tuple.to_elements(comp): value, value_modified = transform_postorder(value, transform) elements.append((key, value)) elements_modified = elements_modified or value_modified if elements_modified: comp = computation_building_blocks.Tuple(elements) comp, comp_modified = transform(comp) return comp, comp_modified or elements_modified elif isinstance(comp, computation_building_blocks.Call): fn, fn_modified = transform_postorder(comp.function, transform) if comp.argument is not None: arg, arg_modified = transform_postorder(comp.argument, transform) else: arg, arg_modified = (None, False) if fn_modified or arg_modified: comp = computation_building_blocks.Call(fn, arg) comp, comp_modified = transform(comp) return comp, comp_modified or fn_modified or arg_modified elif isinstance(comp, computation_building_blocks.Lambda): result, result_modified = transform_postorder(comp.result, transform) if result_modified: comp = computation_building_blocks.Lambda(comp.parameter_name, comp.parameter_type, result) comp, comp_modified = transform(comp) return comp, comp_modified or result_modified elif isinstance(comp, computation_building_blocks.Block): variables = [] variables_modified = False for key, value in comp.locals: value, value_modified = transform_postorder(value, transform) variables.append((key, value)) variables_modified = variables_modified or value_modified result, result_modified = transform_postorder(comp.result, transform) if variables_modified or result_modified: comp = computation_building_blocks.Block(variables, result) comp, comp_modified = transform(comp) return comp, comp_modified or variables_modified or result_modified else: raise NotImplementedError( 'Unrecognized computation building block: {}'.format(str(comp)))
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, 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): 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.to_elements(arg) ]), context_stack) elif py_typecheck.is_named_tuple(arg): result = to_value(arg._asdict(), None, context_stack) 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 = six.iteritems(arg) else: items = sorted(six.iteritems(arg)) 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, 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 construct_named_tuple_setattr_lambda(named_tuple_signature, name, value_comp): """Constructs a building block for replacing one attribute in a named tuple. Returns an instance of `computation_building_blocks.Lambda` which takes an argument of type `computation_types.NamedTupleType` and returns a `computation_building_blocks.Tuple` which contains all the same elements as the argument, except the attribute `name` now has value `value_comp`. The Lambda constructed is the analogue of Python's `setattr` for the concrete type `named_tuple_signature`. Args: named_tuple_signature: Instance of `computation_types.NamedTupleType`, the type of the argument to the constructed `computation_building_blocks.Lambda`. name: String name of the attribute in the `named_tuple_signature` to replace with `value_comp`. Must be present as a name in `named_tuple_signature; otherwise we will raise an `AttributeError`. value_comp: Instance of `computation_building_blocks.ComputationBuildingBlock`, the value to place as attribute `name` in the argument of the returned function. Returns: An instance of `computation_building_blocks.Block` of functional type representing setting attribute `name` to value `value_comp` in its argument of type `named_tuple_signature`. Raises: TypeError: If the types of the arguments don't match the assumptions above. AttributeError: If `name` is not present as a named element in `named_tuple_signature` """ py_typecheck.check_type(named_tuple_signature, computation_types.NamedTupleType) py_typecheck.check_type(name, six.string_types) py_typecheck.check_type( value_comp, computation_building_blocks.ComputationBuildingBlock) value_comp_placeholder = computation_building_blocks.Reference( 'value_comp_placeholder', value_comp.type_signature) lambda_arg = computation_building_blocks.Reference('lambda_arg', named_tuple_signature) if name not in dir(named_tuple_signature): raise AttributeError( 'There is no such attribute as \'{}\' in this federated tuple. ' 'TFF does not allow for assigning to a nonexistent attribute. ' 'If you want to assign to \'{}\', you must create a new named tuple ' 'containing this attribute.'.format(name, name)) elements = [] for idx, (key, element_type) in enumerate( anonymous_tuple.to_elements(named_tuple_signature)): if key == name: if not type_utils.is_assignable_from(element_type, value_comp.type_signature): raise TypeError( '`setattr` has attempted to set element {} of type {} with incompatible type {}' .format(key, element_type, value_comp.type_signature)) elements.append((key, value_comp_placeholder)) else: elements.append((key, computation_building_blocks.Selection(lambda_arg, index=idx))) return_tuple = computation_building_blocks.Tuple(elements) lambda_to_return = computation_building_blocks.Lambda( lambda_arg.name, named_tuple_signature, return_tuple) symbols = ((value_comp_placeholder.name, value_comp), ) return computation_building_blocks.Block(symbols, lambda_to_return)
def to_representation_for_type(value, type_spec=None, device=None): """Verifies or converts the `value` to an eager objct 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 computtion, 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. 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`. """ if device is not None: py_typecheck.check_type(device, six.string_types) with tf.device(device): return to_representation_for_type(value, type_spec=type_spec, device=None) type_spec = type_utils.reconcile_value_with_type_spec(value, type_spec) if isinstance(value, EagerValue): return value.internal_representation if isinstance(value, executor_value_base.ExecutorValue): raise TypeError( 'Cannot accept a value embedded within a non-eager executor.') if isinstance(value, computation_base.Computation): return to_representation_for_type( computation_impl.ComputationImpl.get_proto(value), type_spec, device) if isinstance(value, pb.Computation): return embed_tensorflow_computation(value, type_spec, device) if isinstance(type_spec, computation_types.TensorType): if not isinstance(value, tf.Tensor): value = tf.constant(value, type_spec.dtype, type_spec.shape) value_type = (computation_types.TensorType(value.dtype.base_dtype, value.shape)) if not type_utils.is_assignable_from(type_spec, value_type): raise TypeError( 'The apparent type {} of a tensor {} does not match the expected ' 'type {}.'.format(str(value_type), str(value), str(type_spec))) return value 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( str(len(type_elem)), str(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, el_type, device) result_elem.append((t_name, el_repr)) return anonymous_tuple.AnonymousTuple(result_elem) elif isinstance(type_spec, computation_types.SequenceType): py_typecheck.check_type(value, (tf.data.Dataset, tf.compat.v1.data.Dataset, tf.compat.v2.data.Dataset)) element_type = type_utils.tf_dtypes_and_shapes_to_type( tf.compat.v1.data.get_output_types(value), tf.compat.v1.data.get_output_shapes(value)) value_type = computation_types.SequenceType(element_type) if not type_utils.are_equivalent_types(value_type, type_spec): raise TypeError('Expected a value of type {}, found {}.'.format( str(type_spec), str(value_type))) return value else: raise TypeError('Unexpected type {}.'.format(str(type_spec)))
def _create_fn_to_append_chain_zipped_values(value): r"""Creates a function to append a chain of zipped values. Lambda(arg3) \ append([Call, Sel(1)]) / \ \ Lambda(arg2) Sel(0) Ref(arg3) \ \ \ Ref(arg3) \ append([Call, Sel(1)]) / \ \ Lambda(arg1) Sel(0) Ref(arg2) \ \ \ Ref(arg2) \ Ref(arg1) NOTE: This function is intended to be used in conjunction with `_create_chain_zipped_values` add will add back the names that were dropped when zipping the values. Args: value: A `computation_building_blocks.ComputationBuildingBlock` with a `type_signature` of type `computation_types.NamedTupleType` containing at least two elements. Returns: A `computation_building_blocks.Call`. Raises: TypeError: If any of the types do not match. """ py_typecheck.check_type( value, computation_building_blocks.ComputationBuildingBlock) named_type_signatures = anonymous_tuple.to_elements(value.type_signature) length = len(named_type_signatures) if length < 2: raise ValueError( 'Expected a value with at least two elements, received {} elements.' .format(named_type_signatures)) first_name, first_type_signature = named_type_signatures[0] second_name, second_type_signature = named_type_signatures[1] ref_type = computation_types.NamedTupleType(( (first_name, first_type_signature.member), (second_name, second_type_signature.member), )) ref = computation_building_blocks.Reference('arg', ref_type) fn = computation_building_blocks.Lambda(ref.name, ref.type_signature, ref) for name, type_signature in named_type_signatures[2:]: ref_type = computation_types.NamedTupleType(( fn.type_signature.parameter, (name, type_signature.member), )) ref = computation_building_blocks.Reference('arg', ref_type) sel_0 = computation_building_blocks.Selection(ref, index=0) call = computation_building_blocks.Call(fn, sel_0) sel_1 = computation_building_blocks.Selection(ref, index=1) result = create_computation_appending(call, (name, sel_1)) fn = computation_building_blocks.Lambda(ref.name, ref.type_signature, result) return fn
def type_to_tf_dtypes_and_shapes(type_spec): """Returns nested structures of tensor dtypes and shapes for a given TFF type. The returned dtypes and shapes match those used by `tf.data.Dataset`s to indicate the type and shape of their elements. They can be used, e.g., as arguments in constructing an iterator over a string handle. Args: type_spec: Type specification, either an instance of `computation_types.Type`, or something convertible to it. Ther type specification must be composed of only named tuples and tensors. In all named tuples that appear in the type spec, all the elements must be named. Returns: A pair of parallel nested structures with the dtypes and shapes of tensors defined in `type_spec`. The layout of the two structures returned is the same as the layout of the nesdted type defined by `type_spec`. Named tuples are represented as dictionaries. Raises: ValueError: if the `type_spec` is composed of something other than named tuples and tensors, or if any of the elements in named tuples are unnamed. """ type_spec = computation_types.to_type(type_spec) if isinstance(type_spec, computation_types.TensorType): return (type_spec.dtype, type_spec.shape) elif isinstance(type_spec, computation_types.NamedTupleType): elements = anonymous_tuple.to_elements(type_spec) if elements[0][0] is not None: output_dtypes = collections.OrderedDict() output_shapes = collections.OrderedDict() for e in elements: element_name = e[0] element_spec = e[1] if element_name is None: raise ValueError( 'When a sequence appears as a part of a parameter to a section ' 'of TensorFlow code, in the type signature of elements of that ' 'sequence all named tuples must have their elements explicitly ' 'named, and this does not appear to be the case in {}.' .format(str(type_spec))) element_output = type_to_tf_dtypes_and_shapes(element_spec) output_dtypes[element_name] = element_output[0] output_shapes[element_name] = element_output[1] else: output_dtypes = [] output_shapes = [] for e in elements: element_name = e[0] element_spec = e[1] if element_name is not None: raise ValueError( 'When a sequence appears as a part of a parameter to a section ' 'of TensorFlow code, in the type signature of elements of that ' 'sequence all named tuples must have their elements explicitly ' 'named, and this does not appear to be the case in {}.' .format(str(type_spec))) element_output = type_to_tf_dtypes_and_shapes(element_spec) output_dtypes.append(element_output[0]) output_shapes.append(element_output[1]) return (output_dtypes, output_shapes) else: raise ValueError('Unsupported type {}.'.format( py_typecheck.type_string(type(type_spec))))
def test_construction_from_tuple(self): v = (('a', 1), ('b', 2), (None, 3)) x = anonymous_tuple.AnonymousTuple(v) self.assertSequenceEqual(anonymous_tuple.to_elements(x), v)
def is_assignable_from(target_type, source_type): """Determines whether `target_type` is assignable from `source_type`. Args: target_type: The expected type (that of the target of the assignment). source_type: The actual type (that of the source of the assignment), tested for being a specialization of the `target_type`. Returns: `True` iff `target_type` is assignable from `source_type`, or else `False`. Raises: TypeError: If the arguments are not TFF types. """ target_type = computation_types.to_type(target_type) source_type = computation_types.to_type(source_type) py_typecheck.check_type(target_type, computation_types.Type) py_typecheck.check_type(source_type, computation_types.Type) if isinstance(target_type, computation_types.TensorType): def _shape_is_assignable_from(x, y): def _dimension_is_assignable_from(x, y): return (x.value is None) or (x.value == y.value) # TODO(b/123764922): See if we can pass to TensorShape's # `is_compatible_with`. return ((x.ndims == y.ndims) and ((x.dims is None) or all( _dimension_is_assignable_from(x.dims[k], y.dims[k]) for k in range(x.ndims)))) or x.ndims is None return (isinstance(source_type, computation_types.TensorType) and (target_type.dtype == source_type.dtype) and _shape_is_assignable_from(target_type.shape, source_type.shape)) elif isinstance(target_type, computation_types.NamedTupleType): if not isinstance(source_type, computation_types.NamedTupleType): return False target_elements = anonymous_tuple.to_elements(target_type) source_elements = anonymous_tuple.to_elements(source_type) return ((len(target_elements) == len(source_elements)) and all( ((source_elements[k][0] in [target_elements[k][0], None]) and is_assignable_from(target_elements[k][1], source_elements[k][1])) for k in range(len(target_elements)))) elif isinstance(target_type, computation_types.SequenceType): return (isinstance(source_type, computation_types.SequenceType) and is_assignable_from(target_type.element, source_type.element)) elif isinstance(target_type, computation_types.FunctionType): return (isinstance(source_type, computation_types.FunctionType) and (((source_type.parameter is None) and (target_type.parameter is None)) or ((source_type.parameter is not None) and (target_type.parameter is not None) and is_assignable_from( source_type.parameter, target_type.parameter)) and is_assignable_from(target_type.result, source_type.result))) elif isinstance(target_type, computation_types.AbstractType): # TODO(b/113112108): Revise this to extend the relation of assignability to # abstract types. raise TypeError('Abstract types are not comparable.') elif isinstance(target_type, computation_types.PlacementType): return isinstance(source_type, computation_types.PlacementType) elif isinstance(target_type, computation_types.FederatedType): if (not isinstance(source_type, computation_types.FederatedType) or not is_assignable_from(target_type.member, source_type.member) or target_type.all_equal and not source_type.all_equal): return False for val in [target_type, source_type]: py_typecheck.check_type(val.placement, placement_literals.PlacementLiteral) return target_type.placement is source_type.placement else: raise TypeError('Unexpected target type {}.'.format(str(target_type)))
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_utils.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_utils.check_assignable_from(type_spec, value_type) return value else: raise TypeError('Unexpected type {}.'.format(type_spec))