XLA compilation fails when using tf.sequence_mask with SymbolicTensor loop length

Issue type

Bug

Have you reproduced the bug with TensorFlow Nightly?

Yes

Source

binary

TensorFlow version

2.20.0

Custom code

Yes

OS platform and distribution

linux ubuntu 24.04

Mobile device

No response

Python version

3.9

Bazel version

No response

GCC/compiler version

No response

CUDA/cuDNN version

CUDA Version 13.0

GPU model and memory

No response

Current behavior?

Compiling a model with XLA (@tf.function(jit_compile=True)) fails when tf.sequence_mask is used with a loop length derived from a SymbolicTensor. In eager mode the model works correctly, causing a behavior inconsistency.

Standalone code to reproduce the issue

import tensorflow as tf
import random

class TestModel(tf.keras.Model):

    def __init__(self):
        super().__init__()
        self.embedding = tf.keras.layers.Embedding(100, 32, mask_zero=True)
        self.lstm = tf.keras.layers.LSTM(64, return_sequences=True)
        self.dense = tf.keras.layers.Dense(16)

    def call(self, x, sequence_lengths=None):
        embedded = self.embedding(x)
        if sequence_lengths is not None:
            mask = tf.sequence_mask(tf.constant([random.randint(1, 10 - 1) for _ in range(tf.shape(sequence_lengths)[0])]), maxlen=tf.shape(embedded)[1])
            embedded = embedded * tf.cast(mask[:, :, tf.newaxis], embedded.dtype)
        lstm_out = self.lstm(embedded)
        return self.dense(lstm_out)

def get_default_model():
    return TestModel()

def get_sample_inputs():
    batch_size = 8
    max_len = 10
    x = tf.random.uniform([batch_size, max_len], maxval=100, dtype=tf.int32)
    sequence_lengths = tf.constant([random.randint(1, max_len - 1) for _ in range(batch_size)])
    return (x, sequence_lengths)

def main():
    model = get_default_model()
    inputs = get_sample_inputs()
    output = model(*inputs)
    print('Input shape:', inputs[0].shape)
    print('Sequence lengths:', inputs[1])
    print('Output shape:', output.shape)
    @tf.function(jit_compile=True)
    def compiled_forward(*args):
        return model(*args)
    compiled_out = compiled_forward(*inputs)
    print('XLA Output shape:', compiled_out.shape)

if __name__ == '__main__':
    main()

Relevant log output

Input shape: (8, 10)
Sequence lengths: tf.Tensor([1 4 6 5 1 8 3 9], shape=(8,), dtype=int32)
Output shape: (8, 10, 16)
Traceback (most recent call last):
 line 15, in call
        mask = tf.sequence_mask(tf.constant([random.randint(1, 10 - 1) for _ in range(tf.shape(sequence_lengths)[0])]), maxlen=tf.shape(embedded)[1])

    TypeError: Exception encountered when calling TestModel.call().
    
    'SymbolicTensor' object cannot be interpreted as an integer
    
    Arguments received by TestModel.call():
      • x=tf.Tensor(shape=(8, 10), dtype=int32)
      • sequence_lengths=tf.Tensor(shape=(8,), dtype=int32)