def test_unpack_args_from_tuple_type(self, tuple_with_args, expected_args, expected_kwargs): args, kwargs = function_utils.unpack_args_from_tuple(tuple_with_args) self.assertEqual(len(args), len(expected_args)) for idx, arg in enumerate(args): self.assertTrue( type_utils.are_equivalent_types( arg, computation_types.to_type(expected_args[idx]))) self.assertEqual(set(kwargs.keys()), set(expected_kwargs.keys())) for k, v in six.iteritems(kwargs): self.assertTrue( type_utils.are_equivalent_types(computation_types.to_type(v), expected_kwargs[k]))
def test_unpack_args_from_anonymous_tuple(self, tuple_with_args, expected_args, expected_kwargs): self.assertEqual( function_utils.unpack_args_from_tuple(tuple_with_args), (expected_args, expected_kwargs))
def get_tf_typespec_and_binding(parameter_type, arg_names, unpack=None): """Computes a `TensorSpec` input_signature and bindings for parameter_type. This is the TF2 analog to `stamp_parameter_in_graph`. Args: parameter_type: The TFF type of the input to a tensorflow function. Must be either an instance of computation_types.Type (or convertible to it), or None in the case of a no-arg function. arg_names: String names for any positional arguments to the tensorflow function. unpack: Whether or not to unpack parameter_type into args and kwargs. See e.g. `function_utils.pack_args_into_anonymous_tuple`. Returns: A tuple (args_typespec, kwargs_typespec, binding), where args_typespec is a list and kwargs_typespec is a dict, both containing `tf.TensorSpec` objects. These structures are intended to be passed to the `get_concrete_function` method of a `tf.function`. Note the "binding" is "preliminary" in that it includes the names embedded in the TensorSpecs produced; these must be converted to the names of actual tensors based on the SignatureDef of the SavedModel before the binding is finalized. """ if parameter_type is None: return ([], {}, None) if unpack: arg_types, kwarg_types = function_utils.unpack_args_from_tuple( parameter_type) pack_in_tuple = True else: pack_in_tuple = False arg_types, kwarg_types = [parameter_type], {} py_typecheck.check_type(arg_names, collections.Iterable) if len(arg_names) < len(arg_types): raise ValueError( 'If provided, arg_names must be a list of at least {} strings to ' 'match the number of positional arguments. Found: {}'.format( len(arg_types), arg_names)) get_unique_name = UniqueNameFn() def _get_one_typespec_and_binding(parameter_name, parameter_type): """Returns a (tf.TensorSpec, binding) pair.""" parameter_type = computation_types.to_type(parameter_type) if isinstance(parameter_type, computation_types.TensorType): name = get_unique_name(parameter_name) tf_spec = tf.TensorSpec( shape=parameter_type.shape, dtype=parameter_type.dtype, name=name) binding = pb.TensorFlow.Binding( tensor=pb.TensorFlow.TensorBinding(tensor_name=name)) return (tf_spec, binding) elif isinstance(parameter_type, computation_types.NamedTupleType): element_typespec_pairs = [] element_bindings = [] have_names = False have_nones = False for e_name, e_type in anonymous_tuple.to_elements(parameter_type): if e_name is None: have_nones = True else: have_names = True name = '_'.join([n for n in [parameter_name, e_name] if n]) e_typespec, e_binding = _get_one_typespec_and_binding( name if name else None, e_type) element_typespec_pairs.append((e_name, e_typespec)) element_bindings.append(e_binding) # For a given argument or kwarg, we shouldn't have both: if (have_names and have_nones): raise ValueError( 'A mix of named and unnamed entries are not supported inside ' 'a nested structure representing a single argument in a call ' 'to a TensorFlow or Python function.\n' + str(parameter_type)) tf_typespec = anonymous_tuple.AnonymousTuple(element_typespec_pairs) return (tf_typespec, pb.TensorFlow.Binding( tuple=pb.TensorFlow.NamedTupleBinding( element=element_bindings))) elif isinstance(parameter_type, computation_types.SequenceType): raise NotImplementedError('Sequence iputs not yet supported for TF 2.0.') else: raise ValueError('Parameter type component {} cannot be converted ' 'to a TensorSpec'.format(repr(parameter_type))) def get_arg_name(i): name = arg_names[i] if not isinstance(name, six.string_types): raise ValueError('arg_names must be strings, but got' + str(name)) return name # Main logic --- process arg_types and kwarg_types: arg_typespecs = [] kwarg_typespecs = {} bindings = [] for i, arg_type in enumerate(arg_types): name = get_arg_name(i) typespec, binding = _get_one_typespec_and_binding(name, arg_type) typespec = type_utils.convert_to_py_container(typespec, arg_type) arg_typespecs.append(typespec) bindings.append(binding) for name, kwarg_type in six.iteritems(kwarg_types): typespec, binding = _get_one_typespec_and_binding(name, kwarg_type) typespec = type_utils.convert_to_py_container(typespec, kwarg_type) kwarg_typespecs[name] = typespec bindings.append(binding) assert bindings, 'Given parameter_type {}, but produced no bindings.'.format( parameter_type) if pack_in_tuple: final_binding = pb.TensorFlow.Binding( tuple=pb.TensorFlow.NamedTupleBinding(element=bindings)) else: final_binding = bindings[0] return (arg_typespecs, kwarg_typespecs, final_binding)