tf.sort gives different NaN behaviour in eager mode and XLA

Issue type

Bug

Have you reproduced the bug with TensorFlow Nightly?

No

Source

source

TensorFlow version

2.20.0-dev0+selfbuilt

Custom code

Yes

OS platform and distribution

Linux CPU-only VM, 8GB RAM, no GPU

Mobile device

No response

Python version

3.12.3

Bazel version

bazel 9.1.0

GCC/compiler version

gcc (Ubuntu 13.3.0-6ubuntu2~24.04.1) 13.3.0

CUDA/cuDNN version

No CUDA/cuDNN. CPU-only environment.

GPU model and memory

No GPU. CPU-only VM with 8GB RAM.

Current behavior?

tf.sort gives different results in eager execution and under jit_compile=True when the input contains NaNs.

I understand that NaN comparisons can make sorting behaviour subtle, so I am not sure whether this is intended. However, the eager result looks surprising because the finite values are not sorted either: 3.0 is left after 7.0.

This may be related to #117287, but this is a smaller reproducer focused directly on tf.sort.

Expected behavior: either eager and XLA should agree, or the TensorFlow documentation should make clear that tf.sort with NaNs has undefined or implementation-dependent behaviour.

Standalone code to reproduce the issue

import tensorflow as tf

x = tf.constant([4.0, float("nan"), 7.0, 1.0, float("nan"), 3.0])

eager = tf.sort(x).numpy()

@tf.function(jit_compile=True)
def sort_xla(t):
    return tf.sort(t)

xla = sort_xla(x).numpy()

print("TensorFlow:", tf.__version__)
print("input :", x.numpy())
print("eager :", eager)
print("xla   :", xla)

Relevant log output

TensorFlow: 2.20.0-dev0+selfbuilt
input : [ 4. nan  7.  1. nan  3.]
eager : [ 1.  4. nan  7. nan  3.]
xla   : [ 1.  3.  4.  7. nan nan]