XLA jit_compile=True over-propagates NaN in tf.image.resize bilinear output

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 produce different NaN propagation behavior for tf.image.resize with bilinear interpolation.

The input image contains a single NaN value in channel 0. Eager execution only propagates NaN to the local interpolation region, producing 6 NaN values in the resized channel. However, the same function compiled with jit_compile=True produces NaN for the entire resized channel, resulting in 25 NaN values.

In the reproducer below, eager produces a partially finite 5x5 output, while XLA returns all NaNs for channel 0.

Expected behavior?

tf.image.resize with bilinear interpolation should handle NaN propagation consistently between eager execution and jit_compile=True for the same input.

Standalone code to reproduce the issue

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

import sys
import numpy as np
import tensorflow as tf

data = np.arange(1, 28, dtype=np.float32).reshape(1, 3, 3, 3)
data[0, 1, 2, 0] = np.nan
x = tf.constant(data)

def f(v):
    return tf.image.resize(v, (5, 5), method="bilinear")

eager = f(x).numpy()
xla = tf.function(f, jit_compile=True)(x).numpy()

eager_count = int(np.isnan(eager[0, :, :, 0]).sum())
xla_count = int(np.isnan(xla[0, :, :, 0]).sum())

print("input_channel0:", data[0, :, :, 0])
print("eager_nan_count_channel0:", eager_count)
print("xla_nan_count_channel0:", xla_count)
print("eager_channel0:", eager[0, :, :, 0])
print("xla_channel0:", xla[0, :, :, 0])

if eager_count != xla_count:
    print("BUG REPRODUCED: resize bilinear NaN propagation differs between eager and jit_compile=True")
    sys.exit(0)

print("not reproduced")
sys.exit(1)

Relevant log output

input_channel0:
[[ 1.  4.  7.]
 [10. 13. nan]
 [19. 22. 25.]]

eager_nan_count_channel0: 6
xla_nan_count_channel0: 25
BUG REPRODUCED: resize bilinear NaN propagation differs between eager and jit_compile=True