Keras Discretization layer returns inconsistent result in graph mode.

When using the Discretization layer in tensorflow backend you get a different output when calling in graph mode (model.predict() below) as opposed to eager mode.

Keras version: 3.10.0 Tensorflow version: 2.18.1

Reproducible code (gist):

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]])>

# eager mode
model(x)

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

# graph mode
model.predict(x)

array([[2, 2, 2, 2],
       [2, 2, 2, 5]])

This issue is copied modified upon #20981 as requested.

Comment From: divyashreepathihalli

This is a bug and it is being fixed here - https://github.com/keras-team/keras/pull/21514

Comment From: LakshmiKalaKadali

Hi @takumiohym,

The bug has been resolved in the latest Keras version 3.12.0 through this commit. The graph mode results now exactly align with eager mode, as verified in the provided gist.

Thank You

Comment From: takumiohym

@LakshmiKalaKadali Great! Thanks for taking care of this @danielenricocahall