Why tf.function traces layers twice?

Question

The example code is as follows:

import tensorflow as tf     # tf-2.4 or tf-2.x
from datetime import datetime

# Define a layer with an eager side effect
class EagerLayer(tf.keras.layers.Layer):
  def __init__(self, **kwargs):
    super(EagerLayer, self).__init__(**kwargs)
    # Do some kind of initialization here
    self.dense = tf.keras.layers.Dense(32)

  def call(self, inputs):
    print("\nCurrently running eagerly", str(datetime.now()))
    return self.dense(inputs)

input_data = tf.random.uniform([60, 28, 28])
layer = EagerLayer()
tf_func_layer = tf.function(layer)
print("=============")
_ = tf_func_layer(input_data)

The output is

=============

Currently running eagerly 2020-12-16 20:46:38.482592

Currently running eagerly 2020-12-16 20:46:38.503605

The printing side effect shows up twice, which means the function is traced twice.

I am just wondering why the layer is traced twice.

A colab notebook can be found here.

Answer 1

Switching on verbose gives the below output ~

=============
INFO:tensorflow:Converted call: <__main__.EagerLayer object at 0x7fbf6e34b860>
    args: (<tf.Tensor 'self:0' shape=(60, 28, 28) dtype=float32>,)
    kwargs: {}

INFO:tensorflow:Converted call: <built-in function hasattr>
    args: (<__main__.EagerLayer object at 0x7fbf6e34b860>, '_thread_local')
    kwargs: None

INFO:tensorflow:Converted call: <bound method Layer._split_out_first_arg of <__main__.EagerLayer object at 0x7fbf6e34b860>>
    args: ((<tf.Tensor 'self:0' shape=(60, 28, 28) dtype=float32>,), {})
    kwargs: None

INFO:tensorflow:Converted call: <function flatten at 0x7fbf86ceb598>
    args: (<tf.Tensor 'self:0' shape=(60, 28, 28) dtype=float32>,)
    kwargs: None

INFO:tensorflow:Converted call: <function _in_functional_construction_mode at 0x7fbf82850840>
    args: (<__main__.EagerLayer object at 0x7fbf6e34b860>, <tf.Tensor 'self:0' shape=(60, 28, 28) dtype=float32>, (), {}, [<tf.Tensor 'self:0' shape=(60, 28, 28) dtype=float32>])
    kwargs: None

INFO:tensorflow:Converted call: <function call_context at 0x7fbf8276bf28>
    args: ()
    kwargs: None

INFO:tensorflow:Converted call: <built-in function any>
    args: (<generator object outer_factory.<locals>.inner_factory.<locals>.tf____call__.<locals>.else_body_10.<locals>.<genexpr> at 0x7fbf6e07d4c0>,)
    kwargs: None

INFO:tensorflow:Converted call: <built-in function isinstance>
    args: (<tf.Tensor 'self:0' shape=(60, 28, 28) dtype=float32>, (<class 'tensorflow.python.ops.numpy_ops.np_arrays.ndarray'>, <class 'numpy.ndarray'>, <class 'float'>, <class 'int'>))
    kwargs: None

INFO:tensorflow:Converted call: <bound method Layer._get_input_masks of <__main__.EagerLayer object at 0x7fbf6e34b860>>
    args: (<tf.Tensor 'self:0' shape=(60, 28, 28) dtype=float32>, [<tf.Tensor 'self:0' shape=(60, 28, 28) dtype=float32>], (), {})
    kwargs: None

INFO:tensorflow:Converted call: <bound method Layer._set_training_mode of <__main__.EagerLayer object at 0x7fbf6e34b860>>
    args: ((), {}, <tensorflow.python.keras.engine.base_layer_utils.CallContext object at 0x7fbf79a7c860>)
    kwargs: None

INFO:tensorflow:Converted call: <bound method Layer._clear_losses of <__main__.EagerLayer object at 0x7fbf6e34b860>>
    args: ()
    kwargs: None

