def _visit_type(type_signature, function): def inner_function(inner_type): function(inner_type) return inner_type, False type_transformations.transform_type_postorder(type_signature, inner_function)
def _visit_postorder(type_signature: computation_types.Type, function: Callable[[computation_types.Type], None]): py_typecheck.check_type(type_signature, computation_types.Type) def _visit(inner_type): function(inner_type) return inner_type, False type_transformations.transform_type_postorder(type_signature, _visit)
def test_recurses_under_sequence(self): orig_type = computation_types.SequenceType([tf.int32]) expected_type = computation_types.SequenceType([tf.float32]) result_type, mutated = type_transformations.transform_type_postorder( orig_type, _convert_tensor_to_float) noop_type, not_mutated = type_transformations.transform_type_postorder( orig_type, _convert_abstract_type_to_tensor) self.assertEqual(result_type, expected_type) self.assertEqual(noop_type, orig_type) self.assertTrue(mutated) self.assertFalse(not_mutated)
def test_transforms_unnamed_tuple_type(self): orig_type = computation_types.StructType([tf.int32, tf.float64]) expected_type = computation_types.StructType([tf.float32, tf.float32]) result_type, mutated = type_transformations.transform_type_postorder( orig_type, _convert_tensor_to_float) noop_type, not_mutated = type_transformations.transform_type_postorder( orig_type, _convert_abstract_type_to_tensor) self.assertEqual(result_type, expected_type) self.assertEqual(noop_type, orig_type) self.assertTrue(mutated) self.assertFalse(not_mutated)
def test_transforms_tensor(self): orig_type = computation_types.TensorType(tf.int32) expected_type = computation_types.TensorType(tf.float32) result_type, mutated = type_transformations.transform_type_postorder( orig_type, _convert_tensor_to_float) noop_type, not_mutated = type_transformations.transform_type_postorder( orig_type, _convert_abstract_type_to_tensor) self.assertEqual(result_type, expected_type) self.assertEqual(noop_type, orig_type) self.assertTrue(mutated) self.assertFalse(not_mutated)
def test_transforms_federated_type(self): orig_type = computation_types.FederatedType(tf.int32, placement_literals.CLIENTS) expected_type = computation_types.FederatedType(tf.float32, placement_literals.CLIENTS) result_type, mutated = type_transformations.transform_type_postorder( orig_type, _convert_tensor_to_float) noop_type, not_mutated = type_transformations.transform_type_postorder( orig_type, _convert_abstract_type_to_tensor) self.assertEqual(result_type, expected_type) self.assertEqual(noop_type, orig_type) self.assertTrue(mutated) self.assertFalse(not_mutated)
def test_transforms_named_tuple_type_preserving_tuple_container(self): orig_type = computation_types.NamedTupleTypeWithPyContainerType( [('a', tf.int32), ('b', tf.float64)], dict) expected_type = computation_types.NamedTupleTypeWithPyContainerType( [('a', tf.float32), ('b', tf.float32)], dict) result_type, mutated = type_transformations.transform_type_postorder( orig_type, _convert_tensor_to_float) noop_type, not_mutated = type_transformations.transform_type_postorder( orig_type, _convert_abstract_type_to_tensor) self.assertEqual(result_type, expected_type) self.assertEqual(noop_type, orig_type) self.assertTrue(mutated) self.assertFalse(not_mutated)
def test_recurses_under_named_tuple_type(self): orig_type = computation_types.to_type([[('a', tf.int32), ('b', tf.float64)]]) expected_type = computation_types.to_type([[('a', tf.float32), ('b', tf.float32)]]) result_type, mutated = type_transformations.transform_type_postorder( orig_type, _convert_tensor_to_float) noop_type, not_mutated = type_transformations.transform_type_postorder( orig_type, _convert_abstract_type_to_tensor) self.assertEqual(result_type, expected_type) self.assertEqual(noop_type, orig_type) self.assertTrue(mutated) self.assertFalse(not_mutated)
def is_binary_op_with_upcast_compatible_pair( possibly_nested_type: Optional[computation_types.Type], type_to_upcast: computation_types.Type) -> bool: """Checks unambiguity in applying `type_to_upcast` to `possibly_nested_type`. That is, checks that either these types are equivalent and contain only tuples and tensors, or that `possibly_nested_type` is perhaps a nested structure containing only tensors with `dtype` of `type_to_upcast` at the leaves, where `type_to_upcast` must be a scalar tensor type. Notice that this relationship is not symmetric, since binary operators need not respect this symmetry in general. For example, it makes perfect sence to divide a nested structure of tensors by a scalar, but not the other way around. Args: possibly_nested_type: A `computation_types.Type`, or `None`. type_to_upcast: A `computation_types.Type`, or `None`. Returns: Boolean indicating whether `type_to_upcast` can be upcast to `possibly_nested_type` in the manner described above. """ if possibly_nested_type is not None: py_typecheck.check_type(possibly_nested_type, computation_types.Type) if type_to_upcast is not None: py_typecheck.check_type(type_to_upcast, computation_types.Type) if not (is_generic_op_compatible_type(possibly_nested_type) and is_generic_op_compatible_type(type_to_upcast)): return False if possibly_nested_type is None: return type_to_upcast is None if possibly_nested_type.is_equivalent_to(type_to_upcast): return True if not (isinstance(type_to_upcast, computation_types.TensorType) and type_to_upcast.shape == tf.TensorShape(())): return False types_are_ok = [True] only_allowed_dtype = type_to_upcast.dtype def _check_tensor_types(type_spec): if isinstance(type_spec, computation_types.TensorType ) and type_spec.dtype != only_allowed_dtype: types_are_ok[0] = False return type_spec, False type_transformations.transform_type_postorder(possibly_nested_type, _check_tensor_types) return types_are_ok[0]
async def _embed_tf_secure_sum_mask_value(self, type_spec, extended_bitwidth): """Construct a CompiledComputation with the mask for the secure sum.""" def transform_to_scalar_type_spec(t): """Converts all `tff.TensorType` to scalar shapes.""" if not t.is_tensor(): return t, False return computation_types.TensorType(dtype=t.dtype, shape=[]), True mask_type, _ = type_transformations.transform_type_postorder( type_spec, transform_to_scalar_type_spec) def compute_mask(bits, type_spec): mask_value = 2**bits - 1 if mask_value > type_spec.dtype.max: logging.warning( 'Emulated secure sum mask exceeds maximum value of ' 'dtype. Dtype %s, mask bits: %d', type_spec.dtype, bits) mask_value = type_spec.dtype.max return tf.constant(mask_value, type_spec.dtype) mask_value = _map_numpy_or_structure(extended_bitwidth, mask_type, fn=compute_mask) logging.debug('Emulated secure sum using mask: %s', mask_value) return await executor_utils.embed_tf_constant(self._executor, mask_type, mask_value)
def _check_container_compat_with_tf_nest(type_spec: computation_types.Type): """Asserts that all `StructTypes` with names have OrderedDict containers.""" def _names_are_in_sorted_order(name_sequence: Sequence[str]) -> bool: return sorted(name_sequence) == name_sequence def _check_ordereddict_container_for_struct(type_to_check): if not type_to_check.is_struct(): return type_to_check, False # We can't use `dir` here, since it sorts the names before returning. We # also must filter to names which are actually present. names_in_sequence_order = structure.name_list(type_to_check) names_are_sorted = _names_are_in_sorted_order(names_in_sequence_order) has_no_names = not bool(names_in_sequence_order) if has_no_names or (names_in_sequence_order and names_are_sorted): # If alphabetical order matches sequence order, TFF's deserialization will # traverse the structure correctly; there is no ambiguity here. On the # other hand, if there are no names, sequence order is the only method of # traversal, so there is no ambiguity here either. return type_to_check, False elif not type_to_check.is_struct_with_python(): raise ValueError( 'Attempting to serialize a named struct type with ' 'ambiguous traversal order (sequence order distinct ' 'from alphabetical order) without a Python container; ' 'this is an unsafe operation, as TFF cannot determine ' 'the intended traversal order after deserializing the ' 'proto due to inconsistent behavior of tf.nest.') container_type = computation_types.StructWithPythonType.get_container_type( type_to_check) if (not names_are_sorted ) and container_type is not collections.OrderedDict: raise ValueError( 'Attempted to serialize a dataset yielding named ' 'elements in non-sorted sequence order with ' f'non-OrderedDict container (type {container_type}). ' 'This is an ambiguous operation; `tf.nest` behaves in ' 'a manner which depends on the Python type of this ' 'container, so coercing the dataset reconstructed ' 'from the resulting Value proto depends on assuming a ' 'single Python type here. Please prefer to use ' '`collections.OrderedDict` containers for the elements ' 'your dataset yields.') return type_to_check, False type_transformations.transform_type_postorder( type_spec, _check_ordereddict_container_for_struct)
def _build_reservoir_type( sample_value_type: computation_types.Type) -> computation_types.Type: """Create the TFF type for the reservoir's state. `UnweightedReservoirSamplingFactory` will use this type as the "state" type in a `tff.federated_aggregate` (an input to `accumulate`, `merge` and `report`). Args: sample_value_type: The `tff.Type` of the values that will be aggregated from clients. Returns: A `collection.OrderedDict` with three keys: random_seed: A 2-tuple of `tf.int64` scalars. This keeps track of the `seed` parameter for `tf.random.stateless_uniform` calls for sampling during aggregation. random_values: A 1-d tensor of `int32` values randomly generated from `tf.random.stateless_uniform`. The size of the tensor will be the same as the first dimenion of each leaf in `samples` (these can be thought of as parallel lists). These values are used to determine whether a sample stays in the reservoir, or is evicted, as the values are aggregated. If the i-th value of this list is evicted, then the i-th (in the first dimension) tensor in each of the `samples` structure leaves must be evicted. samples: A tensor or structure of tensors representing the actual sampled values. If a structure, the shape of the structure matches that of `sample_value_type`. All tensors have one additional dimenion prepended which has an unknown size. This will be used to concatenate samples and store them in the reservoir. """ # TODO(b/181365504): relax this to allow `StructType` once a `Struct` can be # returned from `tf.function` decorated methods. def is_tesnor_or_struct_with_py_type(t: computation_types.Type) -> bool: return t.is_tensor() or t.is_struct_with_python() if not type_analysis.contains_only(sample_value_type, is_tesnor_or_struct_with_py_type): raise TypeError( 'Cannot create a reservoir for type structure. Sample type ' 'must only contain `TensorType` or `StructWithPythonType`, ' f'got a {sample_value_type!r}.') def add_unknown_dimension(t): if t.is_tensor(): return (computation_types.TensorType(dtype=t.dtype, shape=[None] + t.shape), True) return t, False # TODO(b/181155367): creating a value from a type for the `zero` is a common # pattern for users of `tff.federated_aggregate` that could be made easier # for TFF users. Replace this once such helper exists. return computation_types.to_type( collections.OrderedDict( random_seed=computation_types.TensorType(tf.int64, shape=[2]), random_values=computation_types.TensorType(tf.int32, shape=[None]), samples=type_transformations.transform_type_postorder( sample_value_type, add_unknown_dimension)[0]))
def _remove_lambda_placement(comp): """Removes placement from Lambda's parameter.""" if comp.parameter_name is None: new_parameter_type = None else: new_parameter_type, _ = type_transformations.transform_type_postorder( comp.parameter_type, _remove_placement_from_type) return building_blocks.Lambda(comp.parameter_name, new_parameter_type, comp.result)
def count_tensors_in_type( type_spec: computation_types.Type, tensor_filter: Optional[Callable[[computation_types.TensorType], bool]] = None ) -> collections.OrderedDict: """Counts tensors and fully-specified elements under `type_spec`. Args: type_spec: Instance of `computation_types.Type` to count tensors under. tensor_filter: Optional filtering function. Callable which takes an argument of type `computation_types.TensorType` and returns a boolean. If specified, only tensor type which pass this filter (IE, on which this function returns `True`) will be counted. Returns: A `collections.OrderedDict` with three parameters. The first, `tensors`, is the count of all `computation_types.TensorType` (passing `tensor_filter` if this argument is specified). The second, `parameters`, is the count of all fully-specified parameters of these tensors. Note that this implies any tensor with a `None` dimension (IE, of unspecified size) will not be counted. The third counts how many tensors fall into this category (that is, now many have unspecified size). """ py_typecheck.check_type(type_spec, computation_types.Type) if tensor_filter is None: tensor_filter = lambda _: True py_typecheck.check_callable(tensor_filter) tensors_and_params = collections.OrderedDict(num_tensors=0, parameters=0, num_unspecified_tensors=0) def _capture_tensors(type_signature): if type_signature.is_tensor() and tensor_filter(type_signature): tensors_and_params['num_tensors'] += 1 num_parameters = type_signature.shape.num_elements() if num_parameters is not None: tensors_and_params['parameters'] += num_parameters else: tensors_and_params['num_unspecified_tensors'] += 1 return type_signature, False type_transformations.transform_type_postorder(type_spec, _capture_tensors) return tensors_and_params
def _remove_client_all_equals_from_type(type_signature): def _transform(inner_type): if (inner_type.is_federated() and inner_type.placement.is_clients() and inner_type.all_equal): return computation_types.FederatedType(inner_type.member, inner_type.placement, False), True return inner_type, False return type_transformations.transform_type_postorder(type_signature, _transform)[0]
def test_raises_on_non_type_first_arg(self): with self.assertRaises(TypeError): type_transformations.transform_type_postorder(tf.int32, None)
def test_raises_on_none_function(self): with self.assertRaises(TypeError): type_transformations.transform_type_postorder( computation_types.TensorType(tf.int32), None)
def test_raises_on_none_type(self): with self.assertRaises(TypeError): type_transformations.transform_type_postorder(None, lambda x: x)
def test_updates_mutated_bit_at_tuple(self): orig_type = computation_types.StructType([tf.int32, tf.float64]) _, mutated = type_transformations.transform_type_postorder( orig_type, _convert_tuple_to_tensor) self.assertTrue(mutated)
def _remove_reference_placement(comp): """Unwraps placement from references and updates unbound reference info.""" new_type, _ = type_transformations.transform_type_postorder( comp.type_signature, _remove_placement_from_type) return building_blocks.Reference(comp.name, new_type)
def test_updates_mutated_bit_at_function(self): orig_type = computation_types.FunctionType(tf.int32, tf.int64) _, mutated = type_transformations.transform_type_postorder( orig_type, _convert_function_to_tensor) self.assertTrue(mutated)
def test_updates_mutated_bit_at_sequence(self): orig_type = computation_types.SequenceType(tf.int32) _, mutated = type_transformations.transform_type_postorder( orig_type, _convert_sequence_to_tensor) self.assertTrue(mutated)
def test_updates_mutated_bit_at_federated(self): orig_type = computation_types.FederatedType(tf.int32, placements.CLIENTS) _, mutated = type_transformations.transform_type_postorder( orig_type, _convert_federated_to_tensor) self.assertTrue(mutated)
def create_constant(scalar_value, type_spec: computation_types.Type) -> ProtoAndType: """Returns a tensorflow computation returning a constant `scalar_value`. The returned computation has the type signature `( -> T)`, where `T` is `type_spec`. `scalar_value` must be a scalar, and cannot be a float if any of the tensor leaves of `type_spec` contain an integer data type. `type_spec` must contain only named tuples and tensor types, but these can be arbitrarily nested. Args: scalar_value: A scalar value to place in all the tensor leaves of `type_spec`. type_spec: A `computation_types.Type` to use as the argument to the constructed binary operator; must contain only named tuples and tensor types. Raises: TypeError: If the constraints of `type_spec` are violated. """ if not type_analysis.is_generic_op_compatible_type(type_spec): raise TypeError( 'Type spec {} cannot be constructed as a TensorFlow constant in TFF; ' ' only nested tuples and tensors are permitted.'.format(type_spec)) inferred_scalar_value_type = type_conversions.infer_type(scalar_value) if (not inferred_scalar_value_type.is_tensor() or inferred_scalar_value_type.shape != tf.TensorShape(())): raise TypeError( 'Must pass a scalar value to `create_tensorflow_constant`; encountered ' 'a value {}'.format(scalar_value)) tensor_dtypes_in_type_spec = [] def _pack_dtypes(type_signature): """Appends dtype of `type_signature` to nonlocal variable.""" if type_signature.is_tensor(): tensor_dtypes_in_type_spec.append(type_signature.dtype) return type_signature, False type_transformations.transform_type_postorder(type_spec, _pack_dtypes) if (any(x.is_integer for x in tensor_dtypes_in_type_spec) and not inferred_scalar_value_type.dtype.is_integer): raise TypeError( 'Only integers can be used as scalar values if our desired constant ' 'type spec contains any integer tensors; passed scalar {} of dtype {} ' 'for type spec {}.'.format(scalar_value, inferred_scalar_value_type.dtype, type_spec)) result_type = type_spec def _create_result_tensor(type_spec, scalar_value): """Packs `scalar_value` into `type_spec` recursively.""" if type_spec.is_tensor(): type_spec.shape.assert_is_fully_defined() result = tf.constant( scalar_value, dtype=type_spec.dtype, shape=type_spec.shape) else: elements = [] for _, type_element in structure.iter_elements(type_spec): elements.append(_create_result_tensor(type_element, scalar_value)) result = elements return result with tf.Graph().as_default() as graph: result = _create_result_tensor(result_type, scalar_value) _, result_binding = tensorflow_utils.capture_result_from_graph( result, graph) type_signature = computation_types.FunctionType(None, result_type) tensorflow = pb.TensorFlow( graph_def=serialization_utils.pack_graph_def(graph.as_graph_def()), parameter=None, result=result_binding) return _tensorflow_comp(tensorflow, type_signature)
def create_constant( value, type_spec: computation_types.Type) -> ComputationProtoAndType: """Returns a tensorflow computation returning a constant `value`. The returned computation has the type signature `( -> T)`, where `T` is `type_spec`. `value` must be a value convertible to a tensor or a structure of values, such that the dtype and shapes match `type_spec`. `type_spec` must contain only named tuples and tensor types, but these can be arbitrarily nested. Args: value: A value to embed as a constant in the tensorflow graph. type_spec: A `computation_types.Type` to use as the argument to the constructed binary operator; must contain only named tuples and tensor types. Raises: TypeError: If the constraints of `type_spec` are violated. """ if not type_analysis.is_generic_op_compatible_type(type_spec): raise TypeError( 'Type spec {} cannot be constructed as a TensorFlow constant in TFF; ' ' only nested tuples and tensors are permitted.'.format(type_spec)) inferred_value_type = type_conversions.infer_type(value) if (inferred_value_type.is_struct() and not type_spec.is_assignable_from(inferred_value_type)): raise TypeError( 'Must pass a only tensor or structure of tensor values to ' '`create_tensorflow_constant`; encountered a value {v} with inferred ' 'type {t!r}, but needed {s!r}'.format(v=value, t=inferred_value_type, s=type_spec)) if inferred_value_type.is_struct(): value = structure.from_container(value, recursive=True) tensor_dtypes_in_type_spec = [] def _pack_dtypes(type_signature): """Appends dtype of `type_signature` to nonlocal variable.""" if type_signature.is_tensor(): tensor_dtypes_in_type_spec.append(type_signature.dtype) return type_signature, False type_transformations.transform_type_postorder(type_spec, _pack_dtypes) if (any(x.is_integer for x in tensor_dtypes_in_type_spec) and (inferred_value_type.is_tensor() and not inferred_value_type.dtype.is_integer)): raise TypeError( 'Only integers can be used as scalar values if our desired constant ' 'type spec contains any integer tensors; passed scalar {} of dtype {} ' 'for type spec {}.'.format(value, inferred_value_type.dtype, type_spec)) result_type = type_spec def _create_result_tensor(type_spec, value): """Packs `value` into `type_spec` recursively.""" if type_spec.is_tensor(): type_spec.shape.assert_is_fully_defined() result = tf.constant(value, dtype=type_spec.dtype, shape=type_spec.shape) else: elements = [] if inferred_value_type.is_struct(): # Copy the leaf values according to the type_spec structure. for (name, elem_type), value in zip( structure.iter_elements(type_spec), value): elements.append( (name, _create_result_tensor(elem_type, value))) else: # "Broadcast" the value to each level of the type_spec structure. for _, elem_type in structure.iter_elements(type_spec): elements.append( (None, _create_result_tensor(elem_type, value))) result = structure.Struct(elements) return result with tf.Graph().as_default() as graph: result = _create_result_tensor(result_type, value) _, result_binding = tensorflow_utils.capture_result_from_graph( result, graph) type_signature = computation_types.FunctionType(None, result_type) tensorflow = pb.TensorFlow(graph_def=serialization_utils.pack_graph_def( graph.as_graph_def()), parameter=None, result=result_binding) return _tensorflow_comp(tensorflow, type_signature)
def _deserialize_dataset_from_graph_def(serialized_graph_def: bytes, element_type: computation_types.Type): """Deserializes a serialized `tf.compat.v1.GraphDef` to a `tf.data.Dataset`. Args: serialized_graph_def: `bytes` object produced by `tensorflow_serialization.serialize_dataset` element_type: a `tff.Type` object representing the type structure of the elements yielded from the dataset. Returns: A `tf.data.Dataset` instance. """ py_typecheck.check_type(element_type, computation_types.Type) type_analysis.check_tensorflow_compatible_type(element_type) def transform_to_tff_known_type( type_spec: computation_types.Type ) -> Tuple[computation_types.Type, bool]: """Transforms `StructType` to `StructWithPythonType`.""" if type_spec.is_struct() and not type_spec.is_struct_with_python(): field_is_named = tuple( name is not None for name, _ in structure.iter_elements(type_spec)) has_names = any(field_is_named) is_all_named = all(field_is_named) if is_all_named: return computation_types.StructWithPythonType( elements=structure.iter_elements(type_spec), container_type=collections.OrderedDict), True elif not has_names: return computation_types.StructWithPythonType( elements=structure.iter_elements(type_spec), container_type=tuple), True else: raise TypeError( 'Cannot represent TFF type in TF because it contains ' f'partially named structures. Type: {type_spec}') return type_spec, False if element_type.is_struct(): # TF doesn't suppor `structure.Strut` types, so we must transform the # `StructType` into a `StructWithPythonType` for use as the # `tf.data.Dataset.element_spec` later. tf_compatible_type, _ = type_transformations.transform_type_postorder( element_type, transform_to_tff_known_type) else: # We've checked this is only a struct or tensors, so we know this is a # `TensorType` here and will use as-is. tf_compatible_type = element_type def type_to_tensorspec(t: computation_types.TensorType) -> tf.TensorSpec: return tf.TensorSpec(shape=t.shape, dtype=t.dtype) element_spec = type_conversions.structure_from_tensor_type_tree( type_to_tensorspec, tf_compatible_type) ds = tf.data.experimental.from_variant( tf.raw_ops.DatasetFromGraph(graph_def=serialized_graph_def), structure=element_spec) # If a serialized dataset had elements of nested structes of tensors (e.g. # `dict`, `OrderedDict`), the deserialized dataset will return `dict`, # `tuple`, or `namedtuple` (loses `collections.OrderedDict` in a conversion). # # Since the dataset will only be used inside TFF, we wrap the dictionary # coming from TF in an `OrderedDict` when necessary (a type that both TF and # TFF understand), using the field order stored in the TFF type stored during # serialization. return tensorflow_utils.coerce_dataset_elements_to_tff_type_spec( ds, tf_compatible_type)