qtransformer_mha_classical.py

class MultiHeadAttentionClassical(MultiHeadAttentionBase):
    def __init__(self, embed_dim, num_heads):
        super(MultiHeadAttentionClassical, self).__init__(embed_dim, num_heads)
        self.wq = tf.keras.layers.Dense(embed_dim)
        self.wk = tf.keras.layers.Dense(embed_dim)
        self.wv = tf.keras.layers.Dense(embed_dim)
        self.dense = tf.keras.layers.Dense(embed_dim)
    
    def apply_dense_layers(self, v, k, q):
        q = self.wq(q)  # (batch_size, seq_len, embed_dim)
        k = self.wk(k)  # (batch_size, seq_len, embed_dim)
        v = self.wv(v)  # (batch_size, seq_len, embed_dim)
        return v, k, q