model.predict() gives different result to model() when using Discretization layer

[georyetti]

When using the Discretization layer in tensorflow backend you get a different output when calling model.predict() as opposed to the layer itself or using the models call method directly.

Keras version: 3.2.1
Tensorflow version: 2.16.1
Numpy version: 1.23.5

Reproducible code

import tensorflow as tf
import keras

layer = keras.layers.Discretization(
    bin_boundaries=[-0.5, 0, 0.1, 0.2, 3],
    name="bucket",
    output_mode="int",
)

x = tf.constant([[0.0, 0.15, 0.21, 0.3], [0.0, 0.17, 0.451, 7.8]])
inputs = keras.layers.Input(name="inp", dtype="float32", shape=(4,))

model_output = layer(inputs)
model = keras.models.Model(inputs=[inputs], outputs=[model_output])

layer(x)
<tf.Tensor: shape=(2, 4), dtype=int64, numpy=
array([[2, 3, 4, 4],
       [2, 3, 4, 5]])>

model(x)
<tf.Tensor: shape=(2, 4), dtype=int64, numpy=
array([[2, 3, 4, 4],
       [2, 3, 4, 5]])>

model.predict(x)
array([[2, 2, 2, 2],
       [2, 2, 2, 5]])

I've checked with later keras versions (3.4.0) and the issue occurs.

[CloudPiyush]

Hi @georyetti

I saw your issue regarding the Discretization layer in Keras producing inconsistent outputs between model.predict() and direct calls to the layer. I investigated the problem and found that it was likely caused by bug in Keras Discretization layer when running in graph mode (model.predict()).
Solution : The best immediate workaround is forcing eager execution or using tf.bucketize manually.

I have a worked solution that ensures consistent outputs across all methods. You can check out the corrected implementation here: fixed version of the code that ensures the Discretization layer works correctly under both eager execution and graph mode (model.predict()).

import tensorflow as tf
import keras


bin_boundaries = [-0.5, 0, 0.1, 0.2, 3]

def discretize(x):
    
    return tf.raw_ops.Bucketize(input=x, boundaries=bin_boundaries)

inputs = keras.layers.Input(name="inp", dtype="float32", shape=(4,))

model_output = keras.layers.Lambda(discretize, output_shape=(4,))(inputs)
model = keras.models.Model(inputs=[inputs], outputs=[model_output])

x = tf.constant([[0.0, 0.15, 0.21, 0.3], [0.0, 0.17, 0.451, 7.8]], dtype=tf.float32)

print("Direct Layer Call:", discretize(x).numpy())
print("Model Call:", model(x).numpy())
print("Model Predict:", model.predict(x))

Let me know if this helps or if you’d like any further clarification

Best,
Piyush

[dhantule]

Hi @georyetti, Thanks for reporting this.

I've tested your code with Keras 3.9.0 and faced the same issue with model() and model.predict() results, however running it eagerly gave consistent results in this gist.

[CloudPiyush]

Hi @dhantule
Code I Committed is also working you can check