Update Mo.py

Mo.py

@@ -124,6 +124,17 @@ class SwiGLU(layers.Layer):
         out = self.W1(tf.nn.silu(a) * b)     
         return tf.cast(out, x.dtype)
 
+class SwiGLU(layers.Layer):
+    def __init__(self, d_model, d_ff):
+        super().__init__()
+        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 LoU(layers.Layer):
     def __init__(self, d_model, clip_value=5.0, eps=1e-6):
         super().__init__()
@@ -137,9 +148,7 @@ class LoU(layers.Layer):
         self.norm1 = layers.LayerNormalization(epsilon=1e-5, dtype='float32')
         
         self.glu = SwiGLU(d_model, 320)
-        self.cross = CrossBlock()
-
-    def call(self, x, z):
+    def call(self, x):
         x_f32 = tf.cast(x, tf.float32)
         residual = x_f32
         x_f32 = self.norm1(x)
@@ -171,7 +180,6 @@ class LoU(layers.Layer):
         x_comb = score_clipped * V
         
         out = self.norm(x_comb + residual)
-        out = self.cross(out, z)
         out = self.glu(out)
         return tf.cast(out, x.dtype)