import os

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

import keras
import numpy as np
from keras import layers

keras.utils.set_random_seed(42)