import os
from pathlib import Path

os.environ.setdefault("KERAS_BACKEND", "jax")

import keras
import numpy as np


keras.utils.set_random_seed(31)
rng = np.random.default_rng(31)

x_train = rng.normal(size=(64, 256)).astype("float32")
y_train = rng.normal(size=(64, 16)).astype("float32")
x_eval = rng.normal(size=(8, 256)).astype("float32")