fix: run lint

Author: mesakhcienet

src/MaxText/layers/deepseek.py

@@ -65,9 +65,7 @@ def __init__(
     self.quant = quant
     self.rngs = rngs
 
-    batch_size, sequence_length = max_utils.get_batch_seq_len_for_mode(
-        self.config, self.model_mode
-    )
+    batch_size, sequence_length = max_utils.get_batch_seq_len_for_mode(self.config, self.model_mode)
     self.dummy_inputs_shape = (batch_size, sequence_length, self.config.emb_dim)
 
     self.pre_self_attention_layer_norm = RMSNorm(
@@ -119,9 +117,7 @@ def __init__(
         rngs=rngs,
     )
 
-    self.dropout = Dropout(
-        rate=self.config.dropout_rate, broadcast_dims=(-2,), rngs=self.rngs
-    )
+    self.dropout = Dropout(rate=self.config.dropout_rate, broadcast_dims=(-2,), rngs=self.rngs)
 
   def __call__(
       self,
@@ -162,9 +158,7 @@ def with_logical_constraint(self, x):
     return nn.with_logical_constraint(x, self.logical_axis_names)
 
   def dropout_op(self, x, deterministic):
-    return self.with_logical_constraint(
-        self.dropout(x, deterministic=deterministic)
-    )
+    return self.with_logical_constraint(self.dropout(x, deterministic=deterministic))
 
   def pre_attention_norm_op(self, x):
     return self.with_logical_constraint(self.pre_self_attention_layer_norm(x))
@@ -311,9 +305,7 @@ def __init__(
     self.DeepSeekMoeBlock_0 = moe.RoutedAndSharedMoE(
         config=self.config,
         mesh=mesh,
-        kernel_init=initializers.nd_dense_init(
-            1.0, "fan_in", "truncated_normal"
-        ),
+        kernel_init=initializers.nd_dense_init(1.0, "fan_in", "truncated_normal"),
         kernel_axes=("embed", None),
         dtype=self.config.dtype,
         weight_dtype=self.config.weight_dtype,

tests/pipeline_parallelism_test.py

@@ -69,13 +69,9 @@ def assert_pipeline_same_output_and_grad(self, config, single_pipeline_stage_cla
     else:
       if issubclass(single_pipeline_stage_class, nnx_wrappers.ToLinen):
         rngs = nnx.Rngs(params=0)
-        single_pipeline_stage = single_pipeline_stage_class(
-            config=config, mesh=mesh, model_mode=model_mode, rngs=rngs
-        )
+        single_pipeline_stage = single_pipeline_stage_class(config=config, mesh=mesh, model_mode=model_mode, rngs=rngs)
       else:
-        single_pipeline_stage = single_pipeline_stage_class(
-            config=config, mesh=mesh, model_mode=model_mode
-        )
+        single_pipeline_stage = single_pipeline_stage_class(config=config, mesh=mesh, model_mode=model_mode)
 
     def get_inputs(batch_size, sequence, features):
       """Get random inputs, and random dummy targets