Autograph: unclear AssertionError when using tf.size(x) inside a @tf.function branch

Issue type

Bug

Have you reproduced the bug with TensorFlow Nightly?

Yes

Source

source

TensorFlow version

2.19

Custom code

Yes

OS platform and distribution

Ubuntu 20.04

Mobile device

No response

Python version

3.8

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.size(x) inside a conditional branch of a @tf.function, Autograph sometimes raises an internal AssertionError instead of a meaningful user-facing error:

AssertionError: [1, 0]

The traceback points to an Autograph-generated internal file (__autograph_generated_fileXXXX.py), which is very confusing for end users.

Standalone code to reproduce the issue

import tensorflow as tf
import random

# Randomly choose input size
size = random.choice([0, 1, 3])
input_data = tf.constant([1] * size, dtype=tf.float32) if size > 0 else tf.constant([], dtype=tf.float32)

@tf.function
def process_data(x):
    # Problematic branch: tf.size(x) used in Python if
    traced_args = tf.unstack(x) if tf.size(x) > 0 else []
    tf.autograph.trace(*traced_args)
    return x * 2 + 1

output = process_data(input_data)
print("Input size:", size, "Output:", output)

Relevant log output

AssertionError: [1, 0]