import os

os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"

import tensorflow as tf

tf.get_logger().setLevel("WARNING")


def sample_batches():
    return [tf.ones((rows, 3), dtype=tf.float32) for rows in range(1, 8)]


@tf.function
def score_unfixed(features):
    print(f"Tracing unfixed graph for shape {features.shape}")
    return tf.reduce_sum(features, axis=1)


@tf.function(
    input_signature=[
        tf.TensorSpec(shape=(None, 3), dtype=tf.float32, name="features")
    ]
)
def score_fixed(features):
    print(f"Tracing fixed graph for shape {features.shape}")
    return tf.reduce_sum(features, axis=1)


print(f"TensorFlow {tf.__version__}")
print("Unfixed calls:")
for batch in sample_batches():
    score_unfixed(batch)
print(f"Unfixed tracing count: {score_unfixed.experimental_get_tracing_count()}")

print("")
print("Fixed calls:")
last_output = None
for batch in sample_batches():
    last_output = score_fixed(batch)
print(f"Fixed tracing count: {score_fixed.experimental_get_tracing_count()}")
print(f"Fixed last output shape: {tuple(last_output.shape.as_list())}")