Keras Inconsistent broadcasting behaviour in layers.Rescaling

It seems that layers.Rescaling will broadcast its input to match the size of the provided gain/offset parameters.

This may or may not be expected behaviour (it's not clear from the documentation).

What definitely seems wrong, however, is that this broadcasting is only occurring if I make a prediction with the model, but not when I print model.summary(), or when I query the model.output_shape parameter.

If I train the model against a target with a different shape, no errors are raised. This also does not seem like it should be expected behaviour.

Code to reproduce:

import numpy as np
import keras
from keras import layers

x = np.zeros((8, 100, 1))
y_target = np.zeros((8, 100, 1))


model = keras.Sequential(
    [
        layers.Input(shape=(None, 1), name="Feature"),
        layers.Rescaling([1.0, 1.0], [0.0, 0.0]),
    ]
)
model.summary()

y_pred = model(x)

print(f"Expected model output shape: {model.output_shape}")    
print(f"True model output shape: {y_pred.shape}")
print(f"Target shape: {y_target.shape}")

model.compile(
    optimizer="adam",
    loss="mse",
)
model.fit(x, y_target, epochs=1, batch_size=8)

print("Fit executed with no errors, despite inconsistent shapes.")

Versions: keras==3.8.0 jax==0.5.0

Comment From: sonali-kumari1

Hi @optiluca -

Thanks for reporting this. I have tested your code with both list inputs ([1.0, 1.0]) and scalar inputs (1.0) in this gist with the latest version of keras(3.10.0). Scalar values for scale and offset produces consistent output shapes which matches model.summary() and model.output_shape. However, when using lists for scale and offset whose length doesn’t match the inputs(as in your case), rescaling layer seems to be silently broadcasting the shapes to match the scale and offset values. We will look into this and update you.