TF Better performance with tf.

function
In TensorFlow 2, eager execution is turned on by default. The user interface is intuitive and
flexible (running one-off operations is much easier and faster), but this can come at the
expense of performance and deployability.

You can use tf.function to make graphs out of your programs. It is a transformation tool that
creates Python-independent dataflow graphs out of your Python code. This will help you create
performant and portable models, and it is required to use SavedModel .

This guide will help you conceptualize how tf.function works under the hood, so you can
use it effectively.

The main takeaways and recommendations are:

Debug in eager mode, then decorate with @tf.function .

Don't rely on Python side effects like object mutation or list appends.
tf.function works best with TensorFlow ops; NumPy and Python calls are converted to
constants.

Basics
Usage
A tf.function that you define (for example by applying the @tf.function decorator) is just
like a core TensorFlow operation: You can execute it eagerly; you can compute gradients; and
so on.

tf.function s can be faster than eager code, especially for graphs with many small ops. But
for graphs with a few expensive ops (like convolutions), you may not see much speedup.

Tracing
This section exposes how tf.function works under the hood, including implementation
details which may change in the future. However, once you understand why and when tracing
happens, it's much easier to use tf.function effectively!

What is "tracing"?
A tf.function runs your program in a TensorFlow Graph. However, a tf.Graph cannot
represent all the things that you'd write in an eager TensorFlow program. For instance, Python
supports polymorphism, but tf.Graph requires its inputs to have a specified data type and
dimension. Or you may perform side tasks like reading command-line arguments, raising an
error, or working with a more complex Python object; none of these things can run in
a tf.Graph .

tf.function bridges this gap by separating your code in two stages:

1. In the first stage, referred to as "tracing", tf.function creates a new tf.Graph . Python
code runs normally, but all TensorFlow operations (like adding two Tensors) are deferred:
they are captured by the tf.Graph and not run.
2. In the second stage, a tf.Graph which contains everything that was deferred in the first
stage is run. This stage is much faster than the tracing stage.

Depending on its inputs, tf.function will not always run the first stage when it is called.
See "Rules of tracing" below to get a better sense of how it makes that determination. Skipping
the first stage and only executing the second stage is what gives you TensorFlow's high
performance.

When tf.function does decide to trace, the tracing stage is immediately followed by the
second stage, so calling the tf.function both creates and runs the tf.Graph . Later you will
see how you can run only the tracing stage with get_concrete_function .

When you pass arguments of different types into a tf.function , both stages are run:

@tf.functiondef
double(a):
print("Tracing with", a)
return a + a
print(double(tf.constant(1)))
print()
print(double(tf.constant(1.1)))
print()
print(double(tf.constant("a")))
print()

Note that if you repeatedly call a tf.function with the same argument type, TensorFlow will
skip the tracing stage and reuse a previously traced graph, as the generated graph would be
identical. You can use pretty_printed_concrete_signatures() to see all of the available
traces

So far, you've seen that tf.function creates a cached, dynamic dispatch layer over
TensorFlow's graph tracing logic. To be more specific about the terminology:
 A tf.Graph is the raw, language-agnostic, portable representation of a TensorFlow
computation.
Tracing is the process through which new tf.Graph s are generated from Python code.
An instance of tf.Graph is specialized to the specific input types it was traced with.
Differing types require retracing.
Each traced tf.Graph has a corresponding ConcreteFunction .
A tf.function manages a cache of ConcreteFunction s and picks the right one for your
inputs.
tf.function wraps the Python function that will be traced, returning
a tf.types.experimental.PolymorphicFunction object.

Rules of tracing

When called, a tf.function first evaluates the type of each input argument using
the tf.types.experimental.TraceType of each argument. This is used to construct
a tf.types.experimental.FunctionType describing the signature of the
desired ConcreteFunction . We compare this FunctionType to the FunctionType s of
existing ConcreteFunction s. If a matching ConcreteFunction is found, the call is dispatched
to it. If no match is found, a new ConcreteFunction is traced for the desired FunctionType .

If multiple matches are found, the most specific signature is chosen. Matching is done
by subtyping, much like normal function calls in C++ or Java, for instance. For
example, TensorShape([1, 2]) is a subtype of TensorShape([None, None]) and so a call to
the tf.function with TensorShape([1, 2]) can be dispatched to
the ConcreteFunction produced with TensorShape([None, None]) but if
a ConcreteFunction with TensorShape([1, None]) also exists then it will be prioritized since
it is more specific.

The TraceType is determined from input arguments as follows:

For Tensor , the type is parameterized by the Tensor 's dtype and shape ; ranked
shapes are a subtype of unranked shapes; fixed dimensions are a subtype of unknown
dimensions
For Variable , the type is similar to Tensor , but also includes a unique resource ID of the
variable, necessary to correctly wire control dependencies
For Python primitive values, the type corresponds to the value itself. For example,
the TraceType of the value 3 is LiteralTraceType<3> , not int .
For Python ordered containers such as list and tuple , etc., the type is parameterized
by the types of their elements; for example, the type of [1,
2] is ListTraceType<LiteralTraceType<1>, LiteralTraceType<2>> and the type
 for [2, 1] is ListTraceType<LiteralTraceType<2>, LiteralTraceType<1>> which is
different.
For Python mappings such as dict , the type is also a mapping from the same keys but to
the types of values instead of the actual values. For example, the type of {1: 2, 3: 4} ,
is MappingTraceType<<KeyValue<1, LiteralTraceType<2>>>, <KeyValue<3,
LiteralTraceType<4>>>> . However, unlike ordered containers, {1: 2, 3: 4} and {3:
4, 1: 2} have equivalent types.
For Python objects which implement the __tf_tracing_type__ method, the type is
whatever that method returns.
For any other Python objects, the type is a generic TraceType , and the matching
precedure is:
First it checks if the object is the same object used in the previous trace (using
Python id() or is ). Note that this will still match if the object has changed, so if you
use Python objects as tf.function arguments it's best to use immutable ones.
Next it checks if the object is equal to the object used in the previous trace (using
Python == ).
Note that this procedure only keeps a weakref to the object and hence only works as long
as the object is in scope/not deleted.

Controlling retracing
Retracing, which is when your tf.function creates more than one trace, helps ensure that
TensorFlow generates correct graphs for each set of inputs. However, tracing is an expensive
operation! If your tf.function retraces a new graph for every call, you'll find that your code
executes more slowly than if you didn't use tf.function .

To control the tracing behavior, you can use the following techniques:

Pass a fixed input_signature to tf.function

This forces tf.function to constrain itself to only
one tf.types.experimental.FunctionType composed of the types enumerated by
the input_signature . Calls that cannot be dispatched to this FunctionType will throw an
error.
Use unknown dimensions for flexibility

Since TensorFlow matches tensors based on their shape, using a None dimension as a
wildcard will allow tf.function s to reuse traces for variably-sized input. Variably-sized input
can occur if you have sequences of different length, or images of different sizes for each batch.
You can check out the Transformer and Deep Dream tutorials for examples.

@tf.function(input_signature=(tf.TensorSpec(shape=[None], dtype=tf.int32),))
def g(x):
print('Tracing with', x)
return x

Use reduce_retracing for automatic flexibility

When reduce_retracing is enabled, tf.function automatically identifies supertypes of the
input types it is observing and chooses to trace more generalized graphs automatically. It is less
efficient than setting the input_signature directly but useful when many types need to be
supported.

@tf.function(reduce_retracing=True)
def g(x):
print('Tracing with', x)
return x