keras/src/backend/tensorflow/numpy.py/argmax
def argmax(x, axis=None, keepdims=False):
    _x = x
    if axis is None:
        x = tf.reshape(x, [-1])
    y = tf.cast(tf.argmax(x, axis=axis), dtype="int32")
    if keepdims:
        y = _keepdims(_x, y, axis)
    return y