XLA simplify-while-loops incorrectly folds an integer tf.while_loop to its initial value

XLA simplify-while-loops incorrectly folds an integer tf.while_loop to its initial value

Issue type

Bug

Have you reproduced the bug with TensorFlow Nightly?

Yes

Source

binary

TensorFlow version

tf-nightly 2.22.0-dev20260531 (v1.12.1-140696-g1f5a6a050b4)

Custom code

Yes

OS platform and distribution

Linux Ubuntu 20.04.4 LTS

Mobile device

N/A

Python version

Python 3.11.15

Bazel version

N/A

GCC/compiler version

N/A

CUDA/cuDNN version

N/A. Reproduced on XLA CPU JIT.

GPU model and memory

N/A. Reproduced on XLA CPU JIT.

Current behavior?

An integer tf.while_loop executes once in eager mode and non-XLA tf.function, returning 2. The same function compiled with jit_compile=True incorrectly returns the initial loop value 0.

The loop starts with i = 0, its condition is i < 1, and its body computes i + 2, so the expected result is 2.

The issue reproduces with both int32 and int64 loop variables. The equivalent float32 loop returns the correct result, which suggests that the incorrect transformation is specific to integral loop variables.

HLO pass dumps show a valid while operation before optimization. The first incorrect dump appears after the second simplify-while-loops pass: the pass removes the loop and replaces its result with the initial value 0. For the float32 control case, the loop is correctly folded to 2.

Standalone code to reproduce the issue

import os

os.environ["CUDA_VISIBLE_DEVICES"] = "-1"

import tensorflow as tf


def f():
    limit = tf.constant(1, dtype=tf.int32)
    (result,) = tf.while_loop(
        lambda i: i < limit,
        lambda i: (i + tf.constant(2, dtype=tf.int32),),
        (tf.constant(0, dtype=tf.int32),),
    )
    return result


graph_f = tf.function(f)
xla_f = tf.function(f, jit_compile=True)

print("TensorFlow:", tf.__version__)
print("Eager:", f().numpy())
print("Graph:", graph_f().numpy())
print("XLA:", xla_f().numpy())

Relevant log output

TensorFlow: 2.22.0-dev20260531
Eager: 2
Graph: 2
XLA: 0

Additional dtype checks:

int32:   eager 2,   graph 2,   XLA 0
int64:   eager 2,   graph 2,   XLA 0
float32: eager 2.0, graph 2.0, XLA 2.0

The HLO pass sequence was captured with:

XLA_FLAGS="--xla_dump_to=/tmp/hlo --xla_dump_hlo_as_text --xla_dump_hlo_pass_re=.*" python repro.py

Immediately before the incorrect transformation, the optimized module still contains the loop:

%tuple = (s32[]) tuple(%initial)
%while = (s32[]) while(%tuple), condition=%condition, body=%body
%result = s32[] get-tuple-element(%while), index=0

Immediately after simplify-while-loops, the loop has been replaced by its initial tuple:

%initial = s32[] constant(0)
%tuple = (s32[]) tuple(%initial)
%result = s32[] get-tuple-element(%tuple), index=0