import numpy as np
import keras

from custom_layers import ScaleLayer


inputs = keras.Input(shape=(3,), name="features")
x = ScaleLayer(0.5, name="scale_features")(inputs)
outputs = keras.layers.Dense(
    1,
    kernel_initializer="ones",
    bias_initializer="zeros",
    name="score",
)(x)
model = keras.Model(inputs, outputs)

sample = np.ones((2, 3), dtype="float32")
model(sample)
model.save("custom_scale_model.keras")

print("saved: custom_scale_model.keras")
print(f"custom layer: {model.get_layer('scale_features').__class__.__name__}")
print(f"prediction shape: {np.asarray(model(sample)).shape}")