XLA jit_compile=True drops NaN in tf.reduce_max while eager returns NaN Issue type

Issue type

Bug

Have you reproduced the bug with TensorFlow Nightly?

Yes

Source

source

TensorFlow version

2.21

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 behave differently for tf.reduce_max when the input contains NaN.

The input tensor contains finite values and one NaN. Eager execution returns NaN, while the same reduction compiled with jit_compile=True returns 11.0.

In the reproducer below, eager returns nan, but XLA returns the maximum finite value.

Expected behavior?

tf.reduce_max should handle NaN consistently between eager execution and jit_compile=True. XLA should return NaN if eager execution returns NaN 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


vals = tf.constant([7.0, np.nan, -2.5, 11.0, 0.25], dtype=tf.float32)

try:
    eager = tf.reduce_max(vals).numpy()
    eager_err = None
except Exception as e:
    eager = None
    eager_err = type(e).__name__ + ": " + str(e)[:120]

try:
    jit = tf.function(lambda x: tf.reduce_max(x), jit_compile=True)(vals).numpy()
    jit_err = None
except Exception as e:
    jit = None
    jit_err = type(e).__name__ + ": " + str(e)[:120]

print(f"eager reduce_max: {eager}  error={eager_err}")
print(f"jit   reduce_max: {jit}  error={jit_err}")

if eager_err is None and jit_err is None:
    if bool(np.isnan(eager)) != bool(np.isnan(jit)):
        print("BUG REPRODUCED: reduce_max NaN behavior differs between eager and jit_compile=True")
    else:
        print("not reproduced")
elif eager_err != jit_err:
    print("BUG REPRODUCED: reduce_max error behavior differs between eager and jit_compile=True")
else:
    print("not reproduced")

Relevant log output

eager reduce_max: nan error=None
jit reduce_max: 11.0 error=None
BUG REPRODUCED: reduce_max NaN behavior differs between eager and jit_compile=True