Adding type suppressions for pyrefly

PiperOrigin-RevId: 947186144

Author: h-joo

src/maxtext/common/checkpointing.py

@@ -146,7 +146,7 @@ def restore_single_process(item, process_index, process_count):
 
     if isinstance(item, list):
       restored_items = []
-      for data_iter, process_idx in zip(item, process_index):
+      for data_iter, process_idx in zip(item, process_index):  # pyrefly: ignore[bad-argument-type]
         restored_items.append(restore_single_process(data_iter, process_idx, process_count))
       return restored_items
     else:
@@ -423,7 +423,7 @@ def create_orbax_checkpoint_manager(
 
   if dataset_type is not None and dataset_type == "grain":
     item_names += ("iter",)
-    item_handlers["iter"] = GrainCheckpointHandler()
+    item_handlers["iter"] = GrainCheckpointHandler()  # pyrefly: ignore[bad-assignment]
 
   # local storage checkpoint needs parent directory created
   p = gcs_utils.mkdir_and_check_permissions(checkpoint_dir)
@@ -762,7 +762,7 @@ def load_state_if_possible(
   if checkpoint_manager is not None:
     max_logging.log("checkpoint manager exists so trying to load this run's existing checkpoint")
 
-    step = checkpoint_manager.latest_step() if step < 0 else step
+    step = checkpoint_manager.latest_step() if step < 0 else step  # pyrefly: ignore[bad-assignment]
     if step is not None:
       max_logging.log(f"restoring from this run's directory step {step}")
 
@@ -1150,10 +1150,10 @@ def save_checkpoint(checkpoint_manager, step, state, config=None, data_iterator=
   if config and config.dataset_type == "grain" and not isinstance(data_iterator, PlaceHolderDataIterator):
     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)
-    elif not isinstance(data_iterator, list) and isinstance(data_iterator.local_iterator, ElasticIterator):
+      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]
       # ElasticIterator checkpoints a single global scalar shared by all shards.
-      save_args_composite["iter"] = GrainCheckpointSave(item=data_iterator.local_iterator)
+      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,8 +1163,8 @@ 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))
-      save_args_composite["iter"] = GrainCheckpointSave(item=grain_iters_to_save)
+        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 = {}
   if config:

src/maxtext/utils/model_creation_utils.py

@@ -777,7 +777,7 @@ def create_models_and_meshes(trainer_config, sampler_config, trainer_devices, sa
           use_no_op_mappings=use_no_op_mappings,
           pad_id=tokenizer_pad_id,
       )
-      actor_model.config = None
+      actor_model.config = None  # pyrefly: ignore[missing-attribute]
     actor_mesh = reference_mesh
   else:
     max_logging.log("Creating policy model with same config as reference model on trainer mesh")
@@ -902,7 +902,7 @@ def from_pretrained(
     load_parameters_path = epath.Path(config.base_output_directory) / "0" / "items"
     # Create a copied Pydantic model with the updated values
     pydantic_config = getattr(config, "_pydantic_config", config)
-    new_config = pydantic_config.model_copy(
+    new_config = pydantic_config.model_copy(  # pyrefly: ignore[missing-attribute]
         update={
             "load_parameters_path": load_parameters_path,
         }
@@ -930,7 +930,7 @@ def from_pretrained(
   sharded_state = nnx.state(model)
 
   if mesh is None:
-    mesh = model.mesh
+    mesh = model.mesh  # pyrefly: ignore[missing-attribute]
 
   with mesh:
     if config.load_parameters_path:
@@ -1114,7 +1114,7 @@ def _free_device_memory(node):
       )
 
       if is_nnx_checkpoint:
-        restored_root = restored["base"] if has_base_key else restored
+        restored_root = restored["base"] if has_base_key else restored  # pyrefly: ignore[unbound-name]
         checkpoint = jax.tree.map(
             lambda v: v["value"],
             restored_root,
@@ -1202,11 +1202,11 @@ def _walk_align(ckpt, model_arr, axes):
       with mesh:
         use_no_op_mappings = "maxtext_config" in config.vllm_additional_config
         model = TunixMaxTextAdapter(
-            base_model=model,
+            base_model=model,  # pyrefly: ignore[bad-argument-type]
             use_no_op_mappings=use_no_op_mappings,
             pad_id=tokenizer_pad_id,
         )
-        model.config = None
+        model.config = None  # pyrefly: ignore[missing-attribute]
 
     if original_mesh:
       return model