`tf.function(jit_compile=True)` **silently returns wrong result** for a size-0 `dynamic_size` TensorArray written in a `while_loop`

Issue type

Bug

Have you reproduced the bug with TensorFlow Nightly?

No

Source

binary

TensorFlow version

v2.21.0-rc1-5-ga481b10260d (2.21.0)

Custom code

Yes

OS platform and distribution

Linux Ubuntu 22.04(kernel 6.8, glibc 2.35)

Mobile device

No response

Python version

3.12.13

Bazel version

No response

GCC/compiler version

No response

CUDA/cuDNN version

No response

GPU model and memory

NVIDIA RTX A6000 (48 GB)

Current behavior?

Writing to a tf.TensorArray(size=0, dynamic_size=True) inside a tf.while_loop works in eager but, under tf.function(jit_compile=True) (XLA), compiles successfully and silently returns an empty tensor instead of the written elements. XLA appears to freeze the array's static size at the initial size (0) and ignore the dynamic growth, so stack() yields shape (0,). Only the size=0 + dynamic_size=True case misbehaves (size=2 or a static array are correct); it reproduces on both CPU and GPU. Expected: under jit_compile=True XLA must compute the correct result ([True, True]) or reject compilation — it must not compile and silently change the output.

Standalone code to reproduce the issue

import tensorflow as tf

def w():
    ta = tf.TensorArray(tf.bool, size=0, dynamic_size=True)
    def cond(i, ta): return i < 2
    def body(i, ta):
        ta = ta.write(i, tf.constant(True))
        return i + 1, ta
    _, out = tf.while_loop(cond, body, [tf.constant(0), ta])
    return out.stack()

print(w())                                   # eager: [ True  True]
print(tf.function(w, jit_compile=True)())    # XLA:   []   <-- silently wrong, no error

Relevant log output

tf.Tensor([ True  True], shape=(2,), dtype=bool)     # eager
tf.Tensor([], shape=(0,), dtype=bool)                # tf.function(jit_compile=True)  -- silent wrong result