Run pre-commit to fix MaxText linter

src/maxtext/layers/moe.py

@@ -252,9 +252,7 @@ def calculate_load_balance_updates(top_k_indices, num_experts, rate):
   flat_indices = top_k_indices.ravel()
   # one_hot rather than bincount: bincount clips out-of-range values, so the -1
   # padding that forced routing uses would all be counted as expert 0.
-  expert_counts = jnp.sum(
-      jax.nn.one_hot(flat_indices, num_experts, dtype=jnp.int32), axis=0
-  )
+  expert_counts = jnp.sum(jax.nn.one_hot(flat_indices, num_experts, dtype=jnp.int32), axis=0)
   total_tokens = jnp.sum(expert_counts)
   average_load = total_tokens / num_experts
   direction = jnp.sign(average_load - expert_counts)
@@ -771,59 +769,37 @@ def get_topk(
       gather_indices = jnp.where(valid_token_mask, top_k_indices, 0)
       if self.config.decoder_block == ctypes.DecoderBlockType.GEMMA4:
         router_probs = jax.nn.softmax(gate_logits.astype(jnp.float32), axis=-1)
-        top_k_weights = jnp.take_along_axis(
-            router_probs, gather_indices, axis=-1
-        ).astype(self.dtype)
+        top_k_weights = jnp.take_along_axis(router_probs, gather_indices, axis=-1).astype(self.dtype)
       else:
-        top_k_weights = jnp.take_along_axis(
-            gate_logits, gather_indices, axis=-1
-        )
+        top_k_weights = jnp.take_along_axis(gate_logits, gather_indices, axis=-1)
     else:
       if self.config.use_random_routing:
         if rngs is None:
           raise ValueError("The random key cannot be None for random routing.")
         # Reuse the 'params' RNG stream to ensure random routing
-        rng = (
-            rngs.params()
-            if hasattr(rngs, "params") and callable(getattr(rngs, "params"))
-            else rngs
-        )
-        top_k_weights, top_k_indices = random_routing(
-            rng, gate_logits, self.num_experts_per_tok
-        )
+        rng = rngs.params() if hasattr(rngs, "params") and callable(getattr(rngs, "params")) else rngs
+        top_k_weights, top_k_indices = random_routing(rng, gate_logits, self.num_experts_per_tok)
         return top_k_weights, top_k_indices
 
       if self.is_hash_routing:
         if input_ids is None:
-          raise ValueError(
-              "input_ids cannot be None when is_hash_routing is True"
-          )
+          raise ValueError("input_ids cannot be None when is_hash_routing is True")
         # Access the static routing table
         tid2eid_int = self.tid2eid.value
         # Cast the float32 array to int32 (JAX automatically assigns 0.0 gradients to integer casts)
         tid2eid_int = tid2eid_int.astype(jnp.int32)
         # Cast input_ids to int32 to safely index the hash routing table
         top_k_indices = tid2eid_int[input_ids.astype(jnp.int32)]
-        top_k_weights = jnp.take_along_axis(
-            pre_bias_logits, top_k_indices, axis=-1
-        )
+        top_k_weights = jnp.take_along_axis(pre_bias_logits, top_k_indices, axis=-1)
       # NOTE: deepseek2 has a different pattern
-      elif self.config.model_name.startswith(
-          ("deepseek3", "deepseek4", "kimi-k2")
-      ):
-        top_k_weights, top_k_indices = self.deepseek_routing(
-            gate_logits, pre_bias_logits
-        )
+      elif self.config.model_name.startswith(("deepseek3", "deepseek4", "kimi-k2")):
+        top_k_weights, top_k_indices = self.deepseek_routing(gate_logits, pre_bias_logits)
       elif self.config.decoder_block == ctypes.DecoderBlockType.GEMMA4:
         router_probs = jax.nn.softmax(gate_logits.astype(jnp.float32), axis=-1)
         _, top_k_indices = jax.lax.top_k(gate_logits, self.num_experts_per_tok)
-        top_k_weights = jnp.take_along_axis(
-            router_probs, top_k_indices, axis=-1
-        ).astype(self.dtype)
+        top_k_weights = jnp.take_along_axis(router_probs, top_k_indices, axis=-1).astype(self.dtype)
       else:
-        top_k_weights, top_k_indices = jax.lax.top_k(
-            gate_logits, self.num_experts_per_tok
-        )
+        top_k_weights, top_k_indices = jax.lax.top_k(gate_logits, self.num_experts_per_tok)
 
     if self.config.decoder_block in (ctypes.DecoderBlockType.DEEPSEEK, ctypes.DecoderBlockType.DEEPSEEK4):
       top_k_weights = self.deepseek_scale_weights(top_k_weights)
@@ -834,9 +810,7 @@ def get_topk(
         if valid_token_mask is not None:
           # Padding must stay out of the softmax denominator or it rescales the
           # real slots. Large-negative, not -inf: a fully-padded token would be NaN.
-          top_k_weights = jnp.where(
-              valid_token_mask, top_k_weights, jnp.finfo(jnp.float32).min / 2
-          )
+          top_k_weights = jnp.where(valid_token_mask, top_k_weights, jnp.finfo(jnp.float32).min / 2)
         top_k_weights = jax.nn.softmax(top_k_weights.astype(jnp.float32), axis=-1).astype(self.dtype)
 
       if valid_token_mask is not None:
@@ -979,9 +953,7 @@ def permute(
     inputs_shape = inputs.shape
     bsz_times_seq_len = inputs_shape[0] * inputs_shape[1]
     inputs_2d = jnp.reshape(inputs, (bsz_times_seq_len, inputs_shape[2]))
-    weights, selected_experts = self.get_topk(
-        gate_logits, pre_bias_logits, rngs, input_ids, forced_routed_experts
-    )
+    weights, selected_experts = self.get_topk(gate_logits, pre_bias_logits, rngs, input_ids, forced_routed_experts)
 
     lb_loss = None
     # Using pre_bias_logits ensures the router bias does not leak into the auxiliary loss gradient
@@ -1061,38 +1033,28 @@ def permute(
       # Must precede roll_to_expert_id: `(-1 - roll) % num_experts` wraps padding
       # onto a real expert id, so a mask computed after it sees no padding at all.
       if forced_routed_experts is not None:
-        valid_mask = valid_expert_mask(
-            flatten_selected_experts, self.num_experts
-        )
+        valid_mask = valid_expert_mask(flatten_selected_experts, self.num_experts)
 
       if roll_to_expert_id is not None:
         flatten_selected_experts = (flatten_selected_experts - roll_to_expert_id) % self.num_experts
 
       if forced_routed_experts is not None:
         # Spread padding round-robin: it carries zero weight but still counts in
         # group_size, so under ragged_buffer_factor > 0 it can evict real tokens.
-        dummy_indices = (
-            jnp.arange(flatten_selected_experts.shape[0]) % self.num_experts
-        )
-        flatten_selected_experts_safe = jnp.where(
-            valid_mask, flatten_selected_experts, dummy_indices
-        )
+        dummy_indices = jnp.arange(flatten_selected_experts.shape[0]) % self.num_experts
+        flatten_selected_experts_safe = jnp.where(valid_mask, flatten_selected_experts, dummy_indices)
       else:
         flatten_selected_experts_safe = flatten_selected_experts
 
       sorted_selected_experts = jnp.argsort(flatten_selected_experts_safe)
       if self.config.moe_use_direct_token_gather:
-        sorted_inputs = _route_activations(
-            inputs_2d, flatten_selected_experts_safe
-        ).astype(self.dtype)
+        sorted_inputs = _route_activations(inputs_2d, flatten_selected_experts_safe).astype(self.dtype)
       else:
         replicated_inputs_2d = jnp.repeat(inputs_2d, self.num_experts_per_tok, axis=0)
         sorted_inputs = _sort_activations(replicated_inputs_2d, sorted_selected_experts, use_custom_sort_vjp).astype(
             self.dtype
         )
-      group_size = jnp.bincount(
-          flatten_selected_experts_safe, length=self.num_experts
-      )
+      group_size = jnp.bincount(flatten_selected_experts_safe, length=self.num_experts)
 
     num_tokens = bsz_times_seq_len * self.num_experts_per_tok
     use_truncated_buffer = use_ragged_in_permute and buffer_size is not None and buffer_size < num_tokens
@@ -2395,21 +2357,19 @@ def _moe_body(
     ):
       batch_size, sequence_length, embed_dim = x.shape
       if self.config.num_moe_emb_chunks > 0:
-        output0, output1, gmm_fn, routing, route_metadata, wo_bias = (
-            moe_emb_chunking(
-                x,
-                logits,
-                pre_bias_logits,
-                w0,
-                w1,
-                w0_bias,
-                w1_bias,
-                wo_bias,
-                sharded_input_ids,
-                rngs,
-                embed_dim,
-                forced_routed_experts=forced_routed_experts,
-            )
+        output0, output1, gmm_fn, routing, route_metadata, wo_bias = moe_emb_chunking(
+            x,
+            logits,
+            pre_bias_logits,
+            w0,
+            w1,
+            w0_bias,
+            w1_bias,
+            wo_bias,
+            sharded_input_ids,
+            rngs,
+            embed_dim,
+            forced_routed_experts=forced_routed_experts,
         )
       else:
         x, routing, route_metadata = route(
@@ -2604,9 +2564,7 @@ def sparse_matmul_route_and_compute(
             wo_bias,
             None if sharded_input_ids is None else sharded_input_ids[:, sl],
             rngs,
-            None
-            if forced_routed_experts is None
-            else forced_routed_experts[:, sl, :],
+            None if forced_routed_experts is None else forced_routed_experts[:, sl, :],
         )
         if self.config.moe_chunk_barrier:
           _prev = out_c
@@ -2709,9 +2667,7 @@ def reshape_and_update_weights(self, weights, indices, safe_updates=False):
             jnp.arange(weights.shape[0])[:, None, None],
             ("activation_batch", None, None),
         ),
-        self._maybe_shard_with_logical(
-            jnp.arange(weights.shape[1])[:, None], ("activation_length", None)
-        ),
+        self._maybe_shard_with_logical(jnp.arange(weights.shape[1])[:, None], ("activation_length", None)),
         safe_indices,
     )
     weight_sharding = (
@@ -2726,13 +2682,9 @@ def reshape_and_update_weights(self, weights, indices, safe_updates=False):
       # so accumulate with `.add()` instead: padding slots always carry a
       # zero weight, so summing duplicates is safe and avoids silently
       # dropping a real weight update.
-      update_weights = update_weights.at[index_update].add(
-          safe_weights, out_sharding=weight_sharding
-      )
+      update_weights = update_weights.at[index_update].add(safe_weights, out_sharding=weight_sharding)
     else:
-      update_weights = update_weights.at[index_update].set(
-          safe_weights, out_sharding=weight_sharding
-      )
+      update_weights = update_weights.at[index_update].set(safe_weights, out_sharding=weight_sharding)
     return update_weights
 
   def get_context_partition_and_sub_seq(self, seq_len):
@@ -3333,10 +3285,7 @@ def fused_moe_matmul(
     It does not compute lb_loss or bias_updates (inference-only).
     """
     if forced_routed_experts is not None:
-      raise NotImplementedError(
-          "Forced routing via forced_routed_experts is not supported with"
-          " fused_moe_matmul."
-      )
+      raise NotImplementedError("Forced routing via forced_routed_experts is not supported with fused_moe_matmul.")
     try:
       # pylint: disable=import-outside-toplevel
       # pytype: disable=import-error

src/maxtext/layers/nnx_decoders.py

@@ -1048,9 +1048,7 @@ def layer_fn(carry, scanned_vars):
         current_params, current_state, kv_cache_layer = scanned_vars
         forced_routed_experts_layer = None
       elif use_forced_routing:
-        current_params, current_state, forced_routed_experts_layer = (
-            scanned_vars
-        )
+        current_params, current_state, forced_routed_experts_layer = scanned_vars
         kv_cache_layer = None
       else:
         current_params, current_state = scanned_vars
@@ -1138,9 +1136,7 @@ def layer_fn(carry, scanned_vars):
         scan_xs = (params, state, forced_routed_experts_scanned)
       else:
         scan_xs = (params, state)
-      final_carry, scanned_state = jax.lax.scan(
-          layer_fn_wrapped, x_in, scan_xs, unroll=unroll
-      )
+      final_carry, scanned_state = jax.lax.scan(layer_fn_wrapped, x_in, scan_xs, unroll=unroll)
       returned_kv_stacked = None
 
       # Move the scan axis to each variable's param_scan_axis and restore its name
@@ -1980,13 +1976,11 @@ def __call__(
               )
             # Only the per-layer slices may reach the layers from here on.
             layer_kwargs.pop("forced_routed_experts", None)
-            forced_routed_experts_scanned = (
-                reshape_forced_routed_experts_for_scan(
-                    forced_routed_experts,
-                    num_layers=cfg.num_decoder_layers,
-                    scan_length=scan_length,
-                    layers_per_cycle=cycle_interval,
-                )
+            forced_routed_experts_scanned = reshape_forced_routed_experts_for_scan(
+                forced_routed_experts,
+                num_layers=cfg.num_decoder_layers,
+                scan_length=scan_length,
+                layers_per_cycle=cycle_interval,
             )
             if cfg.decoder_block == DecoderBlockType.MIXTRAL:
               # Mixtral has no ScannableBlock: one scan iteration is one layer,
@@ -1997,9 +1991,7 @@ def __call__(
                     " inhomogeneous_layer_cycle_interval == 1; got"
                     f" {cycle_interval}."
                 )
-              forced_routed_experts_scanned = jnp.squeeze(
-                  forced_routed_experts_scanned, axis=1
-              )
+              forced_routed_experts_scanned = jnp.squeeze(forced_routed_experts_scanned, axis=1)
           if kv_caches is not None:
             # Pass the kv_caches list directly to avoid copying in jnp.stack,
             # which breaks vLLM PagedAttention in-place memory updates.
@@ -2033,9 +2025,7 @@ def pure_layer_fn(graphdef_in, state_in, y_in, kv_in, valid_kwargs):
                 state_in,
             )
           merged_layer = nnx.merge(graphdef_in, state_in)
-          out_y, out_kv = merged_layer(
-              y_in, *layer_args, kv_cache=kv_in, **valid_kwargs
-          )
+          out_y, out_kv = merged_layer(y_in, *layer_args, kv_cache=kv_in, **valid_kwargs)
           state_out = nnx.state(merged_layer)
 
           if dynamic_graph_init:
@@ -2105,10 +2095,7 @@ def pure_layer_fn(graphdef_in, state_in, y_in, kv_in, valid_kwargs):
                   f" top_k] (4D, per-layer); got ndim={routed_experts.ndim}"
                   f" with shape {routed_experts.shape}."
               )
-            if (
-                routed_experts.ndim == 4
-                and routed_experts.shape[2] != cfg.num_decoder_layers
-            ):
+            if routed_experts.ndim == 4 and routed_experts.shape[2] != cfg.num_decoder_layers:
               # jnp clamps an out-of-range static index, so a short layer axis
               # would silently replay the last slice on every later layer.
               raise ValueError(
@@ -2118,19 +2105,13 @@ def pure_layer_fn(graphdef_in, state_in, y_in, kv_in, valid_kwargs):
                   f" {routed_experts.shape}."
               )
             current_kwargs["forced_routed_experts"] = (
-                routed_experts[:, :, lyr, :]
-                if routed_experts.ndim == 4
-                else routed_experts
+                routed_experts[:, :, lyr, :] if routed_experts.ndim == 4 else routed_experts
             )
 
           if cfg.remat_policy != "none":
-            y, kv_cache, new_state, new_graphdef = checkpointed_fn(
-                graphdef, state, y, kv_cache, current_kwargs
-            )
+            y, kv_cache, new_state, new_graphdef = checkpointed_fn(graphdef, state, y, kv_cache, current_kwargs)
           else:
-            y, kv_cache, new_state, new_graphdef = pure_layer_fn(
-                graphdef, state, y, kv_cache, current_kwargs
-            )
+            y, kv_cache, new_state, new_graphdef = pure_layer_fn(graphdef, state, y, kv_cache, current_kwargs)
 
           if dynamic_graph_init:
             new_layer = nnx.merge(new_graphdef, new_state)

src/maxtext/models/gemma4.py

@@ -140,13 +140,11 @@ def __call__(
     gate_inputs = unscaled_norm * root_size * router_scale
 
     # 3. Pass both to routed_moe
-    routed_experts, load_balance_loss, moe_bias_updates = (
-        self.moe_block.routed_moe(
-            routed_inputs,
-            gate_inputs=gate_inputs,
-            out_sharding=out_sharding,
-            forced_routed_experts=forced_routed_experts,
-        )
+    routed_experts, load_balance_loss, moe_bias_updates = self.moe_block.routed_moe(
+        routed_inputs,
+        gate_inputs=gate_inputs,
+        out_sharding=out_sharding,
+        forced_routed_experts=forced_routed_experts,
     )
     routed_experts = self.post_feedforward_layernorm_2(routed_experts)
 

src/maxtext/models/mixtral.py

@@ -169,9 +169,7 @@ def __call__(
     # NOTE: the naming mismatch here is to ensure reverse compatibility with existing checkpoints.
     # The `name` represents the weight name in JAX/checkpoints and so the class name
     # is just for readability.
-    mlp_lnx, load_balance_loss, _ = self.MoeBlock_0(
-        hidden_states, forced_routed_experts=forced_routed_experts
-    )
+    mlp_lnx, load_balance_loss, _ = self.MoeBlock_0(hidden_states, forced_routed_experts=forced_routed_experts)
     mlp_lnx = nn.with_logical_constraint(mlp_lnx, self.activation_axis_names)
 
     layer_output = mlp_lnx + intermediate_inputs

src/maxtext/models/qwen3.py

@@ -1193,9 +1193,7 @@ def __call__(
         - The load balancing loss from the routed experts, if applicable during training.
     """
     # 1. Apply the routed experts block.
-    routed_output, load_balance_loss, _ = self.routed_experts(
-        hidden_states, forced_routed_experts=forced_routed_experts
-    )
+    routed_output, load_balance_loss, _ = self.routed_experts(hidden_states, forced_routed_experts=forced_routed_experts)
 
     if not self.use_shared_expert:
       return routed_output, load_balance_loss

src/maxtext/models/qwen3_5.py

@@ -99,11 +99,7 @@ def __call__(
       # forced_routed_experts, when present, is shaped
       # [inhomogeneous_layer_cycle_interval, batch, seq, top_k]: one slice per
       # sub-layer in this cycle (see nnx_decoders.py's scan wiring).
-      layer_forced_routed_experts = (
-          forced_routed_experts[i]
-          if forced_routed_experts is not None
-          else None
-      )
+      layer_forced_routed_experts = forced_routed_experts[i] if forced_routed_experts is not None else None
       x, _ = layer(
           x,
           decoder_segment_ids,

src/maxtext/trainers/pre_train/train.py

@@ -148,9 +148,7 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr
   # Only forward the kwarg when router replay is actually in use, so models
   # and adapters whose __call__ predates the feature keep working.
   forced_routing_kwargs = (
-      {"forced_routed_experts": data["forced_routed_experts"]}
-      if "forced_routed_experts" in data
-      else {}
+      {"forced_routed_experts": data["forced_routed_experts"]} if "forced_routed_experts" in data else {}
   )
 
   if is_block_diffusion: