Merge pull request #4446 from AI-Hypercomputer:chengnuojin-fix-lint

PiperOrigin-RevId: 948475711

Author: Google-ML-Automation

src/maxtext/checkpoint_conversion/utils/param_mapping.py

@@ -1269,7 +1269,9 @@ def concat_ba_and_transpose(input_tensor, target_shape=None):
     hooks[f"{mlp_prefix}-shared_expert-wo-kernel"] = transpose
     hooks[f"{mlp_prefix}-shared_expert_gate-kernel"] = transpose
 
-    hooks[(f"{mlp_prefix}-routed_experts-wi_0", f"{mlp_prefix}-routed_experts-wi_1")] = process_wi_0_wi_1  # pyrefly: ignore[unsupported-operation]
+    hooks[(f"{mlp_prefix}-routed_experts-wi_0", f"{mlp_prefix}-routed_experts-wi_1")] = (
+        process_wi_0_wi_1  # pyrefly: ignore[unsupported-operation]
+    )
     hooks[f"{mlp_prefix}-routed_experts-wo"] = transpose_expert
 
   # Vision hooks for Qwen3.5
@@ -1340,7 +1342,9 @@ def reshape_vision_attn_out(input_tensor, target_shape):
         return input_tensor.T.reshape(target_shape)
 
     # Apply vision hooks
-    hooks["params-vision_encoder-Qwen3_5MoeVisionEncoder_0-patch_embed-proj-kernel"] = reshape_conv3d_patch_embed  # pyrefly: ignore[bad-assignment]
+    hooks["params-vision_encoder-Qwen3_5MoeVisionEncoder_0-patch_embed-proj-kernel"] = (
+        reshape_conv3d_patch_embed  # pyrefly: ignore[bad-assignment]
+    )
 
     for i in range(n_vision_layers):
       prefix = f"params-vision_encoder-Qwen3_5MoeVisionEncoder_0-blocks_{i}"
@@ -1355,8 +1359,12 @@ def reshape_vision_attn_out(input_tensor, target_shape):
       hooks[f"{prefix}-mlp_out-kernel"] = reshape_kernel_vision  # pyrefly: ignore[bad-assignment]
 
     # Vision projector
-    hooks["params-vision_encoder-Qwen3_5MoeVisionProjector_0-merger-mlp_0-kernel"] = reshape_kernel_vision  # pyrefly: ignore[bad-assignment]
-    hooks["params-vision_encoder-Qwen3_5MoeVisionProjector_0-merger-mlp_2-kernel"] = reshape_kernel_vision  # pyrefly: ignore[bad-assignment]
+    hooks["params-vision_encoder-Qwen3_5MoeVisionProjector_0-merger-mlp_0-kernel"] = (
+        reshape_kernel_vision  # pyrefly: ignore[bad-assignment]
+    )
+    hooks["params-vision_encoder-Qwen3_5MoeVisionProjector_0-merger-mlp_2-kernel"] = (
+        reshape_kernel_vision  # pyrefly: ignore[bad-assignment]
+    )
 
   return hooks
 
@@ -1384,7 +1392,9 @@ def QWEN3_NEXT_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=F
       prefix = f"params-decoder-layers-layer_{block_idx}"
 
       # Layer norms
-      mapping[f"{prefix}-input_layernorm-scale"] = [f"model.layers.{i}.input_layernorm.weight" for i in hf_indices]  # pyrefly: ignore[bad-assignment]
+      mapping[f"{prefix}-input_layernorm-scale"] = [
+          f"model.layers.{i}.input_layernorm.weight" for i in hf_indices
+      ]  # pyrefly: ignore[bad-assignment]
       mapping[f"{prefix}-post_attention_layernorm-scale"] = [  # pyrefly: ignore[bad-assignment]
           f"model.layers.{i}.post_attention_layernorm.weight" for i in hf_indices
       ]
