Summary

keras.layers.IntegerLookup(output_mode="one_hot") returns different output shapes for eager tensors vs. Keras.Input symbolic tensors.

Reproduction

import keras
import numpy as np
import tensorflow as tf

layer = keras.layers.IntegerLookup(
    vocabulary=[12, 36, 1138, 42],
    num_oov_indices=3,
    mask_token=-1,
    output_mode="one_hot",
)

x = tf.constant([[12, 36, 1138], [42, 999, 12]], dtype=tf.int64)
print(layer(x).shape)  # eager

inp = keras.Input(shape=(3,), dtype="int64")
print(layer(inp).shape)  # symbolic

Expected

Both calls should agree on the output shape. For one_hot, the symbolic shape should preserve the input rank, e.g. (None, 3, 7).

Actual

• Eager tensor output shape: (2, 3, 7)
• Symbolic input output shape: (None, 7)

This indicates incorrect shape inference for IntegerLookup in one_hot mode when used with Keras.Input.