keras.ops.argmax returns incorrect shape for Keras Input when keepdims=True

[tosubmitissue]

Bug report

keras.ops.argmax(..., keepdims=True) produces the correct shape for eager tensors, but returns the original input shape when called on a Keras.Input placeholder.

Reproduction

import keras
import tensorflow as tf

x = tf.random.normal((3, 4, 5))
print(keras.ops.argmax(x, axis=None, keepdims=True).shape)

inp = keras.Input(shape=(4, 5), batch_size=3)
print(keras.ops.argmax(inp, axis=None, keepdims=True).shape)

Expected result

Both calls should return a reduced shape with singleton dimensions, i.e. (1, 1, 1) for this input.

Actual result

• Eager tensor: (1, 1, 1)
• Keras.Input placeholder: (3, 4, 5)

Notes

This makes argmax(..., keepdims=True) behave inconsistently between eager tensors and symbolic Keras tensors.

[daehyun99]

Hi, I’d like to work on this issue. I will summit PR soon. Thanks

[AGFACBNNR]

I found a silimar issue in keras.ops.argmin:

import keras
import tensorflow as tf

# Eager tensor
x = tf.ones((3, 4, 5))
y = keras.ops.argmin(x, axis=None, keepdims=True)
print(y.shape)

# Symbolic tensor
inp = keras.Input(shape=(4, 5))
z = keras.ops.argmin(inp, axis=None, keepdims=True)
print(z.shape)

Actual result

• eager output shape: (1, 1, 1)
• symbolic output shape: (None, 4, 5)

[divyashreepathihalli]

@keras-team /traige