Update AlphaS2S.py

AlphaS2S.py

@@ -4,12 +4,21 @@ from tensorflow.keras import layers, Model
 class SwiGLU(layers.Layer):
     def __init__(self, d_model, d_ff):
         super().__init__()
-        self.proj = layers.Dense(d_ff*2)
+        self.proj = layers.Dense(d_ff)
         self.out = layers.Dense(d_model)
     def call(self, x):
         x_proj = self.proj(x)
         x_val, x_gate = tf.split(x_proj, 2, axis=-1)
         return self.out(x_val * tf.nn.silu(x_gate))
+        
+class CrossBlock(layers.Layer):
+    def __init__(self):
+        super().__init__()
+        self.alpha = layers.Dense(1, activation='sigmoid', dtype='float32')
+    def call(self, x, z):
+        a = self.alpha(x)
+        y = a * x + (1.0 - a) * z
+        return y
 
 class EncoderBlock(layers.Layer):
     def __init__(self, d_model, num_heads, dff, dropout=0.1):
@@ -35,12 +44,14 @@ class LoU(layers.Layer):
         self.Q = layers.Dense(d_model, dtype='float32')
         self.K = layers.Dense(d_model, dtype='float32')
         self.V = layers.Dense(d_model, dtype='float32')
-        self.proj = layers.Dense(d_model, use_bias=True, dtype='float32')
         self.norm = layers.LayerNormalization(epsilon=1e-5, dtype='float32')
         self.norm1 = layers.LayerNormalization(epsilon=1e-5, dtype='float32')
         
         self.alpha_linear = layers.Dense(1, activation='sigmoid', dtype='float32')
 
+        self.cross = CrossBlock()
+        self.glu = SwiGLU(d_model, 512)
+
     def _ema_over_time(self, score, alpha_dynamic):
         seq = tf.transpose(score, perm=[1, 0, 2])
         alpha_seq = tf.transpose(alpha_dynamic, perm=[1, 0, 2])
@@ -60,7 +71,7 @@ class LoU(layers.Layer):
         ema = tf.transpose(ema_seq, perm=[1, 0, 2])
         return ema
 
-    def call(self, x):
+    def call(self, x, z):
         x_f32 = tf.cast(x, tf.float32)
         residual = x_f32
         x_f32 = self.norm1(x)
@@ -83,8 +94,9 @@ class LoU(layers.Layer):
         score_norm = score_ema / denom
         score_clipped = tf.clip_by_value(score_norm, -self.clip_value, self.clip_value)
         x_comb = score_clipped * V
-        out = self.proj(x_comb)
-        out = self.norm(out + residual)
+        out = self.norm(x_comb + residual)
+        out = self.cross(out, z)
+        out = self.glu(out)
         return tf.cast(out, x.dtype)
         
 class Transformer(tf.keras.Model):