Bug Report: XLA compilation fails when tf.Variable is created inside @tf.function-decorated method

Issue type

Bug

Have you reproduced the bug with TensorFlow Nightly?

Yes

Source

source

TensorFlow version

tf 2.20.0

Custom code

Yes

OS platform and distribution

Linux Ubuntu 24.04

Mobile device

No response

Python version

3.12

Bazel version

No response

GCC/compiler version

No response

CUDA/cuDNN version

No response

GPU model and memory

No response

Current behavior?

When using @tf.function(jit_compile=True) to compile a model that creates tf.Variable inside the call() method, a ValueError occurs. The error indicates that tf.Variables cannot be created within a tf.function context.

Standalone code to reproduce the issue

import tensorflow as tf

@tf.function(jit_compile=True)
def broken_function():

    v = tf.Variable(0.0)
    v.assign(1.0)
    return v

def main():
    try:
        result = broken_function()
        print("Success:", result)
    except Exception as e:
        print("Error:", e)

if __name__ == '__main__':
    main()

Relevant log output

Error: in user code:

    File ... line 6, in broken_function  *
        v = tf.Variable(0.0)

    ValueError: tf.function only supports singleton tf.Variables created on the first call. Make sure the tf.Variable is only created once or created outside tf.function. See https://www.tensorflow.org/guide/function#creating_tfvariables for more information.