Reverts 69ed0c5

PiperOrigin-RevId: 830610186

Author: Google-ML-Automation

src/MaxText/layers/decoders.py

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