Embedded Gaussian中的两个嵌入权重变换 的tensorflow2代码

以下是使用TensorFlow 2实现Embedded Gaussian中的两个嵌入权重变换的代码示例：

import tensorflow as tf

class EmbeddingWeightTransform(tf.keras.layers.Layer):
    def __init__(self, d_model, num_heads, name="embedding_weight_transform"):
        super().__init__(name=name)
        self.d_model = d_model
        self.num_heads = num_heads
        
        # define the layers for the embedding weight transform
        self.dense1 = tf.keras.layers.Dense(d_model, name="dense1")
        self.reshape = tf.keras.layers.Reshape((num_heads, d_model // num_heads), name="reshape")
        self.transpose = tf.keras.layers.Permute((2, 1), name="transpose")
        
    def call(self, inputs):
        # apply the embedding weight transform
        x = self.dense1(inputs)
        x = self.reshape(x)
        x = self.transpose(x)
        return x

其中，`EmbeddingWeightTransform`是一个自定义的Keras层，它接收一个形状为`(batch_size, seq_len, d_model)`的输入张量，并返回一个形状为`(batch_size, num_heads, d_model // num_heads, seq_len)`的输出张量。这个层内部实现了Embedded Gaussian中的两个嵌入权重变换：

1. 将输入张量通过一个全连接层`(batch_size, seq_len, d_model) -> (batch_size, seq_len, d_model)`，得到一个形状为`(batch_size, seq_len, d_model)`的张量。
2. 将全连接层的输出张量通过一个`Reshape`层`(batch_size, seq_len, d_model) -> (batch_size, num_heads, d_model // num_heads, seq_len)`，得到一个形状为`(batch_size, num_heads, d_model // num_heads, seq_len)`的张量。
3. 将`Reshape`层的输出张量通过一个`Permute`层`(batch_size, num_heads, d_model // num_heads, seq_len) -> (batch_size, d_model // num_heads, num_heads, seq_len)`，得到一个形状为`(batch_size, d_model // num_heads, num_heads, seq_len)`的张量。

下面是另一个嵌入权重变换的代码示例：

class EmbeddingWeightTransform(tf.keras.layers.Layer):
    def __init__(self, d_model, num_heads, name="embedding_weight_transform"):
        super().__init__(name=name)
        self.d_model = d_model
        self.num_heads = num_heads
        
        # define the layers for the embedding weight transform
        self.dense1 = tf.keras.layers.Dense(d_model, name="dense1")
        self.reshape1 = tf.keras.layers.Reshape((num_heads, -1), name="reshape1")
        self.permute = tf.keras.layers.Permute((2, 1), name="permute")
        self.dense2 = tf.keras.layers.Dense(d_model, name="dense2")
        self.reshape2 = tf.keras.layers.Reshape((-1, num_heads, d_model // num_heads), name="reshape2")
        
    def call(self, inputs):
        # apply the embedding weight transform
        x = self.dense1(inputs)
        x = self.reshape1(x)
        x = self.permute(x)
        x = self.dense2(x)
        x = self.reshape2(x)
        return x

与前面的实现相比，这个实现在`Reshape`层中使用了占位符`-1`，以便根据其他维度自动推断出该维度大小。此外，这个实现还添加了一个额外的全连接层和一个`Reshape`层。这两个层的作用是将输入张量从形状`(batch_size, seq_len, d_model)`转换为形状`(batch_size, num_heads, d_model // num_heads, seq_len)`。