XLA Compilation Error: Python range() fails with symbolic tensors in @tf.function(jit_compile=True)

Issue type

Bug

Have you reproduced the bug with TensorFlow Nightly?

Yes

Source

source

TensorFlow version

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 Python's range() function with a symbolic tensor (from tf.shape()) inside a @tf.function(jit_compile=True) context, XLA compilation fails with the error: 'SymbolicTensor' object cannot be interpreted as an integer. This is a common issue when migrating eager-mode code to XLA compilation.

Standalone code to reproduce the issue

import tensorflow as tf
tf.config.set_soft_device_placement(False)

class LSTMCell(tf.keras.layers.Layer):

    def __init__(self, units=64):
        super(LSTMCell, self).__init__()
        self.units = units
        self.dense_i = tf.keras.layers.Dense(units, activation='sigmoid')
        self.dense_f = tf.keras.layers.Dense(units, activation='sigmoid')
        self.dense_c = tf.keras.layers.Dense(units, activation='tanh')
        self.dense_o = tf.keras.layers.Dense(units, activation='sigmoid')

    def call(self, inputs, states):
        h_prev, c_prev = states
        combined = tf.concat([inputs, h_prev], axis=-1)
        i = self.dense_i(combined)
        f = self.dense_f(combined)
        c_candidate = self.dense_c(combined)
        o = self.dense_o(combined)
        c = f * c_prev + i * c_candidate
        h = o * tf.tanh(c)
        return (h, [h, c])

class LSTMModel(tf.keras.layers.Layer):

    def __init__(self, units=64):
        super(LSTMModel, self).__init__()
        self.units = units
        self.lstm_cell = LSTMCell(units)

    def call(self, x, training=False):
        batch_size = tf.shape(x)[0]
        seq_len = tf.shape(x)[1]
        h = tf.zeros([batch_size, self.units])
        c = tf.zeros([batch_size, self.units])
        outputs = tf.TensorArray(tf.float32, size=seq_len)
        for t in range(seq_len):
            x_t = x[:, t, :]
            h, (h, c) = self.lstm_cell(x_t, [h, c])
            outputs = outputs.write(t, h)
        outputs = tf.transpose(outputs.stack(), [1, 0, 2])
        return outputs

class TestModel(tf.keras.Model):

    def __init__(self):
        super(TestModel, self).__init__()
        self.lstm_model = LSTMModel()

    def call(self, x, training=False):
        output = self.lstm_model(x, training=training)
        return output

def get_default_model():
    return TestModel()

def get_sample_inputs():
    x = tf.random.normal([8, 10, 20])
    return (x,)

def main():
    model = get_default_model()
    inputs = get_sample_inputs()
    output = model(*inputs)
    print(f'input shape: {inputs[0].shape}')
    print(f'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, 20)
output shape: (8, 10, 64)
Traceback (most recent call last):
...

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