XLA produces different output for tf.while_loop when loop body depends on Python mutable state

Issue type

Bug

Have you reproduced the bug with TensorFlow Nightly?

No

Source

source

TensorFlow version

2.20.0

Custom code

Yes

OS platform and distribution

Linux Ubuntu 20.04

Mobile device

No response

Python version

3.9

Bazel version

No response

GCC/compiler version

No response

CUDA/cuDNN version

No response

GPU model and memory

No response

Current behavior?

When a Keras model uses tf.while_loop with Python mutable state, such as a boolean flag toggled inside loop_cond that controls branching in loop_body, eager execution and XLA-compiled execution produce different numerical outputs.

In eager mode, the example below returns [3.]. Under @tf.function(jit_compile=True) or tf.function with XLA autoclustering enabled, the same logic returns [-3.].

The difference appears to come from how tf.while_loop traces cond() and body() to build the computation graph. During tracing, the Python-level self.flag is evaluated at trace time, and the branch selected by if self.flag is embedded into the graph. The eager and XLA paths appear to embed different branches, producing different numerical outputs for the same input.

Expected: all three execution modes, eager, jit_compile=True, and XLA autoclustering, should produce the same output [3.] for the same input.

Actual:

Eager output:                  tf.Tensor([3.], shape=(1,), dtype=float32)
XLA (jit_compile=True) output: tf.Tensor([-3.], shape=(1,), dtype=float32)
AutoCluster output:            tf.Tensor([-3.], shape=(1,), dtype=float32)

Standalone code to reproduce the issue

import os

import tensorflow as tf


class Model(tf.keras.Model):
    def __init__(self):
        super().__init__()
        self.flag = False

    def cond(self, i, x):
        self.flag = not self.flag
        return tf.less(i, tf.constant(3))

    def body(self, i, x):
        if self.flag:
            return tf.add(i, 1), tf.multiply(x, -1.0)
        return tf.add(i, 1), x

    def call(self, x):
        _, result = tf.while_loop(self.cond, self.body, [tf.constant(0), x])
        return result


class ModelXLA(tf.keras.Model):
    def __init__(self):
        super().__init__()
        self.flag = False

    def cond(self, i, x):
        self.flag = not self.flag
        return tf.less(i, tf.constant(3))

    def body(self, i, x):
        if self.flag:
            return tf.add(i, 1), tf.multiply(x, -1.0)
        return tf.add(i, 1), x

    @tf.function(jit_compile=True)
    def call(self, x):
        _, result = tf.while_loop(self.cond, self.body, [tf.constant(0), x])
        return result


class ModelAutoCluster(tf.keras.Model):
    def __init__(self):
        super().__init__()
        self.flag = False

    def cond(self, i, x):
        self.flag = not self.flag
        return tf.less(i, tf.constant(3))

    def body(self, i, x):
        if self.flag:
            return tf.add(i, 1), tf.multiply(x, -1.0)
        return tf.add(i, 1), x

    @tf.function
    def call(self, x):
        _, result = tf.while_loop(self.cond, self.body, [tf.constant(0), x])
        return result


x = tf.constant([3.0])

print("Eager output:                 ", Model()(x))
print("XLA (jit_compile=True) output:", ModelXLA()(x))

os.environ["TF_XLA_FLAGS"] = "--tf_xla_auto_jit=2 --tf_xla_cpu_global_jit"
print("AutoCluster output:           ", ModelAutoCluster()(x))
os.environ["TF_XLA_FLAGS"] = ""

Relevant log output

Eager output:                  tf.Tensor([3.], shape=(1,), dtype=float32)
XLA (jit_compile=True) output: tf.Tensor([-3.], shape=(1,), dtype=float32)
AutoCluster output:            tf.Tensor([-3.], shape=(1,), dtype=float32)