feat: migrate deepseek to nnx

Author: mesakhcienet

src/MaxText/layers/decoders.py

@@ -397,7 +397,7 @@ def get_decoder_layers(self):
         if self.config.use_batch_split_schedule:
           return [deepseek_batchsplit.DeepSeekDenseLayer, deepseek_batchsplit.DeepSeekMoELayer]
         else:
-          return [deepseek.DeepSeekDenseLayer, deepseek.DeepSeekMoELayer]
+          return [deepseek.DeepSeekDenseLayerToLinen, deepseek.DeepSeekMoELayerToLinen]
       case DecoderBlockType.GEMMA:
         return [gemma.GemmaDecoderLayerToLinen]
       case DecoderBlockType.GEMMA2: