Automated Code Change

PiperOrigin-RevId: 944783931

Author: h-joo

src/maxtext/checkpoint_conversion/utils/hf_model_configs.py

@@ -19,7 +19,7 @@
 
 import transformers
 
-if transformers.__version__ >= "5.0.0":
+if transformers.__version__ >= "5.0.0":  # pyrefly: ignore[missing-attribute]
   from transformers.configuration_utils import PreTrainedConfig as PTConfig  # pytype: disable=import-error
 else:
   from transformers.configuration_utils import PretrainedConfig as PTConfig
@@ -120,7 +120,7 @@
 
 gemma4_31b_dict = gemma4_26b_dict.copy()
 gemma4_31b_dict["text_config"] = gemma4_26b_dict["text_config"].copy()
-gemma4_31b_dict["text_config"].update(
+gemma4_31b_dict["text_config"].update(  # pyrefly: ignore[no-matching-overload]
     {
         "enable_moe_block": False,
         "hidden_size": 5376,
@@ -265,10 +265,10 @@
 
 try:
   # Will execute successfully if Transformers is updated with Gemma 4 support
-  gemma4_26b_config = transformers.Gemma4Config(**gemma4_26b_dict)
-  gemma4_31b_config = transformers.Gemma4Config(**gemma4_31b_dict)
-  gemma4_e2b_config = transformers.Gemma4Config(**gemma4_e2b_dict)
-  gemma4_e4b_config = transformers.Gemma4Config(**gemma4_e4b_dict)
+  gemma4_26b_config = transformers.Gemma4Config(**gemma4_26b_dict)  # pyrefly: ignore[missing-attribute]
+  gemma4_31b_config = transformers.Gemma4Config(**gemma4_31b_dict)  # pyrefly: ignore[missing-attribute]
+  gemma4_e2b_config = transformers.Gemma4Config(**gemma4_e2b_dict)  # pyrefly: ignore[missing-attribute]
+  gemma4_e4b_config = transformers.Gemma4Config(**gemma4_e4b_dict)  # pyrefly: ignore[missing-attribute]
 except AttributeError:
   # Graceful fallback to raw dict-based PTConfig if Gemma 4 natively is missing
   gemma4_26b_config = PTConfig(**gemma4_26b_dict)  # pytype: disable=wrong-arg-types
@@ -1011,7 +1011,7 @@
 
 
 # TODO(shuningjin): replace with DeepseekV32Config when available in transformers library
-class DeepseekV32Config(PTConfig):
+class DeepseekV32Config(PTConfig):  # pyrefly: ignore[invalid-inheritance]
   model_type = "deepseek_v32"
 
   def __init__(self, **kwargs):
@@ -1089,7 +1089,7 @@ def __init__(self, **kwargs):
     "use_cache": True,
     "vocab_size": 201088,
 }
-gpt_oss_20b_config = transformers.GptOssConfig(**gpt_oss_20b_dict)
+gpt_oss_20b_config = transformers.GptOssConfig(**gpt_oss_20b_dict)  # pyrefly: ignore[bad-argument-type]
 
 # from https://huggingface.co/openai/gpt-oss-120b/blob/main/config.json
 # remove mxfp4 quantization_config, since we are using bf16
@@ -1171,7 +1171,7 @@ def __init__(self, **kwargs):
     "use_cache": True,
     "vocab_size": 201088,
 }
-gpt_oss_120b_config = transformers.GptOssConfig(**gpt_oss_120b_dict)
+gpt_oss_120b_config = transformers.GptOssConfig(**gpt_oss_120b_dict)  # pyrefly: ignore[bad-argument-type]
 
 
 qwen3_omni_30b_a3b_config = transformers.Qwen3OmniMoeConfig(
@@ -1481,8 +1481,8 @@ def __init__(self, **kwargs):
 
 try:
   # Will execute successfully if Transformers is updated with Qwen3.5 support
-  qwen3_5_35b_a3b_config = transformers.Qwen3_5MoeConfig(**qwen3_5_35b_a3b_dict)
-  qwen3_5_397b_a17b_config = transformers.Qwen3_5MoeConfig(**qwen3_5_397b_a17b_dict)
+  qwen3_5_35b_a3b_config = transformers.Qwen3_5MoeConfig(**qwen3_5_35b_a3b_dict)  # pyrefly: ignore[missing-attribute]
+  qwen3_5_397b_a17b_config = transformers.Qwen3_5MoeConfig(**qwen3_5_397b_a17b_dict)  # pyrefly: ignore[missing-attribute]
 except AttributeError:
   qwen3_5_35b_a3b_config = PTConfig(**qwen3_5_35b_a3b_dict)  # pytype: disable=wrong-arg-types
   qwen3_5_397b_a17b_config = PTConfig(**qwen3_5_397b_a17b_dict)  # pytype: disable=wrong-arg-types

src/maxtext/checkpoint_conversion/utils/param_mapping.py

@@ -139,7 +139,7 @@ def main(argv: Sequence[str]) -> None:
 
   if config.use_multimodal:
     tokens = mm_processor.prepare_text_for_image_fusion(tokens=tokens, config=config, processor_output=processor_outputs)
-    true_length += image_offsets
+    true_length += image_offsets  # pyrefly: ignore[unbound-name]
 
     if config.use_mrope:
       from maxtext.multimodal import processor_qwen3_omni  # pylint: disable=import-outside-toplevel
@@ -181,13 +181,13 @@ def main(argv: Sequence[str]) -> None:
       prefill_result, first_token = engine.prefill(
           params=params,
           padded_tokens=tokens,
-          positions=position_ids,
-          mrope_deltas=mrope_position_deltas,
-          images=processor_outputs.pixel_values if config.use_multimodal else None,
-          image_masks=processor_outputs.pixel_mask if config.use_multimodal and "llama4" in config.model_name else 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]
           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,
-          audio_masks=processor_outputs.audio_mask if config.use_audio 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]
           true_length=true_length,
           rng=rng_prefill,
           slot=i,

src/maxtext/inference/decode.py

@@ -400,8 +400,8 @@ def run_benchmarks(config):
           config,
           multisampling_prefill_executable[prefill_length],
           params,
-          prefill_tokens[prefill_length],
-          prefill_true_lengths[prefill_length],
+          prefill_tokens[prefill_length],  # pyrefly: ignore[unbound-name]
+          prefill_true_lengths[prefill_length],  # pyrefly: ignore[unbound-name]
           benchmark_loop_iters,
       )
 