@@ -1943,7 +1953,9 @@ def interleave(input_tensor, target_shape=None):
     hooks[f"{prefix}-GptOssMlp-gate-kernel"] = transpose
     # `composite_mt_key`: A hook for combining multiple MaxText params.
     hooks[(f"{prefix}-GptOssMlp-wi_0", f"{prefix}-GptOssMlp-wi_1")] = interleave  # pyrefly: ignore[unsupported-operation]
-    hooks[(f"{prefix}-GptOssMlp-wi_0_bias", f"{prefix}-GptOssMlp-wi_1_bias")] = interleave  # pyrefly: ignore[unsupported-operation]
+    hooks[(f"{prefix}-GptOssMlp-wi_0_bias", f"{prefix}-GptOssMlp-wi_1_bias")] = (
+        interleave  # pyrefly: ignore[unsupported-operation]
+    )
 
   return hooks
 

src/maxtext/common/checkpointing.py

@@ -1151,9 +1151,13 @@ def save_checkpoint(checkpoint_manager, step, state, config=None, data_iterator=
     if isinstance(data_iterator, RemoteIteratorWrapper):
       # Pass the wrapper directly; GrainCheckpointHandler will call save_state with the step
       save_args_composite["iter"] = GrainCheckpointSave(item=data_iterator)  # pyrefly: ignore[bad-assignment]
-    elif not isinstance(data_iterator, list) and isinstance(data_iterator.local_iterator, ElasticIterator):  # pyrefly: ignore[missing-attribute]
+    elif not isinstance(data_iterator, list) and isinstance(
+        data_iterator.local_iterator, ElasticIterator
+    ):  # pyrefly: ignore[missing-attribute]
       # ElasticIterator checkpoints a single global scalar shared by all shards.
-      save_args_composite["iter"] = GrainCheckpointSave(item=data_iterator.local_iterator)  # pyrefly: ignore[bad-assignment]
+      save_args_composite["iter"] = GrainCheckpointSave(
+          item=data_iterator.local_iterator
+      )  # pyrefly: ignore[bad-assignment]
     else:
       if not isinstance(data_iterator, list):
         data_iterator = [data_iterator]
@@ -1163,7 +1167,9 @@ def save_checkpoint(checkpoint_manager, step, state, config=None, data_iterator=
         process_count_total = process_count_total // config.expansion_factor_real_data
       for i, data_iter in enumerate(data_iterator):
         process_index = jax.process_index() + i * jax.process_count()
-        grain_iters_to_save.append((data_iter.local_iterator, process_index, process_count_total))  # pyrefly: ignore[missing-attribute]
+        grain_iters_to_save.append(
+            (data_iter.local_iterator, process_index, process_count_total)
+        )  # pyrefly: ignore[missing-attribute]
       save_args_composite["iter"] = GrainCheckpointSave(item=grain_iters_to_save)  # pyrefly: ignore[bad-assignment]
 
   custom_metadata = {}

src/maxtext/inference/decode.py

@@ -184,7 +184,9 @@ def main(argv: Sequence[str]) -> None:
           positions=position_ids,  # pyrefly: ignore[bad-argument-type]
           mrope_deltas=mrope_position_deltas,  # pyrefly: ignore[bad-argument-type]
           images=processor_outputs.pixel_values if config.use_multimodal else None,  # pyrefly: ignore[bad-argument-type]
-          image_masks=processor_outputs.pixel_mask if config.use_multimodal and "llama4" in config.model_name else None,  # pyrefly: ignore[bad-argument-type]
+          image_masks=processor_outputs.pixel_mask
+          if config.use_multimodal and "llama4" in config.model_name
+          else None,  # pyrefly: ignore[bad-argument-type]
           videos=getattr(processor_outputs, "video_values", None) if config.use_multimodal else None,
           audio_values=processor_outputs.audio_values if config.use_audio else None,  # pyrefly: ignore[bad-argument-type]
           audio_masks=processor_outputs.audio_mask if config.use_audio else None,  # pyrefly: ignore[bad-argument-type]

src/maxtext/inference/kvcache.py

@@ -844,7 +844,9 @@ def get_cached_values(self, cache_vars, target_dtype, cache_axis_order) -> jax.A
       elif dtype == jnp.float8_e4m3fn:
         scale_value /= E4M3_MAX
 
-      cache_value = KVTensor(qvalue=cache_value, scale=[scale_value], scale_t=None, dequant_dtype=target_dtype, bias=[])  # pyrefly: ignore[unexpected-keyword]
+      cache_value = KVTensor(
+          qvalue=cache_value, scale=[scale_value], scale_t=None, dequant_dtype=target_dtype, bias=[]
+      )  # pyrefly: ignore[unexpected-keyword]
     cache_value_in_logical_shape = jax.tree.map(lambda x: reverse_transpose(x, cache_axis_order), cache_value)
     return cache_value_in_logical_shape
 

src/maxtext/inference/maxengine/maxengine.py

@@ -219,7 +219,9 @@ def _nnx_run_model(
     nnx.replace_by_pure_dict(cache_state, cache_dict)
     # copy=True avoids reusing Variable objects across traces (TraceContextError),
     # mirroring the workaround in train.py's diff_wrapper.
-    model = nnx.merge(self.graphdef, params, cache_state, self._nnx_rest_state, copy=True)  # pyrefly: ignore[no-matching-overload]
+    model = nnx.merge(
+        self.graphdef, params, cache_state, self._nnx_rest_state, copy=True
+    )  # pyrefly: ignore[no-matching-overload]
     logits = model(
         decoder_input_tokens,
         decoder_positions,
@@ -280,7 +282,9 @@ def _layout(x, s, l):
       if x.format == l:
         return x
       # Somehow this can be None sometimes.
-      dll = (l.layout if jax.__version_info__ >= (0, 6, 3) else l.device_local_layout) if isinstance(l, Format) else l  # pyrefly: ignore[missing-attribute]
+      dll = (
+          (l.layout if jax.__version_info__ >= (0, 6, 3) else l.device_local_layout) if isinstance(l, Format) else l
+      )  # pyrefly: ignore[missing-attribute]
       f = jax.jit(self._identity, out_shardings=Format(dll, s)).lower(x).compile(compiler_options=xla_flags)
       y = f(x)
       # Achieves donation of the input argument, but allows for different memory
@@ -409,7 +413,10 @@ def _load_params_nnx(self, params, rng):
       with nn_partitioning.axis_rules(self.config.logical_axis_rules):
         full_sharding = sharding.nnx_construct_named_sharding(full_abs, self._mesh)
       concrete_model = maxtext_utils_nnx.create_nnx_sharded_model(
-          self.model, self._create_model_fn, mesh=self._mesh, named_sharding=full_sharding  # pyrefly: ignore[bad-argument-type]
+          self.model,
+          self._create_model_fn,
+          mesh=self._mesh,
+          named_sharding=full_sharding,  # pyrefly: ignore[bad-argument-type]
       )
       graphdef, _, _, rest_state = nnx.split(concrete_model, nnx.Param, nnx.Cache, ...)
       self.graphdef = graphdef

src/maxtext/input_pipeline/grain_tokenizer.py

@@ -71,7 +71,9 @@ class TokenizeAndTrim(TokenizerTransformBase, grain.MapTransform):
 
   def map(self, element: dict[str, Any]) -> dict[str, Any]:
     """Maps to each element."""
-    for feature_name, max_length in zip(self.feature_names, self.sequence_length, strict=True):  # pyrefly: ignore[bad-argument-type]
+    for feature_name, max_length in zip(
+        self.feature_names, self.sequence_length, strict=True
+    ):  # pyrefly: ignore[bad-argument-type]
       text = element[feature_name]
       token_ids = self._encode(text)[:max_length]
       element[feature_name] = np.asarray(token_ids, dtype=np.int32)
@@ -88,7 +90,9 @@ def __post_init__(self):
     super().__post_init__()
     # TokenizeAndChunk only supports single feature for chunking
     assert len(self.feature_names) == 1, "TokenizeAndChunk only supports single feature name"
-    assert len(self.sequence_length) == 1, "TokenizeAndChunk only supports single sequence length"  # pyrefly: ignore[bad-argument-type]
+    assert (
+        len(self.sequence_length) == 1
+    ), "TokenizeAndChunk only supports single sequence length"  # pyrefly: ignore[bad-argument-type]
     self.feature_name = self.feature_names[0]  # For backward compatibility
     self.sequence_length = self.sequence_length[0]  # Convert back to int for chunking  # pyrefly: ignore[bad-index]
 
@@ -104,7 +108,9 @@ def flat_map(self, element: dict[str, Any]) -> list[dict[str, Any]]:
 
     output_elements = []
     for start_idx in range(0, len(token_ids), chunk_size):  # pyrefly: ignore[bad-argument-type]
-      chunk = np.asarray(token_ids[start_idx : start_idx + chunk_size], dtype=np.int32)  # pyrefly: ignore[unsupported-operation]
+      chunk = np.asarray(
+          token_ids[start_idx : start_idx + chunk_size], dtype=np.int32
+      )  # pyrefly: ignore[unsupported-operation]
       new_element = {self.feature_name: chunk}
       output_elements.append(new_element)
 

src/maxtext/input_pipeline/input_pipeline_utils.py

@@ -172,7 +172,9 @@ def is_conversational(features, data_columns):
   for column in data_columns:
     messages = features[column]
     if isinstance(messages, datasets.Sequence):
-      if isinstance(messages.feature, dict) and "role" in messages.feature and "content" in messages.feature:  # pyrefly: ignore[missing-attribute]
+      if (
+          isinstance(messages.feature, dict) and "role" in messages.feature and "content" in messages.feature
+      ):  # pyrefly: ignore[missing-attribute]
         return True
 
   return False
@@ -840,11 +842,19 @@ def map(
         if isinstance(element[data_column], mm_utils.PreprocessorOutput):
           raise TypeError("Only 'images' column can be of type PreprocessorOutput.")
 
-        element[f"{data_column}_segmentation"] = element[data_column] != self.pad_id  # pyrefly: ignore[unsupported-operation]
-        element[f"{data_column}_segmentation"] = element[f"{data_column}_segmentation"].astype(np.int32)  # pyrefly: ignore[missing-attribute]
-        element[f"{data_column}_position"] = np.arange(element[data_column].shape[0], dtype=np.int32)  # pyrefly: ignore[missing-attribute]
+        element[f"{data_column}_segmentation"] = (
+            element[data_column] != self.pad_id
+        )  # pyrefly: ignore[unsupported-operation]
+        element[f"{data_column}_segmentation"] = element[f"{data_column}_segmentation"].astype(
+            np.int32
+        )  # pyrefly: ignore[missing-attribute]
+        element[f"{data_column}_position"] = np.arange(
+            element[data_column].shape[0], dtype=np.int32
+        )  # pyrefly: ignore[missing-attribute]
         if self.add_true_length:
-          element[f"{data_column}_true_length"] = np.array([element[data_column].shape[0]], dtype=np.int32)  # pyrefly: ignore[missing-attribute]
+          element[f"{data_column}_true_length"] = np.array(
+              [element[data_column].shape[0]], dtype=np.int32
+          )  # pyrefly: ignore[missing-attribute]
 
     for key, _ in element.items():
       if key == "images":

src/maxtext/integration/vllm/maxtext_vllm_adapter/adapter.py

@@ -345,9 +345,7 @@ def load_weights(self, rng_key: jax.Array) -> None:
           self.maxtext_config, mesh=self.mesh, model_mode=self.model_mode, rng_key=rng_key
       )
       if self.maxtext_config.lora.enable_lora:
-        model = lora_utils.apply_lora_to_model(
-            model, self.mesh, self.maxtext_config
-        )
+        model = lora_utils.apply_lora_to_model(model, self.mesh, self.maxtext_config)
         if self.maxtext_config.lora.lora_restore_path:
           lora_utils.restore_lora_from_path(model, self.maxtext_config)
       self.model = nnx.data(model)

src/maxtext/kernels/attention/splash_attention_kernel.py

@@ -624,7 +624,9 @@ def _apply_mask_and_soft_cap(
 
       repeats, rem = divmod(k_slice.size, NUM_LANES)
       assert rem == 0
-      q_sequence = jnp.tile(q_sequence_ref[...], (1, repeats))  # [bq, k_slice.size]  # pyrefly: ignore[unsupported-operation]
+      q_sequence = jnp.tile(
+          q_sequence_ref[...], (1, repeats)
+      )  # [bq, k_slice.size]  # pyrefly: ignore[unsupported-operation]
     else:
       assert q_sequence_ref.shape == (NUM_SUBLANES, bq)  # pyrefly: ignore[missing-attribute]
 
@@ -1694,7 +1696,9 @@ def run():
   if is_mqa:
     should_write = jnp.logical_and(should_write, q_head_index == num_q_heads - 1)
   elif num_kv_heads < num_q_heads:
-    should_write = jnp.logical_and(should_write, q_head_index_per_kv_head == q_heads_per_kv_heads - 1)  # pyrefly: ignore[unsupported-operation]
+    should_write = jnp.logical_and(
+        should_write, q_head_index_per_kv_head == q_heads_per_kv_heads - 1
+    )  # pyrefly: ignore[unsupported-operation]
 
   @pl.when(should_write)
   def end():
@@ -2425,8 +2429,12 @@ def manual_sharding_spec(self, sharding: jax.sharding.NamedSharding):
         data_next=spec if self.fwd_mask_info.data_next is not None else None,  # pyrefly: ignore[bad-argument-type]
         mask_next=spec if self.fwd_mask_info.mask_next is not None else None,  # pyrefly: ignore[bad-argument-type]
         block_mask=spec if self.fwd_mask_info.block_mask is not None else None,  # pyrefly: ignore[bad-argument-type]
-        partial_mask_blocks=partial_mask_blocks_spec if self.fwd_mask_info.partial_mask_blocks is not None else None,  # pyrefly: ignore[bad-argument-type]
-        q_sequence=q_sequence_spec if self.fwd_mask_info.q_sequence is not None else None,  # pyrefly: ignore[bad-argument-type]
+        partial_mask_blocks=partial_mask_blocks_spec
+        if self.fwd_mask_info.partial_mask_blocks is not None
+        else None,  # pyrefly: ignore[bad-argument-type]
+        q_sequence=q_sequence_spec
+        if self.fwd_mask_info.q_sequence is not None
+        else None,  # pyrefly: ignore[bad-argument-type]
     )
     return SplashAttentionKernel(
         mask_info_specs,

src/maxtext/kernels/gather_reduce_sc.py

@@ -570,7 +570,9 @@ def get_row_val(row_idx):
             row0 = arith.mulf(row0, weights_vecs[0])  # pyrefly: ignore[unsupported-operation]
             if topk_wgt_zero_nan:
               row0 = arith.select(
-                  arith.cmpf(arith.CmpFPredicate.OEQ, weights_vecs[0], zero_vec_f32),  # pyrefly: ignore[unbound-name, unsupported-operation]
+                  arith.cmpf(
+                      arith.CmpFPredicate.OEQ, weights_vecs[0], zero_vec_f32
+                  ),  # pyrefly: ignore[unbound-name, unsupported-operation]
                   zero_vec_f32,
                   row0,
               )
@@ -580,7 +582,9 @@ def get_row_val(row_idx):
             row8 = arith.mulf(row8, weights_vecs[8])  # pyrefly: ignore[unsupported-operation]
             if topk_wgt_zero_nan:
               row8 = arith.select(
-                  arith.cmpf(arith.CmpFPredicate.OEQ, weights_vecs[8], zero_vec_f32),  # pyrefly: ignore[unsupported-operation]
+                  arith.cmpf(
+                      arith.CmpFPredicate.OEQ, weights_vecs[8], zero_vec_f32
+                  ),  # pyrefly: ignore[unsupported-operation]
                   zero_vec_f32,
                   row8,
               )