def create_dummy_value_unplaced(): """Returns a Python value and unplaced type.""" value = 10.0 type_signature = computation_types.TensorType(tf.float32) return value, type_signature
def test_make_dummy_element_TensorType(self): type_spec = computation_types.TensorType(tf.float32, [None, 10, None, 10, 10]) elem = tensorflow_utils.make_dummy_element_for_type_spec(type_spec) correct_elem = np.zeros([0, 10, 0, 10, 10], np.float32) self.assertAllClose(elem, correct_elem)
ip = computation_utils.IterativeProcess(init_fn, next_fn) cf = canonical_form_utils.get_canonical_form_for_iterative_process(ip) self.assertIsInstance(cf, canonical_form.CanonicalForm) INIT_TYPE = computation_types.FederatedType(tf.float32, placements.SERVER) S1_TYPE = INIT_TYPE C1_TYPE = computation_types.FederatedType(tf.float32, placements.CLIENTS) S6_TYPE = computation_types.FederatedType(tf.float64, placements.SERVER) S7_TYPE = computation_types.FederatedType(tf.bool, placements.SERVER) C6_TYPE = computation_types.FederatedType(tf.int64, placements.CLIENTS) S2_TYPE = computation_types.FederatedType([tf.float32], placements.SERVER) C2_TYPE = computation_types.FederatedType(S2_TYPE.member, placements.CLIENTS) C5_TYPE = computation_types.FederatedType([tf.float64], placements.CLIENTS) ZERO_TYPE = computation_types.TensorType(tf.int64) ACCUMULATE_TYPE = computation_types.FunctionType([ZERO_TYPE, C5_TYPE.member], ZERO_TYPE) MERGE_TYPE = computation_types.FunctionType([ZERO_TYPE, ZERO_TYPE], ZERO_TYPE) REPORT_TYPE = computation_types.FunctionType(ZERO_TYPE, tf.int64) S3_TYPE = computation_types.FederatedType(REPORT_TYPE.result, placements.SERVER) def _create_next_type_with_s1_type(x): param_type = computation_types.NamedTupleType([x, C1_TYPE]) result_type = computation_types.NamedTupleType([S6_TYPE, S7_TYPE, C6_TYPE]) return computation_types.FunctionType(param_type, result_type) def _create_before_broadcast_type_with_s1_type(x):
def create_dummy_computation_intrinsic(): value = pb.Computation( type=type_serialization.serialize_type(tf.int32), intrinsic=pb.Intrinsic(uri=intrinsic_defs.GENERIC_ZERO.uri)) type_signature = computation_types.TensorType(tf.int32) return value, type_signature
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
def test_init_does_not_raise_type_error_with_unknown_dimensions(self): server_state_type = computation_types.TensorType(shape=[None], dtype=tf.int32) @computations.tf_computation def initialize(): # Return a value of a type assignable to, but not equal to # `server_state_type` return tf.constant([1, 2, 3]) @computations.tf_computation(server_state_type) def prepare(server_state): del server_state # Unused return tf.constant(1.0) @computations.tf_computation(computation_types.SequenceType( tf.float32), tf.float32) def work(client_data, client_input): del client_data # Unused del client_input # Unused return True, [] @computations.tf_computation def zero(): return tf.constant([], dtype=tf.string) @computations.tf_computation( computation_types.TensorType(shape=[None], dtype=tf.string), tf.bool) def accumulate(accumulator, client_update): del accumulator # Unused del client_update # Unused return tf.constant(['abc']) @computations.tf_computation( computation_types.TensorType(shape=[None], dtype=tf.string), computation_types.TensorType(shape=[None], dtype=tf.string)) def merge(accumulator1, accumulator2): del accumulator1 # Unused del accumulator2 # Unused return tf.constant(['abc']) @computations.tf_computation( computation_types.TensorType(shape=[None], dtype=tf.string)) def report(accumulator): del accumulator # Unused return tf.constant(1.0) @computations.tf_computation def bitwidth(): return [] @computations.tf_computation( server_state_type, (tf.float32, computation_types.StructType([]))) def update(server_state, global_update): del server_state # Unused del global_update # Unused # Return a new server state value whose type is assignable but not equal # to `server_state_type`, and which is different from the type returned # by `initialize`. return tf.constant([1]), [] try: forms.MapReduceForm(initialize, prepare, work, zero, accumulate, merge, report, bitwidth, update) except TypeError: self.fail('Raised TypeError unexpectedly.')
def create_dummy_unplaced_type(): value = 10 type_signature = computation_types.TensorType(tf.int32) return value, type_signature
def _convert_federated_to_tensor(type_spec): if isinstance(type_spec, computation_types.FederatedType): return computation_types.TensorType(tf.float32), True return type_spec, False
def _convert_sequence_to_tensor(type_spec): if isinstance(type_spec, computation_types.SequenceType): return computation_types.TensorType(tf.float32), True return type_spec, False
def _convert_abstract_type_to_tensor(type_spec): if isinstance(type_spec, computation_types.AbstractType): return computation_types.TensorType(tf.float32), True return type_spec, False
def _convert_placement_type_to_tensor(type_spec): if isinstance(type_spec, computation_types.PlacementType): return computation_types.TensorType(tf.float32), True return type_spec, False
def _convert_tensor_to_float(type_spec): if isinstance(type_spec, computation_types.TensorType): return computation_types.TensorType(tf.float32, shape=type_spec.shape), True return type_spec, False
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 = computation_types.to_type(type_spec) if isinstance(value, typed_object.TypedObject): if type_spec is not None: if not type_utils.are_equivalent_types(value.type_signature, type_spec): raise TypeError( 'Expected a value of type {}, found {}.'.format( str(type_spec), str(value.type_signature))) else: type_spec = value.type_signature if type_spec is None: raise ValueError( 'Cannot derive an eager representation for a value of an unknown type.' ) 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) else: raise TypeError('Unexpected type {}.'.format(str(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. """ 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( struct=pb.TensorFlow.StructBinding( 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, 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 = result._asdict().items() # pylint: enable=protected-access return _get_bindings_for_elements( name_value_pairs, graph, functools.partial( computation_types.StructWithPythonType, 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( name_value_pairs.items(), graph, functools.partial( computation_types.StructWithPythonType, container_type=type(result))) elif isinstance(result, structure.Struct): return _get_bindings_for_elements( structure.to_elements(result), graph, computation_types.StructType) elif isinstance(result, collections.abc.Mapping): if isinstance(result, collections.OrderedDict): name_value_pairs = result.items() else: name_value_pairs = sorted(result.items()) return _get_bindings_for_elements( name_value_pairs, graph, functools.partial( computation_types.StructWithPythonType, 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.StructWithPythonType( [e[0] for e in element_type_binding_pairs], type(result)), pb.TensorFlow.Binding( struct=pb.TensorFlow.StructBinding( element=[e[1] for e in element_type_binding_pairs]))) elif isinstance(result, type_conversions.TF_DATASET_REPRESENTATION_TYPES): variant_tensor = tf.data.experimental.to_variant(result) element_structure = result.element_spec try: element_type = computation_types.to_type(element_structure) except TypeError as e: raise TypeError( 'TFF does not support Datasets that yield elements of structure {!s}' .format(element_structure)) from e return (computation_types.SequenceType(element_type), pb.TensorFlow.Binding( sequence=pb.TensorFlow.SequenceBinding( variant_tensor_name=variant_tensor.name))) else: raise TypeError('Cannot capture a result of an unsupported type {}.'.format( py_typecheck.type_string(type(result))))
def test_tensor_type(self): s = computation_types.TensorType(tf.int32) t = computation_types.to_type(s) self.assertIsInstance(t, computation_types.TensorType) self.assertEqual(str(t), 'int32')
def _convert_tuple_to_tensor(type_spec): if isinstance(type_spec, computation_types.NamedTupleType): return computation_types.TensorType(tf.float32), True return type_spec, False
def test_dtype_and_shape(self): t = computation_types.TensorType(tf.int32, [10]) self.assertEqual(t.dtype, tf.int32) self.assertEqual(t.shape, tf.TensorShape([10]))
def foo(): return intrinsic_utils.zero_for( computation_types.TensorType(tf.float32, [2, 3]), context_stack_impl.context_stack)
def create_dummy_computation_reference(): value = pb.Computation(type=type_serialization.serialize_type(tf.int32), reference=pb.Reference(name='a')) type_signature = computation_types.TensorType(tf.int32) return value, type_signature
class TypeUtilsTest(test_case.TestCase, parameterized.TestCase): def test_to_canonical_value_with_none(self): self.assertIsNone(type_utils.to_canonical_value(None)) def test_to_canonical_value_with_int(self): self.assertEqual(type_utils.to_canonical_value(1), 1) def test_to_canonical_value_with_float(self): self.assertEqual(type_utils.to_canonical_value(1.0), 1.0) def test_to_canonical_value_with_bool(self): self.assertEqual(type_utils.to_canonical_value(True), True) self.assertEqual(type_utils.to_canonical_value(False), False) def test_to_canonical_value_with_string(self): self.assertEqual(type_utils.to_canonical_value('a'), 'a') def test_to_canonical_value_with_list_of_ints(self): self.assertEqual(type_utils.to_canonical_value([1, 2, 3]), [1, 2, 3]) def test_to_canonical_value_with_list_of_floats(self): self.assertEqual(type_utils.to_canonical_value([1.0, 2.0, 3.0]), [1.0, 2.0, 3.0]) def test_to_canonical_value_with_list_of_bools(self): self.assertEqual(type_utils.to_canonical_value([True, False]), [True, False]) def test_to_canonical_value_with_list_of_strings(self): self.assertEqual(type_utils.to_canonical_value(['a', 'b', 'c']), ['a', 'b', 'c']) def test_to_canonical_value_with_list_of_dict(self): self.assertEqual(type_utils.to_canonical_value([{ 'a': 1, 'b': 0.1, }]), [structure.Struct([ ('a', 1), ('b', 0.1), ])]) def test_to_canonical_value_with_list_of_ordered_dict(self): self.assertEqual( type_utils.to_canonical_value( [collections.OrderedDict([ ('a', 1), ('b', 0.1), ])]), [structure.Struct([ ('a', 1), ('b', 0.1), ])]) def test_to_canonical_value_with_dict(self): self.assertEqual(type_utils.to_canonical_value({ 'a': 1, 'b': 0.1, }), structure.Struct([ ('a', 1), ('b', 0.1), ])) self.assertEqual(type_utils.to_canonical_value({ 'b': 0.1, 'a': 1, }), structure.Struct([ ('a', 1), ('b', 0.1), ])) def test_to_canonical_value_with_ordered_dict(self): self.assertEqual( type_utils.to_canonical_value( collections.OrderedDict([ ('a', 1), ('b', 0.1), ])), structure.Struct([ ('a', 1), ('b', 0.1), ])) self.assertEqual( type_utils.to_canonical_value( collections.OrderedDict([ ('b', 0.1), ('a', 1), ])), structure.Struct([ ('b', 0.1), ('a', 1), ])) # pyformat: disable @parameterized.named_parameters([ ('buiding_block_and_type_spec', building_block_factory.create_compiled_identity( computation_types.TensorType(tf.int32)), computation_types.FunctionType(tf.int32, tf.int32), computation_types.FunctionType(tf.int32, tf.int32)), ('buiding_block_and_none', building_block_factory.create_compiled_identity( computation_types.TensorType(tf.int32)), None, computation_types.FunctionType(tf.int32, tf.int32)), ('int_and_type_spec', 10, computation_types.TensorType(tf.int32), computation_types.TensorType(tf.int32)), ]) # pyformat: enable def test_reconcile_value_with_type_spec_returns_type( self, value, type_spec, expected_type): actual_type = type_utils.reconcile_value_with_type_spec( value, type_spec) self.assertEqual(actual_type, expected_type) # pyformat: disable @parameterized.named_parameters([ ('building_block_and_bad_type_spec', building_block_factory.create_compiled_identity( computation_types.TensorType(tf.int32)), computation_types.TensorType(tf.int32)), ('int_and_none', 10, None), ]) # pyformat: enable def test_reconcile_value_with_type_spec_raises_type_error( self, value, type_spec): with self.assertRaises(TypeError): type_utils.reconcile_value_with_type_spec(value, type_spec) # pyformat: disable @parameterized.named_parameters([ ('value_type_and_type_spec', computation_types.TensorType(tf.int32), computation_types.TensorType(tf.int32), computation_types.TensorType(tf.int32)), ('value_type_and_none', computation_types.TensorType(tf.int32), None, computation_types.TensorType(tf.int32)), ]) # pyformat: enable def test_reconcile_value_type_with_type_spec_returns_type( self, value_type, type_spec, expected_type): actual_type = type_utils.reconcile_value_type_with_type_spec( value_type, type_spec) self.assertEqual(actual_type, expected_type) def test_reconcile_value_type_with_type_spec_raises_type_error_value_type_and_bad_type_spec( self): value_type = computation_types.TensorType(tf.int32) type_spec = computation_types.TensorType(tf.string) with self.assertRaises(TypeError): type_utils.reconcile_value_type_with_type_spec( value_type, type_spec)
def test_raises_type_error_with_value_and_bad_type(self, value, _): executor = create_test_executor(num_clients=3) bad_type_signature = computation_types.TensorType(tf.string) with self.assertRaises(TypeError): self.run_sync(executor.create_value(value, bad_type_signature))
def to_representation_for_type( value: Any, tf_function_cache: MutableMapping[str, Any], type_spec: Optional[computation_types.Type] = None, device: Optional[tf.config.LogicalDevice] = None) -> Any: """Verifies or converts the `value` to an eager object matching `type_spec`. WARNING: This function is only partially implemented. It does not support data sets at this point. The output of this function is always an eager tensor, eager dataset, a representation of a TensorFlow computation, or a nested structure of those that matches `type_spec`, and when `device` has been specified, everything is placed on that device on a best-effort basis. TensorFlow computations are represented here as zero- or one-argument Python callables that accept their entire argument bundle as a single Python object. Args: value: The raw representation of a value to compare against `type_spec` and potentially to be converted. tf_function_cache: A cache obeying `dict` semantics that can be used to look up previously embedded TensorFlow functions. type_spec: An instance of `tff.Type`, can be `None` for values that derive from `typed_object.TypedObject`. device: An optional `tf.config.LogicalDevice` to place the value on (for tensor-level values). Returns: Either `value` itself, or a modified version of it. Raises: TypeError: If the `value` is not compatible with `type_spec`. """ type_spec = 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.name if device else None) cached_fn = tf_function_cache.get(key) if cached_fn is not None: return cached_fn embedded_fn = embed_tensorflow_computation(value, type_spec, device) tf_function_cache[key] = embedded_fn return embedded_fn elif type_spec.is_struct(): type_elem = structure.to_elements(type_spec) value_elem = (structure.to_elements(structure.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 structure.Struct(result_elem) elif device is not None: py_typecheck.check_type(device, tf.config.LogicalDevice) with tf.device(device.name): return to_representation_for_type(value, tf_function_cache, type_spec=type_spec, device=None) elif isinstance(value, EagerValue): return value.internal_representation elif isinstance(value, executor_value_base.ExecutorValue): raise TypeError( 'Cannot accept a value embedded within a non-eager executor.') elif type_spec.is_tensor(): 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_spec.is_assignable_from(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 type_spec.is_sequence(): 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_spec.check_assignable_from(value_type) return value else: raise TypeError('Unexpected type {}.'.format(type_spec))
# limitations under the License. import collections import attr import tensorflow as tf from tensorflow_federated.python.common_libs import structure from tensorflow_federated.python.core.api import computation_types from tensorflow_federated.python.core.api import test_case from tensorflow_federated.python.core.impl.context_stack import context_base from tensorflow_federated.python.core.impl.context_stack import context_stack_impl from tensorflow_federated.python.core.impl.utils import function_utils from tensorflow_federated.python.core.impl.wrappers import computation_wrapper tffint32 = computation_types.TensorType(tf.int32) tffstring = computation_types.TensorType(tf.string) def build_zero_argument(parameter_type): if parameter_type is None: return None elif parameter_type.is_struct(): return structure.map_structure(build_zero_argument, parameter_type) elif parameter_type == tffint32: return 0 elif parameter_type == tffstring: return '' else: raise NotImplementedError(f'Unsupported type: {parameter_type}')
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 test_coerce_dataset_elements_noop(self): x = tf.data.Dataset.range(5) y = tensorflow_utils.coerce_dataset_elements_to_tff_type_spec( x, computation_types.TensorType(tf.int64)) self.assertEqual(x.element_spec, y.element_spec)
def test_make_dummy_element_TensorType(self): type_spec = computation_types.TensorType(tf.float32, [None, 10, None, 10, 10]) elem = graph_utils._make_dummy_element_for_type_spec(type_spec) correct_elem = np.zeros([0, 10, 0, 10, 10], np.float32) self.assertTrue(np.array_equal(elem, correct_elem))
def deserialize_type( type_proto: Optional[pb.Type]) -> Optional[computation_types.Type]: """Deserializes 'type_proto' as a computation_types.Type. Note: Currently only deserialization for tensor, named tuple, sequence, and function types is implemented. Args: type_proto: An instance of pb.Type or None. Returns: The corresponding instance of computation_types.Type (or None if the argument was None). Raises: TypeError: if the argument is of the wrong type. NotImplementedError: for type variants for which deserialization is not implemented. """ if type_proto is None: return None py_typecheck.check_type(type_proto, pb.Type) type_variant = type_proto.WhichOneof('type') if type_variant is None: return None elif type_variant == 'tensor': tensor_proto = type_proto.tensor return computation_types.TensorType( dtype=tf.dtypes.as_dtype(tensor_proto.dtype), shape=_to_tensor_shape(tensor_proto)) elif type_variant == 'sequence': return computation_types.SequenceType( deserialize_type(type_proto.sequence.element)) elif type_variant == 'struct': def empty_str_to_none(s): if s == '': # pylint: disable=g-explicit-bool-comparison return None return s return computation_types.StructType( [(empty_str_to_none(e.name), deserialize_type(e.value)) for e in type_proto.struct.element], convert=False) elif type_variant == 'function': return computation_types.FunctionType( parameter=deserialize_type(type_proto.function.parameter), result=deserialize_type(type_proto.function.result)) elif type_variant == 'placement': return computation_types.PlacementType() elif type_variant == 'federated': placement_oneof = type_proto.federated.placement.WhichOneof('placement') if placement_oneof == 'value': return computation_types.FederatedType( member=deserialize_type(type_proto.federated.member), placement=placements.uri_to_placement_literal( type_proto.federated.placement.value.uri), all_equal=type_proto.federated.all_equal) else: raise NotImplementedError( 'Deserialization of federated types with placement spec as {} ' 'is not currently implemented yet.'.format(placement_oneof)) else: raise NotImplementedError('Unknown type variant {}.'.format(type_variant))
class CheckWellFormedTest(parameterized.TestCase): # pyformat: disable @parameterized.named_parameters([ ('abstract_type', computation_types.AbstractType('T')), ('federated_type', computation_types.FederatedType(tf.int32, placement_literals.CLIENTS)), ('function_type', computation_types.FunctionType(tf.int32, tf.int32)), ('named_tuple_type', computation_types.NamedTupleType([tf.int32] * 3)), ('placement_type', computation_types.PlacementType()), ('sequence_type', computation_types.SequenceType(tf.int32)), ('tensor_type', computation_types.TensorType(tf.int32)), ]) # pyformat: enable def test_does_not_raise_type_error(self, type_signature): try: type_analysis.check_well_formed(type_signature) except TypeError: self.fail('Raised TypeError unexpectedly.') @parameterized.named_parameters([ ('federated_function_type', computation_types.FederatedType( computation_types.FunctionType(tf.int32, tf.int32), placement_literals.CLIENTS)), ('federated_federated_type', computation_types.FederatedType( computation_types.FederatedType(tf.int32, placement_literals.CLIENTS), placement_literals.CLIENTS)), ('sequence_sequence_type', computation_types.SequenceType( computation_types.SequenceType([tf.int32]))), ('sequence_federated_type', computation_types.SequenceType( computation_types.FederatedType(tf.int32, placement_literals.CLIENTS))), ('tuple_federated_function_type', computation_types.NamedTupleType([ computation_types.FederatedType( computation_types.FunctionType(tf.int32, tf.int32), placement_literals.CLIENTS) ])), ('tuple_federated_federated_type', computation_types.NamedTupleType([ computation_types.FederatedType( computation_types.FederatedType(tf.int32, placement_literals.CLIENTS), placement_literals.CLIENTS) ])), ('federated_tuple_function_type', computation_types.FederatedType( computation_types.NamedTupleType( [computation_types.FunctionType(tf.int32, tf.int32)]), placement_literals.CLIENTS)), ]) # pyformat: enable def test_raises_type_error(self, type_signature): with self.assertRaises(TypeError): type_analysis.check_well_formed(type_signature)
def test_with_int_vector(self): type_signature = computation_types.TensorType(tf.int32, [10]) dtypes, shapes = type_conversions.type_to_tf_dtypes_and_shapes( type_signature) self.assert_nested_struct_eq(dtypes, tf.int32) self.assert_nested_struct_eq(shapes, tf.TensorShape([10]))
def create_dummy_computation_tensorflow_identity(): """Returns a tensorflow computation and type `(float32 -> float32)`.""" tensor_type = computation_types.TensorType(tf.float32) value = tensorflow_computation_factory.create_identity(tensor_type) type_signature = computation_types.FunctionType(tensor_type, tensor_type) return value, type_signature