Understanding the tf.function in TensorFlow

Hello coders,

In this tutorial, we will look at the tf.function in the TensorFlow library in Python. We use the tf.function to make graphs out of our programs. In Python, it is used as a transformation tool that creates Python-independent dataflow graphs out of Python code. It helps in creating portable and performant models.

You can install the TensorFlow library using the following command:

pip install tensorflow

And we will use the TensorFlow library as follows:

import tensorflow as tf

tf.function complies a function into callable graphs.

Syntax:

tf.function(
func=None,
input_signature=None,
autograph=True,
jit_compile=None,
reduce_retracing=False,
experimental_implements=None,
experimental_autograph_options=None,
experimental_relax_shapes=None,
experimental_compile=None,
experimental_follow_type_hints=None
) -> tf.types.experimental.GenericFunction

Usage of tf.function in Examples

Now let’s just see a basic example of this function:

In this example we can see with trace-compilation we can execute non-TensorFlow operations, but only under special conditions.

@tf.function
def f(x, y):
    return x ** 4 + y
x = tf.constant([9, 6])
y = tf.constant([5, -7])
f(x, y)

Output:

<tf.Tensor: shape=(2,), dtype=int32, numpy=array([6566, 1289])>

Features

In this function we can use control flow statements like while, break, continue, if-else etc.

@tf.function
def f(x):
    if tf.reduce_sum(x) > 0:
        return x * 2
    else:
        return -x // 4
f(tf.constant(9))

Output:

<tf.Tensor: shape=(), dtype=int32, numpy=18>

We can also include tf.Variable and tf.Tensor in this function:

But we can only create tf.Variable object only time when it is called for the first time. That is why it is recommended to crate it outside the tf.function.

@tf.function
def f(x, y):
    return x ** 4 + y
x = tf.constant([9, 6])
y = tf.Variable([5, -7])
f(x, y)

Output:

<tf.Tensor: shape=(), dtype=int32, numpy=18>