import os

os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "2")
os.environ.setdefault("TF_USE_LEGACY_KERAS", "1")

import numpy as np
import tensorflow as tf

model = tf.keras.Sequential(
    [
        tf.keras.Input(shape=(3,), name="features"),
        tf.keras.layers.Dense(2, activation="softmax", name="score"),
    ]
)

prediction = model(np.ones((1, 3), dtype="float32"), training=False)

print(f"keras_module={tf.keras.layers.Dense.__module__.split('.')[0]}")
print(f"prediction_shape={tuple(prediction.shape)}")