def test_returns(self): t1 = computation_types.TensorType(tf.int32, [None]) t2 = computation_types.TensorType(tf.int32, [10]) t3 = computation_types.TensorType(tf.int32, [10]) self.assertTrue(type_analysis.are_equivalent_types(t1, t1)) self.assertTrue(type_analysis.are_equivalent_types(t2, t3)) self.assertTrue(type_analysis.are_equivalent_types(t3, t2)) self.assertFalse(type_analysis.are_equivalent_types(t1, t2)) self.assertFalse(type_analysis.are_equivalent_types(t2, t1))
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_analysis.are_equivalent_types( arg, computation_types.to_type(expected_args[idx]))) self.assertEqual(set(kwargs.keys()), set(expected_kwargs.keys())) for k, v in kwargs.items(): self.assertTrue( type_analysis.are_equivalent_types( computation_types.to_type(v), expected_kwargs[k]))
def reconcile_value_type_with_type_spec(value_type, type_spec): """Reconciles a pair of types. Args: value_type: An instance of `tff.Type` or something convertible to it. Must not be `None`. type_spec: An instance of `tff.Type`, something convertible to it, or `None`. Returns: Either `value_type` if `type_spec` is `None`, or `type_spec` if `type_spec` is not `None` and rquivalent with `value_type`. Raises: TypeError: If arguments are of incompatible types. """ value_type = computation_types.to_type(value_type) py_typecheck.check_type(value_type, computation_types.Type) if type_spec is None: return value_type else: type_spec = computation_types.to_type(type_spec) if type_analysis.are_equivalent_types(value_type, type_spec): return type_spec else: raise TypeError('Expected a value of type {}, found {}.'.format( type_spec, value_type))
async def create_value(self, value, type_spec=None): type_spec = computation_types.to_type(type_spec) if isinstance(value, computation_impl.ComputationImpl): return await self.create_value( computation_impl.ComputationImpl.get_proto(value), type_utils.reconcile_value_with_type_spec(value, type_spec)) py_typecheck.check_type(type_spec, computation_types.Type) hashable_key = _get_hashable_key(value, type_spec) try: identifier = self._cache.get(hashable_key) except TypeError as err: raise RuntimeError( 'Failed to perform a hash table lookup with a value of Python ' 'type {} and TFF type {}, and payload {}: {}'.format( py_typecheck.type_string(type(value)), type_spec, value, err)) if isinstance(identifier, CachedValueIdentifier): cached_value = self._cache.get(identifier) # If may be that the same payload appeared with a mismatching type spec, # which may be a legitimate use case if (as it happens) the payload alone # does not uniquely determine the type, so we simply opt not to reuse the # cache value and fallback on the regular behavior. if (cached_value is not None and type_spec is not None and not type_analysis.are_equivalent_types( cached_value.type_signature, type_spec)): identifier = None else: identifier = None if identifier is None: self._num_values_created = self._num_values_created + 1 identifier = CachedValueIdentifier(str(self._num_values_created)) self._cache[hashable_key] = identifier target_future = asyncio.ensure_future( self._target_executor.create_value(value, type_spec)) cached_value = None if cached_value is None: cached_value = CachedValue(identifier, hashable_key, type_spec, target_future) self._cache[identifier] = cached_value try: await cached_value.target_future except Exception as e: # Invalidate the entire cache in the inner executor had an exception. # TODO(b/145514490): This is a bit heavy handed, there maybe caches where # only the current cache item needs to be invalidated; however this # currently only occurs when an inner RemoteExecutor has the backend go # down. self._cache = {} raise e # No type check is necessary here; we have either checked # `type_analysis.are_equivalent_types` or just constructed `target_value` # explicitly with `type_spec`. return cached_value
def _serialize_deserialize_roundtrip_test(self, type_list): """Performs roundtrip serialization/deserialization of computation_types. Args: type_list: A list of instances of computation_types.Type or things convertible to it. """ for t in type_list: t1 = computation_types.to_type(t) p1 = type_serialization.serialize_type(t1) t2 = type_serialization.deserialize_type(p1) p2 = type_serialization.serialize_type(t2) self.assertEqual(repr(t1), repr(t2)) self.assertEqual(repr(p1), repr(p2)) self.assertTrue(type_analysis.are_equivalent_types(t1, t2))
def __add__(self, other): other = to_value(other, None, self._context_stack) if not type_analysis.are_equivalent_types(self.type_signature, other.type_signature): raise TypeError('Cannot add {} and {}.'.format( self.type_signature, other.type_signature)) return ValueImpl( building_blocks.Call( building_blocks.Intrinsic( intrinsic_defs.GENERIC_PLUS.uri, computation_types.FunctionType( [self.type_signature, self.type_signature], self.type_signature)), ValueImpl.get_comp( to_value([self, other], None, self._context_stack))), self._context_stack)
def __eq__(self, other): """Base class equality checks names and values equal.""" # TODO(b/130890785): Delegate value-checking to # `building_blocks.ComputationBuildingBlock`. if self is other: return True if not isinstance(other, BoundVariableTracker): return NotImplemented if self.name != other.name: return False if (isinstance(self.value, building_blocks.ComputationBuildingBlock) and isinstance(other.value, building_blocks.ComputationBuildingBlock)): return (self.value.compact_representation() == other.value.compact_representation() and type_analysis.are_equivalent_types(self.value.type_signature, other.value.type_signature)) return self.value is other.value
def embed_tensorflow_computation(comp, type_spec=None, device=None): """Embeds a TensorFlow computation for use in the eager context. Args: comp: An instance of `pb.Computation`. type_spec: An optional `tff.Type` instance or something convertible to it. device: An optional device name. Returns: Either a one-argument or a zero-argument callable that executes the computation in eager mode. Raises: TypeError: If arguments are of the wrong types, e.g., in `comp` is not a TensorFlow computation. """ # TODO(b/134543154): Decide whether this belongs in `tensorflow_utils.py` # since it deals exclusively with eager mode. Incubate here, and potentially # move there, once stable. py_typecheck.check_type(comp, pb.Computation) comp_type = type_serialization.deserialize_type(comp.type) type_spec = computation_types.to_type(type_spec) if type_spec is not None: if not type_analysis.are_equivalent_types(type_spec, comp_type): raise TypeError('Expected a computation of type {}, got {}.'.format( type_spec, comp_type)) else: type_spec = comp_type which_computation = comp.WhichOneof('computation') if which_computation != 'tensorflow': raise TypeError('Expected a TensorFlow computation, found {}.'.format( which_computation)) if isinstance(type_spec, computation_types.FunctionType): param_type = type_spec.parameter result_type = type_spec.result else: param_type = None result_type = type_spec if param_type is not None: input_tensor_names = tensorflow_utils.extract_tensor_names_from_binding( comp.tensorflow.parameter) else: input_tensor_names = [] output_tensor_names = tensorflow_utils.extract_tensor_names_from_binding( comp.tensorflow.result) def function_to_wrap(*args): # pylint: disable=missing-docstring if len(args) != len(input_tensor_names): raise RuntimeError('Expected {} arguments, found {}.'.format( len(input_tensor_names), len(args))) graph_def = serialization_utils.unpack_graph_def(comp.tensorflow.graph_def) init_op = comp.tensorflow.initialize_op if init_op: graph_def = tensorflow_utils.add_control_deps_for_init_op( graph_def, init_op) def _import_fn(): return tf.import_graph_def( graph_merge.uniquify_shared_names(graph_def), input_map=dict(list(zip(input_tensor_names, args))), return_elements=output_tensor_names) if device is not None: with tf.device(device): return _import_fn() else: return _import_fn() signature = [] param_fns = [] if param_type is not None: for spec in anonymous_tuple.flatten(type_spec.parameter): if isinstance(spec, computation_types.TensorType): signature.append(tf.TensorSpec(spec.shape, spec.dtype)) param_fns.append(lambda x: x) else: py_typecheck.check_type(spec, computation_types.SequenceType) signature.append(tf.TensorSpec([], tf.variant)) param_fns.append(tf.data.experimental.to_variant) wrapped_fn = tf.compat.v1.wrap_function(function_to_wrap, signature) result_fns = [] for spec in anonymous_tuple.flatten(result_type): if isinstance(spec, computation_types.TensorType): result_fns.append(lambda x: x) else: py_typecheck.check_type(spec, computation_types.SequenceType) structure = type_conversions.type_to_tf_structure(spec.element) def fn(x, structure=structure): return tf.data.experimental.from_variant(x, structure) result_fns.append(fn) def _fn_to_return(arg, param_fns, wrapped_fn): # pylint:disable=missing-docstring param_elements = [] if arg is not None: arg_parts = anonymous_tuple.flatten(arg) if len(arg_parts) != len(param_fns): raise RuntimeError('Expected {} arguments, found {}.'.format( len(param_fns), len(arg_parts))) for arg_part, param_fn in zip(arg_parts, param_fns): param_elements.append(param_fn(arg_part)) result_parts = wrapped_fn(*param_elements) # There is a tf.wrap_function(...) issue b/144127474 that variables created # from tf.import_graph_def(...) inside tf.wrap_function(...) is not # destroyed. So get all the variables from `wrapped_fn` and destroy # manually. # TODO(b/144127474): Remove this manual cleanup once tf.wrap_function(...) # is fixed. resources = [] for op in wrapped_fn.graph.get_operations(): if op.type == 'VarHandleOp': resources += op.outputs if resources: for resource in wrapped_fn.prune(feeds={}, fetches=resources)(): tf.raw_ops.DestroyResourceOp(resource=resource) result_elements = [] for result_part, result_fn in zip(result_parts, result_fns): result_elements.append(result_fn(result_part)) return anonymous_tuple.pack_sequence_as(result_type, result_elements) fn_to_return = lambda arg, p=param_fns, w=wrapped_fn: _fn_to_return(arg, p, w) if device is not None: old_fn_to_return = fn_to_return # pylint: disable=function-redefined def fn_to_return(x): with tf.device(device): return old_fn_to_return(x) # pylint: enable=function-redefined if param_type is not None: return lambda arg: fn_to_return(arg) # pylint: disable=unnecessary-lambda else: return lambda: fn_to_return(None)
def import_tensorflow_computation(comp, type_spec=None, name='fn'): """Converts a TF computation into an MLIR module that can be compiled by IREE. WARNING: This helper function is under construction, and most capabilities are not implemented at this stage: * The parameter and result of `comp` can only be a single tensor. Named tuples, sequences, or functional types are not currently supported. * Only tensorflow code can be imported. TODO(b/153499219): Add support for named tuples, sequences, and functions. Args: comp: An instance of a `pb.Computation` with TensorFlow code to import. type_spec: An optional `tff.Type` instance. name: An optional `str` name of the (single) function in the IREE module. Returns: An instance of IREE compiler's `CompilerModule` class with the imported function present. Raises: TypeError: If arguments are of the wrong types, e.g., in `comp` is not a TensorFlow computation. """ py_typecheck.check_type(comp, pb.Computation) comp_type = type_serialization.deserialize_type(comp.type) if type_spec is not None: py_typecheck.check_type(type_spec, computation_types.Type) if not type_analysis.are_equivalent_types(type_spec, comp_type): raise TypeError( 'Expected a computation of type {}, got {}.'.format( type_spec, comp_type)) else: type_spec = comp_type if isinstance(type_spec, computation_types.FunctionType): param_type = type_spec.parameter result_type = type_spec.result else: param_type = None result_type = type_spec # TODO(b/153499219): Replace this with a recursive check of the signature # after relaxing the type restrictions and introducing nested structures. py_typecheck.check_type(result_type, computation_types.TensorType) if param_type is not None: py_typecheck.check_type(param_type, computation_types.TensorType) which_computation = comp.WhichOneof('computation') if which_computation != 'tensorflow': raise TypeError('Expected a TensorFlow computation, found {}.'.format( which_computation)) output_tensor_names = tensorflow_utils.extract_tensor_names_from_binding( comp.tensorflow.result) if param_type is not None: input_tensor_names = tensorflow_utils.extract_tensor_names_from_binding( comp.tensorflow.parameter) else: input_tensor_names = [] graph_def = serialization_utils.unpack_graph_def(comp.tensorflow.graph_def) init_op = comp.tensorflow.initialize_op return_elements = input_tensor_names + output_tensor_names if init_op: graph_def = tensorflow_utils.add_control_deps_for_init_op( graph_def, init_op) return_elements.append(init_op) with tf.Graph().as_default() as graph: # TODO(b/153499219): See if we can reintroduce uniquify_shared_names(). # Right now, it causes loader breakage, and unclear if still necessary. import_results = tf.import_graph_def(graph_def, input_map={}, return_elements=return_elements, name='') if init_op: initializer = import_results[-1] import_results.pop() else: initializer = None inputs = import_results[0:len(input_tensor_names)] outputs = import_results[len(input_tensor_names):] with graph.as_default(): # TODO(b/153499219): Find a way to reflect the nested parameter and result # structure here after relaxing the restrictions. if inputs: assert len(inputs) < 2 input_dict = { 'parameter': tf.compat.v1.saved_model.utils.build_tensor_info(inputs[0]) } else: input_dict = {} assert len(outputs) == 1 output_dict = { 'result': tf.compat.v1.saved_model.utils.build_tensor_info(outputs[0]) } sig_def = tf.compat.v1.saved_model.signature_def_utils.build_signature_def( inputs=input_dict, outputs=output_dict, method_name=name) with tempfile.TemporaryDirectory() as model_dir: builder = tf.compat.v1.saved_model.Builder(model_dir) with tf.compat.v1.Session(graph=graph) as sess: builder.add_meta_graph_and_variables( sess, ['unused'], signature_def_map={name: sig_def}, legacy_init_op=initializer, strip_default_attrs=True) builder.save() return iree_compiler.tf_load_signature_def_saved_model( model_dir, tags=set(['unused']), exported_names=[name])
def _wrap(fn, parameter_type, wrapper_fn): """Wraps a possibly-polymorphic `fn` in `wrapper_fn`. If `parameter_type` is `None` and `fn` takes any arguments (even with default values), `fn` is inferred to be polymorphic and won't be passed to `wrapper_fn` until invocation time (when concrete parameter types are available). `wrapper_fn` must accept three positional arguments and one defaulted argument `name`: * `target_fn`, the Python function to be wrapped. * `parameter_type`, the optional type of the computation's parameter (an instance of `computation_types.Type`). * `unpack`, an argument which will be passed on to `function_utils.wrap_as_zero_or_one_arg_callable` when wrapping `target_fn`. See that function for details. * Optional `name`, the name of the function that is being wrapped (only for debugging purposes). Args: fn: The function or defun to wrap as a computation. parameter_type: Optional type of any arguments to `fn`. wrapper_fn: The Python callable that performs actual wrapping. The object to be returned by this function should be an instance of a `ConcreteFunction`. Returns: Either the result of wrapping (an object that represents the computation), or a polymorphic callable that performs wrapping upon invocation based on argument types. The returned function still may accept multiple arguments (it has not yet had `function_uils.wrap_as_zero_or_one_arg_callable` applied to it). Raises: TypeError: if the arguments are of the wrong types, or the `wrapper_fn` constructs something that isn't a ConcreteFunction. """ try: fn_name = fn.__name__ except AttributeError: fn_name = None signature = function_utils.get_signature(fn) parameter_type = computation_types.to_type(parameter_type) if parameter_type is None and signature.parameters: # There is no TFF type specification, and the function/defun declares # parameters. Create a polymorphic template. def _wrap_polymorphic(parameter_type: computation_types.Type, unpack: Optional[bool]): return wrapper_fn(fn, parameter_type, unpack=unpack, name=fn_name) polymorphic_fn = function_utils.PolymorphicFunction(_wrap_polymorphic) # When applying a decorator, the __doc__ attribute with the documentation # in triple-quotes is not automatically transferred from the function on # which it was applied to the wrapped object, so we must transfer it here # explicitly. polymorphic_fn.__doc__ = getattr(fn, '__doc__', None) return polymorphic_fn # Either we have a concrete parameter type, or this is no-arg function. concrete_fn = wrapper_fn(fn, parameter_type, unpack=None) py_typecheck.check_type(concrete_fn, function_utils.ConcreteFunction, 'value returned by the wrapper') if not type_analysis.are_equivalent_types( concrete_fn.type_signature.parameter, parameter_type): raise TypeError( 'Expected a concrete function that takes parameter {}, got one ' 'that takes {}.'.format(str(parameter_type), str(concrete_fn.type_signature.parameter))) # When applying a decorator, the __doc__ attribute with the documentation # in triple-quotes is not automatically transferred from the function on concrete_fn.__doc__ = getattr(fn, '__doc__', None) return concrete_fn