import os

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

import keras
import numpy as np

keras.utils.set_random_seed(11)

x_train = np.array(
    [
        [0.0, 0.0, 0.1],
        [0.0, 1.0, 0.4],
        [1.0, 0.0, 0.6],
        [1.0, 1.0, 0.9],
    ],
    dtype="float32",
)
y_train = np.array([0.0, 1.0, 1.0, 0.0], dtype="float32")