INFO:tensorflow:Converted call: <function executing_eagerly at 0x7fbf86cb58c8>
    args: ()
    kwargs: None

INFO:tensorflow:Converted call: <bound method Layer._maybe_cast_inputs of <__main__.EagerLayer object at 0x7fbf6e34b860>>
    args: (<tf.Tensor 'self:0' shape=(60, 28, 28) dtype=float32>, [<tf.Tensor 'self:0' shape=(60, 28, 28) dtype=float32>])
    kwargs: None

INFO:tensorflow:Converted call: <function assert_input_compatibility at 0x7fbf8276f2f0>
    args: (None, <tf.Tensor 'self:0' shape=(60, 28, 28) dtype=float32>, 'eager_layer_7')
    kwargs: None

INFO:tensorflow:Converted call: <bound method Layer._name_scope of <__main__.EagerLayer object at 0x7fbf6e34b860>>
    args: ()
    kwargs: None

INFO:tensorflow:Converted call: <bound method Layer._autographed_call of <__main__.EagerLayer object at 0x7fbf6e34b860>>
    args: ()
    kwargs: None

INFO:tensorflow:Converted call: <bound method Layer._maybe_build of <__main__.EagerLayer object at 0x7fbf6e34b860>>
    args: (<tf.Tensor 'self:0' shape=(60, 28, 28) dtype=float32>,)
    kwargs: None

INFO:tensorflow:Converted call: <function EagerLayer.call at 0x7fbf6ccc1bf8>
    args: (<tf.Tensor 'self:0' shape=(60, 28, 28) dtype=float32>,)
    kwargs: {}

INFO:tensorflow:Converted call: <bound method EagerLayer.call of <__main__.EagerLayer object at 0x7fbf6e34b860>>
    args: (<tf.Tensor 'self:0' shape=(60, 28, 28) dtype=float32>,)
    kwargs: {}

INFO:tensorflow:Converted call: <built-in method now of type object at 0xa33e60>
    args: ()
    kwargs: None

INFO:tensorflow:Converted call: <class 'str'>
    args: (datetime.datetime(2020, 12, 31, 15, 15, 24, 101021),)
    kwargs: None


Currently running eagerly 2020-12-31 15:15:24.101021
INFO:tensorflow:Converted call: <tensorflow.python.keras.layers.core.Dense object at 0x7fbf6cccb198>
    args: (<tf.Tensor 'self:0' shape=(60, 28, 28) dtype=float32>,)
    kwargs: None

INFO:tensorflow:Converted call: <bound method Layer._set_save_spec of <__main__.EagerLayer object at 0x7fbf6e34b860>>
    args: (<tf.Tensor 'self:0' shape=(60, 28, 28) dtype=float32>,)
    kwargs: None

INFO:tensorflow:Converted call: <__main__.EagerLayer object at 0x7fbf6e34b860>
    args: (<tf.Tensor 'self:0' shape=(60, 28, 28) dtype=float32>,)
    kwargs: {}

INFO:tensorflow:Converted call: <built-in function hasattr>
    args: (<__main__.EagerLayer object at 0x7fbf6e34b860>, '_thread_local')
    kwargs: None

INFO:tensorflow:Converted call: <bound method Layer._split_out_first_arg of <__main__.EagerLayer object at 0x7fbf6e34b860>>
    args: ((<tf.Tensor 'self:0' shape=(60, 28, 28) dtype=float32>,), {})
    kwargs: None

INFO:tensorflow:Converted call: <function flatten at 0x7fbf86ceb598>
    args: (<tf.Tensor 'self:0' shape=(60, 28, 28) dtype=float32>,)
    kwargs: None

INFO:tensorflow:Converted call: <function _in_functional_construction_mode at 0x7fbf82850840>
    args: (<__main__.EagerLayer object at 0x7fbf6e34b860>, <tf.Tensor 'self:0' shape=(60, 28, 28) dtype=float32>, (), {}, [<tf.Tensor 'self:0' shape=(60, 28, 28) dtype=float32>])
    kwargs: None

