Fix linter

Signed-off-by: Vladimir Suvorov <suvorovv@google.com>

Author: mydatascience

src/MaxText/layerwise_quantization.py

@@ -38,7 +38,6 @@
 import jax.numpy as jnp
 from absl import app
 
-from flax import nnx
 from flax.linen import partitioning as nn_partitioning
 
 
@@ -94,29 +93,15 @@ def load_and_quantize(self, rng: None | PRNGKeyType = None) -> None:
 
     self.quant.quant_mode = quantizations.get_quant_mode("convert")
 
-    if rng is None:
-      rng = jax.random.PRNGKey(0)
-
-    dense_init_rng, moe_init_rng = jax.random.split(rng)
-    dense_params_rng, dense_dropout_rng = jax.random.split(dense_init_rng)
-    moe_params_rng, moe_dropout_rng = jax.random.split(moe_init_rng)
-
-    dense_layer_rngs = nnx.Rngs(params=dense_params_rng, dropout=dense_dropout_rng)
-    moe_layer_rngs = nnx.Rngs(params=moe_params_rng, dropout=moe_dropout_rng)
-
     layers = [
-        deepseek.DeepSeekDenseLayer(
+        deepseek.DeepSeekDenseLayer(  # pylint: disable=no-value-for-parameter
             config,
-            model_mode=common_types.MODEL_MODE_TRAIN,
             mesh=self._mesh,
-            rngs=dense_layer_rngs,
             quant=self.quant,
         ),
-        deepseek.DeepSeekMoELayer(
+        deepseek.DeepSeekMoELayer(  # pylint: disable=no-value-for-parameter
             config,
-            model_mode=common_types.MODEL_MODE_TRAIN,
             mesh=self._mesh,
-            rngs=moe_layer_rngs,
             quant=self.quant,
         ),
     ]