Migrate Gpt3 to NNX.

Author: hsuan-lun-chiang

src/MaxText/layers/decoders.py

@@ -405,7 +405,7 @@ def get_decoder_layers(self):
       case DecoderBlockType.GEMMA3:
         return [gemma3.Gemma3DecoderLayerToLinen]
       case DecoderBlockType.GPT3:
-        return [gpt3.Gpt3DecoderLayer]
+        return [gpt3.Gpt3DecoderLayerToLinen]
       case DecoderBlockType.GPT_OSS:
         return [gpt_oss.GptOssScannableBlockToLinen] if self.config.scan_layers else [gpt_oss.GptOssDecoderLayerToLinen]
       case DecoderBlockType.QWEN3:
@@ -584,7 +584,7 @@ def _apply_embedding(
           name="position_embedder",
           config=cfg,
           mesh=self.mesh,
-      )(decoder_positions, model_mode=model_mode)
+      )(decoder_positions.astype("int32"), model_mode=model_mode)
     return y
 
   @nn.compact
@@ -837,9 +837,7 @@ def __call__(
           # Iterate over the two layer groups (dense and MoE) and apply layer transformation
           for layer, num_layers, layer_prefix in zip(layers, num_layers_list, layer_prefixes):
             for index in range(num_layers):
-              y = layer(
-                  config=cfg, mesh=mesh, name=f"{layer_prefix}_{index}", quant=self.quant, model_mode=self.model_mode
-              )(
+              y = layer(config=cfg, mesh=mesh, name=f"{layer_prefix}_{index}", quant=self.quant, model_mode=self.model_mode)(
                   y,
                   decoder_segment_ids,
                   decoder_positions,