Merge pull request #2522 from CIeNET-International:feat/migrate-deepseek-split-batch-to-nnx

PiperOrigin-RevId: 829304686

Author: Google-ML-Automation

src/MaxText/layers/decoders.py

@@ -395,9 +395,15 @@ def get_decoder_layers(self):
         return [mixtral.MixtralDecoderLayerToLinen]
       case DecoderBlockType.DEEPSEEK:
         if self.config.use_batch_split_schedule:
-          return [deepseek_batchsplit.DeepSeekDenseLayer, deepseek_batchsplit.DeepSeekMoELayer]
+          return [
+              deepseek_batchsplit.DeepSeekDenseLayerToLinen,
+              deepseek_batchsplit.DeepSeekMoELayerToLinen,
+          ]
         else:
-          return [deepseek.DeepSeekDenseLayer, deepseek.DeepSeekMoELayer]
+          return [
+              deepseek.DeepSeekDenseLayerToLinen,
+              deepseek.DeepSeekMoELayerToLinen,
+          ]
       case DecoderBlockType.GEMMA:
         return [gemma.GemmaDecoderLayerToLinen]
       case DecoderBlockType.GEMMA2: