import os

os.environ["KERAS_BACKEND"] = "jax"

import keras
import numpy as np


keras.utils.set_random_seed(5)

model = keras.Sequential(
    [
        keras.Input(shape=(4,), name="features"),
        keras.layers.Dense(8, activation="relu", name="hidden"),
        keras.layers.Dense(3, activation="softmax", name="class_probs"),
    ],
    name="ticket_classifier",
)

sample = np.array([[0.1, 0.4, 0.2, 0.8], [0.7, 0.3, 0.6, 0.1]], dtype="float32")
prediction = model.predict(sample, verbose=0)

print(f"backend: {keras.backend.backend()}")
print(f"model name: {model.name}")
print(f"layer count: {len(model.layers)}")
print(f"parameter count: {model.count_params()}")
print(f"input shape: {model.input_shape}")
print(f"output shape: {model.output_shape}")
print(f"prediction shape: {prediction.shape}")
print(f"first row probability sum: {prediction[0].sum():.6f}")