Update V2.py

V2.py

@@ -128,18 +128,20 @@ ds = ds.batch(BATCH_SIZE, drop_remainder=True)
 ds = ds.map(lambda v1, v2: ((v1, v2), tf.zeros([BATCH_SIZE], dtype=tf.float32)), num_parallel_calls=tf.data.AUTOTUNE)
 ds = ds.prefetch(tf.data.AUTOTUNE)
 
-
-
 class MixerBlock(layers.Layer):
+    """
+    TPU / mixed-precision 친화적 MLP-Mixer 블록 (토큰-믹싱 + 채널-믹싱).
+    내부 연산은 float32로 수행하여 안정성 확보, 출력은 입력 dtype으로 복원.
+    """
     def __init__(self, seq_len, dim, token_mlp_dim, channel_mlp_dim, dropout=0.0):
         super().__init__()
         self.seq_len = seq_len
         self.dim = dim
-        self.token_mlp_dim = token_mlp_dim
-        self.channel_mlp_dim = channel_mlp_dim
 
+        # LayerNorm은 float32로 안정화
         self.ln1 = layers.LayerNormalization(epsilon=1e-6, dtype=tf.float32)
         # token-mixing MLP: operate over tokens => apply Dense on transposed axis
+        # Dense도 float32로 강제
         self.token_fc1 = layers.Dense(token_mlp_dim, activation='gelu', dtype=tf.float32)
         self.token_fc2 = layers.Dense(seq_len, dtype=tf.float32)
 
@@ -151,26 +153,31 @@ class MixerBlock(layers.Layer):
         self.dropout = layers.Dropout(dropout)
 
     def call(self, x, training=None):
-        # x: (B, L, D)
-        B = tf.shape(x)[0]
-        L = tf.shape(x)[1]
-        D = tf.shape(x)[2]
-
-        # Token-mixing
-        y = self.ln1(x)                          # (B, L, D)
-        y_t = tf.transpose(y, perm=[0,2,1])     # (B, D, L)
-        y_t = self.token_fc1(y_t)               # (B, D, token_mlp_dim)
-        y_t = self.token_fc2(y_t)               # (B, D, L)
-        y = tf.transpose(y_t, perm=[0,2,1])     # (B, L, D)
-        x = x + self.dropout(y, training=training)
-
-        # Channel-mixing
-        z = self.ln2(x)
-        z = self.channel_fc1(z)
-        z = self.channel_fc2(z)
-        x = x + self.dropout(z, training=training)
-
-        return x
+        """
+        x: (B, L, D) — dtype can be bfloat16/float32 depending on policy.
+        내부 계산은 float32로 수행하고, 반환은 원래 x.dtype으로 캐스팅.
+        """
+        orig_dtype = x.dtype
+        # 안정적 연산을 위해 float32로 변환
+        x_f = tf.cast(x, tf.float32)  # (B, L, D)
+
+        # ---- Token-mixing (Dense on token axis) ----
+        y = self.ln1(x_f)                      # (B, L, D) in float32
+        y_t = tf.transpose(y, perm=[0, 2, 1])  # (B, D, L)
+        y_t = self.token_fc1(y_t)              # (B, D, token_mlp_dim)
+        y_t = self.token_fc2(y_t)              # (B, D, L)
+        y = tf.transpose(y_t, perm=[0, 2, 1])  # (B, L, D)
+        x_f = x_f + self.dropout(y, training=training)
+
+        # ---- Channel-mixing (per-token MLP) ----
+        z = self.ln2(x_f)                      # (B, L, D)
+        z = self.channel_fc1(z)                # (B, L, channel_mlp_dim)
+        z = self.channel_fc2(z)                # (B, L, D)
+        x_f = x_f + self.dropout(z, training=training)
+
+        # 최종: 원래 dtype으로 복원 (mixed-precision 이득 유지)
+        return tf.cast(x_f, orig_dtype)
+
 
 class L2NormLayer(layers.Layer):
     def __init__(self, axis=1, epsilon=1e-10, **kwargs):