XLA compilation failure: tf.function fails due to tf.Variable creation inside call()

import tensorflow as tf

class TestModel(tf.keras.Model):
    def __init__(self):
        super().__init__()
        self.embedding_column_0 = tf.feature_column.embedding_column(
            tf.feature_column.categorical_column_with_vocabulary_list(
                key="feature_0", vocabulary_list=["a", "b", "c"]
            ),
            dimension=8
        )
        self.embedding_column_1 = tf.feature_column.embedding_column(
            tf.feature_column.categorical_column_with_vocabulary_list(
                key="feature_1", vocabulary_list=["x", "y"]
            ),
            dimension=5
        )
        self.embedding_column_2 = tf.feature_column.embedding_column(
            tf.feature_column.categorical_column_with_vocabulary_list(
                key="feature_2", vocabulary_list=["p", "q", "r", "s"]
            ),
            dimension=10
        )
        self.dense = tf.keras.layers.Dense(16, activation='relu')
        self.output_layer = tf.keras.layers.Dense(1)

    def call(self, inputs):
        features = {
            'feature_0': inputs[:, 0:1],
            'feature_1': inputs[:, 1:2],
            'feature_2': inputs[:, 2:3]
        }
        embedded_0 = tf.nn.embedding_lookup(
            tf.Variable(tf.random.normal([3, 8])), 
            tf.cast(features['feature_0'], tf.int32)
        )
        embedded_1 = tf.nn.embedding_lookup(
            tf.Variable(tf.random.normal([2, 5])), 
            tf.cast(features['feature_1'], tf.int32)
        )
        embedded_2 = tf.nn.embedding_lookup(
            tf.Variable(tf.random.normal([4, 10])), 
            tf.cast(features['feature_2'], tf.int32)
        )
        combined = tf.concat([embedded_0, embedded_1, embedded_2], axis=-1)
        x = self.dense(combined)
        return self.output_layer(x)

def get_default_model():
    return TestModel()

def get_sample_inputs():
    x = tf.random.uniform([8, 3], minval=0, maxval=10, dtype=tf.int32) % tf.constant([3, 2, 4])
    return (x,)

def main():
    model = get_default_model()
    inputs = get_sample_inputs()
    output = model(*inputs)
    print("Input shape:", inputs[0].shape)
    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()

Input shape: (8, 3)
Output shape: (8, 1, 1)
Traceback (most recent call last):
 line 36, in call
        tf.Variable(tf.random.normal([3, 8])),

    ValueError: Exception encountered when calling TestModel.call().
    
    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.
    
    Arguments received by TestModel.call():
      • inputs=tf.Tensor(shape=(8, 3), dtype=int32)