src/maxtext/inference/inference_microbenchmark.py

@@ -277,7 +277,7 @@ def __init__(
       *,
       # Not used in KVCache but passed in by nnx_wrappers.to_linen.
       # TODO: Remove when bridge no longer needed
-      rngs: nnx.Rngs = None,
+      rngs: nnx.Rngs = None,  # pyrefly: ignore[bad-function-definition]
   ):
     """Initializes the KVCache module.
 
@@ -612,8 +612,8 @@ def kv_cache_chunked_prefill(
     # For quantized kv cached. Could be get without transpose twice.
     cached_key = self.get_cached_values(cached_prefill_key_vars, key.dtype, self.prefill_cache_axis_order)
     cached_value = self.get_cached_values(cached_prefill_value_vars, value.dtype, self.prefill_cache_axis_order)
-    cached_key_value = jnp.transpose(cached_key, self.prefill_cache_axis_order)
-    cached_value_value = jnp.transpose(cached_value, self.prefill_cache_axis_order)
+    cached_key_value = jnp.transpose(cached_key, self.prefill_cache_axis_order)  # pyrefly: ignore[bad-argument-type]
+    cached_value_value = jnp.transpose(cached_value, self.prefill_cache_axis_order)  # pyrefly: ignore[bad-argument-type]
 
     seq_axis = self.prefill_cache_logical_axis_names.index(CACHE_SEQUENCE)
     cache_seq_axis = self.prefill_cache_axis_order.index(seq_axis)
@@ -767,17 +767,17 @@ def update_ar_key_value(
     if use_ragged_attention:
       cache_locations = [slice(None)] * 4
       new_token_locations = [slice(None)] * 4
-      new_token_locations[ar_cache_sequence_axis] = 0
+      new_token_locations[ar_cache_sequence_axis] = 0  # pyrefly: ignore[unsupported-operation]
 
       def key_body(i, val):
         cache_locations[ar_cache_batch_axis] = i
-        cache_locations[ar_cache_sequence_axis] = lengths[i]
+        cache_locations[ar_cache_sequence_axis] = lengths[i]  # pyrefly: ignore[unsupported-operation]
         new_token_locations[ar_cache_batch_axis] = i
         return val.at[tuple(cache_locations)].set(one_token_key_shaped_for_cache[tuple(new_token_locations)])
 
       def value_body(i, val):
         cache_locations[ar_cache_batch_axis] = i
-        cache_locations[ar_cache_sequence_axis] = lengths[i]
+        cache_locations[ar_cache_sequence_axis] = lengths[i]  # pyrefly: ignore[unsupported-operation]
         new_token_locations[ar_cache_batch_axis] = i
         return val.at[tuple(cache_locations)].set(one_token_value_shaped_for_cache[tuple(new_token_locations)])
 
@@ -815,15 +815,15 @@ def value_body(i, val):
       cached_key_scale.set_value(
           jax.lax.dynamic_update_index_in_dim(
               cached_key_scale.get_value(),
-              one_token_key_scale_shaped_for_cache,
+              one_token_key_scale_shaped_for_cache,  # pyrefly: ignore[unbound-name]
               ar_cache_update_idx,
               ar_cache_scale_update_axis,
           )
       )
       cached_value_scale.set_value(
           jax.lax.dynamic_update_index_in_dim(
               cached_value_scale.get_value(),
-              one_token_value_scale_shaped_for_cache,
+              one_token_value_scale_shaped_for_cache,  # pyrefly: ignore[unbound-name]
               ar_cache_update_idx,
               ar_cache_scale_update_axis,
           )
@@ -844,7 +844,7 @@ 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=[])
+      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
 
@@ -1061,7 +1061,7 @@ def __init__(
       *,
       # Not used in MlaKVCache but passed in by nnx_wrappers.to_linen.
       # TODO: Remove when bridge no longer needed
-      rngs: nnx.Rngs = None,
+      rngs: nnx.Rngs = None,  # pyrefly: ignore[bad-function-definition]
   ):
     """Initializes the MlaKVCache module.
 

