Discretization using one_hot as output_model doesn't work as expected

[kianyangchn]

Environment: Py310, TF2.16.1, keras 3.7

Hi I'm trying to add a simple discretization layer(with pre-defined bin_boundries and one_hot output_mode) in a toy model.
The expected output shape of this layer in the model is (None, 1) instead of (None, len(bin_boundries)+1).

import tensorflow as tf
from tensorflow.keras.layers import Layer, Input, Discretization, Dense
from tensorflow.keras.models import Model

# Define the input layer
inputs = Input(shape=(1,), dtype=tf.float32, name="f1")

# Add the custom discretization layer
discretization_output = Discretization(bin_boundaries=[1, 2, 3, 4, 5], output_mode="one_hot")(inputs)

# Print the shape of the discretization output
print(discretization_output.shape)

outputs = Dense(1)(discretization_output)

# Create the model
model = Model(inputs=inputs, outputs=outputs)
model.summary()

┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┓
┃ Layer (type)                    ┃ Output Shape           ┃       Param # ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━┩
│ f1 (InputLayer)                 │ (None, 1)              │             0 │
├─────────────────────────────────┼────────────────────────┼───────────────┤
│ discretization_2                │ (None, 1)              │             0 │
│ (Discretization)                │                        │               │
├─────────────────────────────────┼────────────────────────┼───────────────┤
│ dense_1 (Dense)                 │ (None, 1)              │             2 │
└─────────────────────────────────┴────────────────────────┴───────────────┘

The actual output shape is correct, just the model expects the output to be (None, 1), which leading to the error
expected axis -1 of input shape to have value 1, but received input with shape (None, 6)
Seems a bug in compute_output_shape