tf.math.argmax throws InvalidArgumentError with valid axis of int16 dtype #97096
Description
Issue type
Bug
Have you reproduced the bug with TensorFlow Nightly?
Yes
Source
source
TensorFlow version
2.20.0-dev20250715
Custom code
Yes
OS platform and distribution
Linux Ubuntu 20.04
Mobile device
No response
Python version
3.12
Bazel version
No response
GCC/compiler version
No response
CUDA/cuDNN version
No response
GPU model and memory
No response
Current behavior?
Running tf.math.argmax on a float16 tensor with large (overflowing) values and an axis=0 of int16 type gives an unexpected error about invalid dimension even though the dimension is a valid one. The issue can be recreated without overflowing tensors on tf.experimental.numpy.argmax, although inconsistently (sometimes the error is there, sometimes it isn't. See the colab below.)
Also reproducible with Tensorflow 2.19 colab
Standalone code to reproduce the issue
import tensorflow as tf
import numpy as np
print(tf.__version__) # 2.20.0-dev20250715
rng = np.random.default_rng(311)
input_tensor = tf.constant(rng.uniform(-np.finfo(np.float16).max-1, 0., size=(1, 2)), dtype=tf.float16)
axis = tf.constant(0, dtype=tf.int16)
output_type = tf.int16
output = tf.math.argmax(input_tensor, axis=axis, output_type=output_type)Relevant log output
File ".../venv/lib/python3.12/site-packages/tensorflow/python/util/traceback_utils.py", line 153, in error_handler
raise e.with_traceback(filtered_tb) from None
File ".../venv/lib/python3.12/site-packages/tensorflow/python/framework/ops.py", line 6027, in raise_from_not_ok_status
raise core._status_to_exception(e) from None # pylint: disable=protected-access
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
tensorflow.python.framework.errors_impl.InvalidArgumentError: {{function_node __wrapped__ArgMax_device_/job:localhost/replica:0/task:0/device:CPU:0}} Expected dimension in the range [-2, 2), but got -52363264 [Op:ArgMax] name: