import os
from pathlib import Path

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

import keras
import numpy as np


keras.utils.set_random_seed(17)

model = keras.Sequential(
    [
        keras.Input(shape=(6,), name="features"),
        keras.layers.Dense(12, activation="relu", name="feature_projection"),
        keras.layers.Dense(3, activation="softmax", name="class_score"),
    ]
)

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