tf.math.unsorted_segment_sum with num_segments=0 fails to compile under XLA ("Scatter dimension 0 is of size zero") while eager returns an empty tensor

Issue type

Bug

Have you reproduced the bug with TensorFlow Nightly?

Yes

Source

source

TensorFlow version

2.21.0 and tf-nightly 2.23.0-dev20260930 (CPU)

Custom code

Yes

OS platform and distribution

No response

Mobile device

No response

Python version

No response

Bazel version

No response

GCC/compiler version

No response

CUDA/cuDNN version

No response

GPU model and memory

No response

Current behavior?

tf.math.unsorted_segment_sum(data, segment_ids, num_segments=0) where all segment_ids are negative (negative ids are dropped) returns an empty tensor of shape (0,) in eager. Under jit_compile=True it fails to compile with InvalidArgumentError: Scatter dimension 0 is of size zero in tensor with shape f32[0].

The tf2xla segment-reduction kernel lowers this to an XLA scatter into a zero-size operand, which XLA's shape inference rejects. The kernel could return an empty result directly when num_segments == 0.

Expected: XLA returns an empty (0,) tensor like eager.

Standalone code to reproduce the issue

import tensorflow as tf

def f(ids):
    # negative segment ids are dropped, so the result is empty
    return tf.math.unsorted_segment_sum(tf.ones([1]), ids, num_segments=0)

ids = tf.constant([-2])
print("eager", tuple(f(ids).shape))
print("XLA", tuple(tf.function(f, jit_compile=True)(ids).shape))

Relevant log output

eager (0,)
tensorflow.python.framework.errors_impl.InvalidArgumentError: Scatter dimension 0 is of size zero in tensor with shape f32[0]