src/maxtext/inference/kvcache.py

@@ -101,7 +101,7 @@ def get_keys(self):
 _BaseEngine = engine_api.Engine if (not is_decoupled() and hasattr(engine_api, "Engine")) else object
 
 
-class MaxEngine(_BaseEngine):
+class MaxEngine(_BaseEngine):  # pyrefly: ignore[invalid-inheritance]
   """The computational core of the generative model server.
 
   Engine defines an API that models must adhere to as they plug into the
@@ -192,7 +192,7 @@ def _nnx_init_cache_dict(self, mode: str = MODEL_MODE_PREFILL) -> dict:
     """Zero-filled pure-dict cache matching the abstract NNX model."""
     src = self._abstract_model_for_mode(mode)
     _, cache_state, _ = nnx.split(src, nnx.Cache, ...)
-    cache_dict = cache_state.to_pure_dict()
+    cache_dict = cache_state.to_pure_dict()  # pyrefly: ignore[missing-attribute]
     return jax.tree.map(lambda x: jnp.zeros(x.shape, x.dtype), cache_dict)
 
   def _nnx_run_model(
@@ -219,7 +219,7 @@ 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)
+    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 +280,7 @@ 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
+      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 +409,7 @@ 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
+          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
@@ -435,16 +435,16 @@ def _load_params_nnx(self, params, rng):
       # PREFILL/AR attention ops have different cache variable shapes, and a
       # mismatch trips the `assert prefill_kv_cache` check inside attention_op.
       with nn_partitioning.axis_rules(self.config.logical_axis_rules):
-        concrete_model = self._create_model_fn()
+        concrete_model = self._create_model_fn()  # pyrefly: ignore[not-callable]
       graphdef, _, _, rest_state = nnx.split(concrete_model, nnx.Param, nnx.Cache, ...)
       # Overlay loaded non-Param/non-Cache leaves (e.g. AQT qrhs.frozen) onto
       # the PREFILL-mode rest_state. The PREFILL concrete_model already has
       # placeholder qrhs vars at the right paths; we just swap in the loaded
       # values. Anything only in `loaded_rest_state` (e.g. AR-only RNG slots)
       # is ignored. We keep PREFILL rest_state as the base so RNG variables
       # match the PREFILL graphdef's expectations.
-      loaded_rest_dict = loaded_rest_state.to_pure_dict()
-      rest_dict = rest_state.to_pure_dict()
+      loaded_rest_dict = loaded_rest_state.to_pure_dict()  # pyrefly: ignore[missing-attribute]
+      rest_dict = rest_state.to_pure_dict()  # pyrefly: ignore[missing-attribute]
 
       def _overlay(dst, src):
         if isinstance(dst, dict) and isinstance(src, dict):
@@ -768,7 +768,7 @@ def _prefill_jit(
     one_d_output = ones_to_keep * DECODING_ACTIVE_SEQUENCE_INDICATOR
     sequence_indicator = jnp.expand_dims(one_d_output, 0)
 
-    rng, new_rng = jax.random.split(rng)
+    rng, new_rng = jax.random.split(rng)  # pyrefly: ignore[bad-argument-type]
     if self.config.pure_nnx:
       # Prefill always operates on batch=1 (one padded prompt at a time).
       nnx_cache = (
@@ -1026,7 +1026,7 @@ def _prefill_multisampling_jit(
     one_d_output = ones_to_keep * DECODING_ACTIVE_SEQUENCE_INDICATOR
     sequence_indicator = jnp.expand_dims(one_d_output, 0)
 
-    rng, new_rng = jax.random.split(rng)
+    rng, new_rng = jax.random.split(rng)  # pyrefly: ignore[bad-argument-type]
     if self.config.pure_nnx:
       # Prefill is batch=1 (one prompt); multi-sampling only draws several first
       # tokens from the shared logits below. Mirror the _prefill_jit NNX branch.

src/maxtext/inference/maxengine/maxengine.py

@@ -203,9 +203,9 @@ def create_dataset_from_pattern(pattern):
     dataset = dataset.repeat(num_epoch)
     dataset = dataset[file_slice]
     if data_file_type == "tfrecord":
-      dataset = dataset.map(input_pipeline_utils.make_tfrecord_iter_dataset)
+      dataset = dataset.map(input_pipeline_utils.make_tfrecord_iter_dataset)  # pyrefly: ignore[missing-attribute]
     else:
-      dataset = dataset.map(grain.experimental.ParquetIterDataset)
+      dataset = dataset.map(grain.experimental.ParquetIterDataset)  # pyrefly: ignore[missing-attribute]
     cycle_length = min(files_per_host, grain_num_threads)
     dataset = grain.experimental.InterleaveIterDataset(dataset, cycle_length=cycle_length)
     if row_shard is not None:

src/maxtext/input_pipeline/grain_data_processing.py

@@ -71,7 +71,7 @@ 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):
+    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,9 +88,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"
+    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
+    self.sequence_length = self.sequence_length[0]  # Convert back to int for chunking  # pyrefly: ignore[bad-index]
 
   def flat_map(self, element: dict[str, Any]) -> list[dict[str, Any]]:
     """Tokenize and chunk text into multiple examples of sequence length."""
@@ -103,8 +103,8 @@ def flat_map(self, element: dict[str, Any]) -> list[dict[str, Any]]:
       return []
 
     output_elements = []
-    for start_idx in range(0, len(token_ids), chunk_size):
-      chunk = np.asarray(token_ids[start_idx : start_idx + chunk_size], dtype=np.int32)
+    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]
       new_element = {self.feature_name: chunk}
       output_elements.append(new_element)