INFO:tensorflow:Converted call: <function call_context at 0x7fbf8276bf28>
    args: ()
    kwargs: None

INFO:tensorflow:Converted call: <built-in function any>
    args: (<generator object outer_factory.<locals>.inner_factory.<locals>.tf____call__.<locals>.else_body_10.<locals>.<genexpr> at 0x7fbf6e3615c8>,)
    kwargs: None

INFO:tensorflow:Converted call: <built-in function isinstance>
    args: (<tf.Tensor 'self:0' shape=(60, 28, 28) dtype=float32>, (<class 'tensorflow.python.ops.numpy_ops.np_arrays.ndarray'>, <class 'numpy.ndarray'>, <class 'float'>, <class 'int'>))
    kwargs: None

INFO:tensorflow:Converted call: <bound method Layer._get_input_masks of <__main__.EagerLayer object at 0x7fbf6e34b860>>
    args: (<tf.Tensor 'self:0' shape=(60, 28, 28) dtype=float32>, [<tf.Tensor 'self:0' shape=(60, 28, 28) dtype=float32>], (), {})
    kwargs: None

INFO:tensorflow:Converted call: <bound method Layer._set_training_mode of <__main__.EagerLayer object at 0x7fbf6e34b860>>
    args: ((), {}, <tensorflow.python.keras.engine.base_layer_utils.CallContext object at 0x7fbf79a7c860>)
    kwargs: None

INFO:tensorflow:Converted call: <bound method Layer._clear_losses of <__main__.EagerLayer object at 0x7fbf6e34b860>>
    args: ()
    kwargs: None

INFO:tensorflow:Converted call: <function executing_eagerly at 0x7fbf86cb58c8>
    args: ()
    kwargs: None

INFO:tensorflow:Converted call: <bound method Layer._maybe_cast_inputs of <__main__.EagerLayer object at 0x7fbf6e34b860>>
    args: (<tf.Tensor 'self:0' shape=(60, 28, 28) dtype=float32>, [<tf.Tensor 'self:0' shape=(60, 28, 28) dtype=float32>])
    kwargs: None

INFO:tensorflow:Converted call: <function assert_input_compatibility at 0x7fbf8276f2f0>
    args: (None, <tf.Tensor 'self:0' shape=(60, 28, 28) dtype=float32>, 'eager_layer_7')
    kwargs: None

INFO:tensorflow:Converted call: <bound method Layer._name_scope of <__main__.EagerLayer object at 0x7fbf6e34b860>>
    args: ()
    kwargs: None

INFO:tensorflow:Converted call: <bound method Layer._autographed_call of <__main__.EagerLayer object at 0x7fbf6e34b860>>
    args: ()
    kwargs: None

INFO:tensorflow:Converted call: <function EagerLayer.call at 0x7fbf6ccfcd08>
    args: (<tf.Tensor 'self:0' shape=(60, 28, 28) dtype=float32>,)
    kwargs: {}

INFO:tensorflow:Converted call: <bound method EagerLayer.call of <__main__.EagerLayer object at 0x7fbf6e34b860>>
    args: (<tf.Tensor 'self:0' shape=(60, 28, 28) dtype=float32>,)
    kwargs: {}

INFO:tensorflow:Converted call: <built-in method now of type object at 0xa33e60>
    args: ()
    kwargs: None

INFO:tensorflow:Converted call: <class 'str'>
    args: (datetime.datetime(2020, 12, 31, 15, 15, 24, 153569),)
    kwargs: None


Currently running eagerly 2020-12-31 15:15:24.153569
INFO:tensorflow:Converted call: <tensorflow.python.keras.layers.core.Dense object at 0x7fbf6cccb198>
    args: (<tf.Tensor 'self:0' shape=(60, 28, 28) dtype=float32>,)
    kwargs: None

I am assuming that the reason for it calling twice is for running set_save_spec, but still not sure...