XLA 编译失败: tf.function 因在 call() 中创建 tf.Variable 而失败

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[0]) print(output)