import os

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

import keras
import numpy as np


def seeded_prediction(seed):
    keras.utils.set_random_seed(seed)
    model = keras.Sequential(
        [
            keras.Input(shape=(3,)),
            keras.layers.Dense(5, activation="relu"),
            keras.layers.Dense(1, activation="sigmoid"),
        ]
    )
    sample = np.array([[0.25, 0.50, 0.75]], dtype="float32")
    return model.predict(sample, verbose=0), model.get_weights()[0].copy()


prediction_a, kernel_a = seeded_prediction(123)
prediction_b, kernel_b = seeded_prediction(123)
prediction_c, kernel_c = seeded_prediction(124)

print(f"backend: {keras.backend.backend()}")
print("seed used for first two runs: 123")
print(f"matching initial kernel: {np.allclose(kernel_a, kernel_b)}")
print(f"matching prediction: {np.allclose(prediction_a, prediction_b)}")
print(f"different seed changes kernel: {not np.allclose(kernel_a, kernel_c)}")
print(f"prediction with seed 123: {prediction_a[0, 0]:.6f}")
print(f"prediction with seed 124: {prediction_c[0, 0]:.6f}")