import os

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

import keras
import numpy as np

keras.utils.set_random_seed(7)

rng = np.random.default_rng(7)
x_train = rng.uniform(size=(64, 2)).astype("float32")
y_train = (x_train[:, 0] > x_train[:, 1]).astype("float32")

x_test = rng.uniform(size=(16, 2)).astype("float32")
y_test = (x_test[:, 0] > x_test[:, 1]).astype("float32")