import os
from pathlib import Path

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

import keras
import numpy as np
from keras import ops


keras.utils.set_random_seed(29)

sample = np.array(
    [
        [0.2, 0.5, 0.1],
        [0.8, 0.3, 0.4],
    ],
    dtype="float32",
)