import os

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

import tensorflow as tf

tf.get_logger().setLevel("ERROR")
tf.keras.utils.set_random_seed(42)

mixed_precision = tf.keras.mixed_precision
mixed_precision.set_global_policy("mixed_float16")

inputs = tf.keras.Input(shape=(4,), name="features")
x = tf.keras.layers.Dense(16, activation="relu", name="hidden")(inputs)
outputs = tf.keras.layers.Dense(1, activation="sigmoid", dtype="float32", name="score")(x)
model = tf.keras.Model(inputs, outputs)

model.compile(
    optimizer="adam",
    loss="binary_crossentropy",
    metrics=["accuracy"],
)

features = tf.random.normal((32, 4))
labels = tf.cast(tf.reduce_sum(features, axis=1, keepdims=True) > 0, tf.float32)
history = model.fit(features, labels, epochs=1, batch_size=8, verbose=0)

print(f"TensorFlow {tf.__version__}")
print(f"global_policy={mixed_precision.global_policy().name}")
print(f"hidden_compute_dtype={model.get_layer('hidden').compute_dtype}")
print(f"hidden_variable_dtype={model.get_layer('hidden').variable_dtype}")
print(f"output_compute_dtype={model.get_layer('score').compute_dtype}")
print(f"prediction_dtype={model(features[:2]).dtype.name}")
print(f"loss={history.history['loss'][-1]:.4f}")
print(f"accuracy={history.history['accuracy'][-1]:.4f}")