XLA jit_compile=True returns different argmax index on NaN input than eager execution

Issue type

Bug

Have you reproduced the bug with TensorFlow Nightly?

Yes

Source

source

TensorFlow version

2.21.0

Custom code

Yes

OS platform and distribution

Linux Ubuntu 22.04

Mobile device

No response

Python version

3.11

Bazel version

No response

GCC/compiler version

No response

CUDA/cuDNN version

No response

GPU model and memory

No response

Current behavior?

TensorFlow eager and XLA return different results for tf.math.argmax when the input contains NaN.

The input tensor is [4.0, 5.5, NaN]. Eager execution returns index 1, corresponding to the finite value 5.5. The same function compiled with jit_compile=True returns index 2, corresponding to the NaN value.

In the reproducer below, eager returns 1, while XLA returns 2.

Expected behavior?

tf.math.argmax should handle NaN consistently between eager execution and jit_compile=True for the same input.

Standalone code to reproduce the issue

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"

import tensorflow as tf
import numpy as np


t = tf.constant([4.0, 5.5, np.nan], dtype=tf.float32)

eager = int(tf.math.argmax(t).numpy())
jit = int(tf.function(lambda v: tf.math.argmax(v), jit_compile=True)(t).numpy())

print(f"eager argmax: {eager}")
print(f"jit   argmax: {jit}")

if eager != jit:
    print("BUG REPRODUCED: argmax NaN behavior differs between eager and jit_compile=True")
else:
    print("not reproduced")

Relevant log output

eager argmax: 1
jit argmax: 2
BUG REPRODUCED: argmax NaN behavior differs between eager and jit_compile=True