import os

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

import tensorflow as tf

tf.get_logger().setLevel("ERROR")

weights = tf.constant([0.7, -0.2, 0.5], dtype=tf.float32)
bias = tf.constant(0.1, dtype=tf.float32)


@tf.function(
    input_signature=[
        tf.TensorSpec(shape=(None, 3), dtype=tf.float32, name="features")
    ]
)
def score_batch(features):
    return tf.reduce_sum(features * weights, axis=1) + bias


compiled_score = score_batch.get_concrete_function()

first_batch = tf.constant(
    [
        [1.0, 0.2, 0.5],
        [0.3, 0.9, 1.2],
    ],
    dtype=tf.float32,
)
second_batch = tf.constant([[0.8, 0.1, 0.4]], dtype=tf.float32)

first_scores = score_batch(first_batch)
second_scores = score_batch(second_batch)
same_graph = score_batch.get_concrete_function() is compiled_score

try:
    score_batch(tf.constant([[1.0, 2.0]], dtype=tf.float32))
except TypeError:
    bad_shape_error = "TypeError"
else:
    bad_shape_error = "none"

graph_ops = ", ".join(op.name for op in compiled_score.graph.get_operations())

print(f"TensorFlow {tf.__version__}")
print(f"Input signature: {compiled_score.structured_input_signature[0][0]}")
print(f"Output shape: {tuple(compiled_score.output_shapes.as_list())}")
print(f"Graph operations: {graph_ops}")
print(
    "First batch scores: "
    f"{[round(float(value), 2) for value in first_scores.numpy()]}"
)
print(
    "Second batch scores: "
    f"{[round(float(value), 2) for value in second_scores.numpy()]}"
)
print(f"Reused compiled graph: {same_graph}")
print(f"Bad shape error: {bad_shape_error}")