Update Mo.py

Mo.py

@@ -124,52 +124,6 @@ class SwiGLU(layers.Layer):
         x_val, x_gate = tf.split(x_proj, 2, axis=-1)
         return self.out(x_val * tf.nn.silu(x_gate))
         
-
-class MHLA(layers.Layer):
-    def __init__(self, embed_dim, num_heads=2, dropout=0.0):
-        super().__init__()
-        assert embed_dim % num_heads == 0, "embed_dim must be divisible by num_heads"
-        self.embed_dim = embed_dim
-        self.num_heads = num_heads
-        self.head_dim = embed_dim // num_heads
-        self.Wq = layers.Dense(embed_dim, use_bias=False)
-        self.Wk = layers.Dense(embed_dim, use_bias=False)
-        self.Wv = layers.Dense(embed_dim, use_bias=False)
-        self.out = layers.Dense(embed_dim)
-        self.dropout = layers.Dropout(dropout)
-
-    def split_heads(self, x):
-        # [B, L, D] -> [B, num_heads, L, head_dim]
-        B, L, D = tf.shape(x)[0], tf.shape(x)[1], tf.shape(x)[2]
-        x = tf.reshape(x, (B, L, self.num_heads, self.head_dim))
-        return tf.transpose(x, perm=[0, 2, 1, 3])
-
-    def combine_heads(self, x):
-        # [B, num_heads, L, head_dim] -> [B, L, D]
-        x = tf.transpose(x, perm=[0, 2, 1, 3])
-        B, L, H, D = tf.shape(x)[0], tf.shape(x)[1], tf.shape(x)[2], tf.shape(x)[3]
-        return tf.reshape(x, (B, L, H*D))
-
-    def call(self, x, training=False):
-        q = tf.nn.elu(self.Wq(x)) + 1
-        k = tf.nn.elu(self.Wk(x)) + 1
-        v = self.Wv(x)
-
-        q = self.split_heads(q)
-        k = self.split_heads(k)
-        v = self.split_heads(v)
-
-        # causal linear attention cumulative sum
-        k_cum = tf.cumsum(k, axis=2)
-        kv_cum = tf.cumsum(k * v, axis=2)
-
-        z = 1.0 / tf.reduce_sum(q * k_cum, axis=-1, keepdims=True)
-        out = (q * kv_cum) * z
-        out = self.combine_heads(out)
-        out = self.dropout(out, training=training)
-        return self.out(out)
-
-
 class Lo(layers.Layer):
     def __init__(self, d_model):
         super().__init__()
@@ -187,12 +141,13 @@ class Lo(layers.Layer):
 class Block(layers.Layer):
     def __init__(self, d_model):
         super().__init__()
-        self.lou = MHLA(d_model, 2)
+        self.mha = layers.MultiHeadAttention(8, 384//8)
         self.glu = SwiGLU(d_model, 1048)
         self.lo = Lo(d_model)
 
     def call(self, x):
-        x = self.lou(x)
+        x = self.mha(x)
+        x = self.glu(x)
         x = self.lo(x)
         return x