feat: reorder import

Author: mesakhcienet

src/MaxText/layers/deepseek.py

@@ -25,14 +25,19 @@
 from flax import linen as nn
 from flax import nnx
 
-from MaxText.layers import initializers, linears, moe, nnx_wrappers, quantizations
-from MaxText.layers import attention_mla
+from MaxText import max_utils
 from MaxText.common_types import Config
-from MaxText.layers.normalizations import RMSNorm
-from MaxText.layers.linears import Dropout
-from MaxText.inference import page_manager
 from MaxText.common_types import MODEL_MODE_PREFILL
-from MaxText import max_utils
+from MaxText.inference import page_manager
+from MaxText.layers import attention_mla
+from MaxText.layers import initializers
+from MaxText.layers import linears
+from MaxText.layers import moe
+from MaxText.layers import nnx_wrappers
+from MaxText.layers import quantizations
+from MaxText.layers.linears import Dropout
+from MaxText.layers.normalizations import RMSNorm
+
 # -----------------------------------------
 # The Decoder Layer for DeepSeek v3
 # -----------------------------------------