import os

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

import keras
import numpy as np


keras.utils.set_random_seed(42)
keras.config.set_dtype_policy("mixed_float16")

x_train = np.linspace(0.0, 1.0, 64, dtype="float32").reshape(16, 4)
y_train = ((x_train[:, 0] + x_train[:, 1]) > 0.7).astype("float32")

model = keras.Sequential(
    [
        keras.Input(shape=(4,), name="features"),
        keras.layers.Dense(8, activation="relu", name="hidden"),
        keras.layers.Dense(1, activation="sigmoid", dtype="float32", name="score"),
    ]
)
model.compile(
    optimizer=keras.optimizers.Adam(learning_rate=0.02),
    loss=keras.losses.BinaryCrossentropy(),
    metrics=[keras.metrics.BinaryAccuracy(name="accuracy")],
)
history = model.fit(x_train, y_train, epochs=3, batch_size=4, verbose=0)
prediction = model.predict(x_train[:2], verbose=0)

print(f"backend: {keras.backend.backend()}")
print(f"global dtype policy: {keras.config.dtype_policy().name}")
print(f"hidden layer policy: {model.get_layer('hidden').dtype_policy.name}")
print(f"output layer policy: {model.get_layer('score').dtype_policy.name}")
print(f"optimizer: {model.optimizer.__class__.__name__}")
print(f"epochs completed: {len(history.history['loss'])}")
print(f"prediction dtype: {prediction.dtype}")
print(f"prediction shape: {prediction.shape}")

keras.config.set_dtype_policy("float32")