class PositionalEncoding(tf.keras.layers.Layer): def __init__(self, max_steps, max_dims, dtype=tf.float32, **kwargs): super().__init__(dtype=dtype, **kwargs) if max_dims % 2 == 1: max_dims += 1 p, i = np.meshgrid(np.arange(max_steps), np.arange(max_dims // 2)) pos_emb = np.empty((1, max_steps, max_dims)) pos_emb[0, :, ::2] = np.sin(p / 10000 ** (2 * i / max_dims)).T pos_emb[0, :, 1::2] = np.cos(p / 10000 ** (2 * i / max_dims)).T self.positional_embedding = tf.constant(pos_emb.astype(self.dtype)) def call(self, inputs): shape = tf.shape(inputs) return inputs + self.positional_embedding[:, :shape[-2], :shape[-1]] def get_config(self): config = super().get_config().copy() config.update({ 'dtype': self.dtype, }) return config

这是一个实现位置编码功能的自定义层`PositionalEncoding`，它可以用于Transformer等模型中。该层根据输入序列的位置信息，为每个位置生成一个固定的编码向量，从而使得模型能够更好地处理序列信息。

该层的初始化函数`__init__`接受三个参数：`max_steps`表示输入序列的最大长度，`max_dims`表示编码向量的最大维度，`dtype`表示数据类型。

在初始化函数中，首先判断`max_dims`是否为奇数，如果是，则将其加1，这是因为接下来要生成的编码向量需要将其拆分为两部分，每部分的维度应该是偶数。

接着，使用`np.meshgrid`函数生成两个二维数组`p`和`i`，其中`p`表示不同位置的序号，`i`表示编码向量的维度编号。然后，根据公式`PE(pos, 2i) = sin(pos / 10000^(2i/dims))`和`PE(pos, 2i+1) = cos(pos / 10000^(2i/dims))`，生成一个形状为`(1, max_steps, max_dims)`的数组`pos_emb`，其中每个元素表示一个位置的编码向量。

最后，使用`tf.constant`函数将数组`pos_emb`转换为一个常量张量，并保存在实例变量`positional_embedding`中。

在`call`函数中，首先获取输入张量`inputs`的形状，并使用`tf.shape`函数获取其动态形状。然后，根据输入张量的形状，从编码向量中选取对应的位置和维度，并将其加到输入张量上，返回结果。

最后，`get_config`函数用于获取该层的配置信息，其中包括数据类型`dtype`。