Copybara import of the project:

--
b6dc761 by mesakhcienet <mesakh.christian@cienet.com>:

feat: migrate deepseek to nnx
COPYBARA_INTEGRATE_REVIEW=#2658 from CIeNET-International:feat/migrate-deepseek-to-nnx b6dc761
PiperOrigin-RevId: 852094911

Author: mesakhcienet

src/MaxText/layers/decoders.py

@@ -404,9 +404,9 @@ 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: