import os

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

import jax
import jax.numpy as jnp
import keras
import numpy as np


keras.utils.set_random_seed(7)

x_train = np.array(
    [
        [0.0, 0.0],
        [0.0, 1.0],
        [1.0, 0.0],
        [1.0, 1.0],
    ],
    dtype="float32",
)
y_train = np.array([[0.0], [1.0], [1.0], [2.0]], dtype="float32")