我在Keras中定义了以下层：

class TextDec(tf.keras.Model):
    def __init__(
        self,
        n_vocab: int,
        n_ctxs: int,
        n_states: int,
        n_heads: int,
        n_layers: int,
    ):
        super(TextDec, self).__init__(name="dec")

        self.token_emb = tf.keras.layers.Embedding(
            n_vocab, n_states, name="dec-token-emb"
        )

        self.position_emb = tf.Variable(
            np.zeros((n_ctxs, n_states)),
            name="dec-position-emb",
            dtype=tf.float32,
        )

        self.add = tf.keras.layers.Add(name="dec-add")

        # ...

    def call(self, inputs: List[tf.Tensor]):
        tokens, audio_features = inputs
        offset = 0

        x = self.add(
            [
                self.token_emb(tokens),
                tf.slice(self.position_emb, [offset, 0], [tokens.shape[-1], -1]),
            ]
        )

        # ...

        return x

当我用n_ctxs = 448, n_states = 512调用这个函数时，它不起作用，我得到了：
ValueError：无法合并具有不同批大小的Tensor。获取形状为[（5，1，512），（1，512）]的Tensor
我还尝试使用tf.expand_dims，但导致了不同的错误：

x = self.add([ self.token_emb(tokens),  tf.expand_dims(tf.slice(self.position_emb, [offset, 0], [tokens.shape[-1], -1]), 0) ])

ValueError：无法合并具有不同批大小的Tensor。获取形状为[（5，1，512），（1，1，512）]的Tensor
我还尝试了this answer中建议的tf.compat.v1.placeholder_with_default，尽管它已被弃用：

z = tf.slice(self.position_emb, [offset, 0], [tokens.shape[-1], -1])
z = tf.compat.v1.placeholder_with_default(z, [None, 1, 512])
x = self.add([ self.token_emb(tokens),  z ])

那也没用。
值错误：形状的排名必须相等，但对于具有输入形状的“{{node PlaceholderWithDefault}} = PlaceholderWithDefaultdtype=DT_FLOAT，shape=[？，1，512]”，排名必须为2和3：[1 512]。
如果形状包含None，则tf.broadcast_to不起作用。
但是，如果我将这行代码改为：

x = self.token_emb(tokens) + tf.slice(self.position_emb, [offset, 0], [tokens.shape[-1], -1])

那么它就可以工作了。为什么当我使用Add时它不能工作，我该如何解决这个问题？
我想使用Add的原因是，当我总结我的模型时，我可以看到加法显示为这样的一行：

Model: "dec"
__________________________________________________________________________________________________
 Layer (type)                   Output Shape         Param #     Connected to                     
==================================================================================================
 tokens (InputLayer)            [(None, 1)]          0           []                               
                                                                                                  
 dec-token-emb (Embedding)      (None, 1, 512)       26554368    ['tokens[0][0]']                 
                                                                                                  
 tf.__operators__.add (TFOpLamb  (None, 1, 512)      0           ['dec-token-emb[0][0]']          
 da)

看到TFOpLambda让我担心这会阻止我序列化我的模型。