Add pyrefly suppressions

examples/deepscaler/train_deepscaler_nb.py

@@ -373,7 +373,7 @@ def process_item(item):
 train_dataset, test_dataset = create_datasets()
 train_dataset, val_dataset = data_lib.post_init_dataset(
     train_dataset,
-    tokenizer,
+    tokenizer,  # pyrefly: ignore[bad-argument-type]
     batch_size=BATCH_SIZE,
     num_batches=NUM_BATCHES,
     max_prompt_length=MAX_PROMPT_LENGTH,
@@ -383,7 +383,7 @@ def process_item(item):
 
 test_dataset, _ = data_lib.post_init_dataset(
     test_dataset,
-    tokenizer,
+    tokenizer,  # pyrefly: ignore[bad-argument-type]
     batch_size=BATCH_SIZE,
     num_batches=NUM_TEST_BATCHES,
     max_prompt_length=MAX_PROMPT_LENGTH,

examples/deepswe/train_deepswe_nb.py

@@ -645,7 +645,7 @@ def mixed_type_batch_fn(elements):
 
 train_dataset, _ = data_lib.post_init_dataset(
     grain_dataset,
-    tokenizer,
+    tokenizer,  # pyrefly: ignore[bad-argument-type]
     batch_size=BATCH_SIZE,
     num_batches=None,
     max_prompt_length=MAX_PROMPT_LENGTH,
@@ -856,7 +856,7 @@ def get_lora_model(base_model, model_mesh):
     "temperature": TEMPERATURE,
     "top_p": TOP_P,
     "top_k": TOP_K,
-    "eos_tokens": [tokenizer.encode("<|im_end|>")[0]],
+    "eos_tokens": [tokenizer.encode("<|im_end|>")[0]],  # pyrefly: ignore[missing-attribute]
     "return_logprobs": USE_ROLLOUT_LOGPS,
     "max_tokens_to_generate": MAX_RESPONSE_LENGTH,
 }
@@ -911,9 +911,9 @@ def get_lora_model(base_model, model_mesh):
   }
   # Force no-op mappings for weight sync if both trainer and sampler use MaxText
   if hasattr(qwen_reference, "use_no_op_mappings"):
-    qwen_reference.use_no_op_mappings = True
+    qwen_reference.use_no_op_mappings = True  # pyrefly: ignore[missing-attribute]
   if hasattr(qwen_actor, "use_no_op_mappings"):
-    qwen_actor.use_no_op_mappings = True
+    qwen_actor.use_no_op_mappings = True  # pyrefly: ignore[missing-attribute]
     logging.info("Forced use_no_op_mappings=True on actor/reference models.")
 
 

examples/frozenlake/agent.py

@@ -198,7 +198,7 @@ def update_from_model(self, response: str, **kwargs) -> agent_types.Action:
     # Record complete step with conversation context and parsed action.
     cur_step = self._trajectory.steps[-1]
     cur_step.thought = thought
-    cur_step.action = action_str
+    cur_step.action = action_str  # pyrefly: ignore[bad-assignment]
     cur_step.model_response = response
 
     self.step += 1

examples/frozenlake/data.py

@@ -21,8 +21,8 @@
 
 import numpy as np
 import pandas as pd
-import datasets as datasets_lib
-import grain
+import datasets as datasets_lib  # pyrefly: ignore[missing-import]
+import grain  # pyrefly: ignore[missing-import]
 
 
 DEFAULT_DIR = os.getcwd()

examples/frozenlake/train_frozenlake.py

@@ -334,8 +334,8 @@ def process_item(item):
     item["prompts"] = ""
     return item
 
-  train_ds = grain.MapDataset.source(train_ds).map(process_item)
-  test_ds = grain.MapDataset.source(test_ds).map(process_item)
+  train_ds = grain.MapDataset.source(train_ds).map(process_item)  # pyrefly: ignore[bad-argument-type]
+  test_ds = grain.MapDataset.source(test_ds).map(process_item)  # pyrefly: ignore[bad-argument-type]
   return train_ds, test_ds
 
 
@@ -349,7 +349,7 @@ def process_item(item):
 train_dataset, test_dataset = create_datasets()
 train_dataset, val_dataset = data_lib.post_init_dataset(
     train_dataset,
-    tokenizer,
+    tokenizer,  # pyrefly: ignore[bad-argument-type]
     batch_size=BATCH_SIZE,
     num_batches=NUM_BATCHES,
     max_prompt_length=MAX_PROMPT_LENGTH,
@@ -358,7 +358,7 @@ def process_item(item):
 )
 test_dataset, _ = data_lib.post_init_dataset(
     test_dataset,
-    tokenizer,
+    tokenizer,  # pyrefly: ignore[bad-argument-type]
     batch_size=BATCH_SIZE,
     num_batches=NUM_TEST_BATCHES,
     max_prompt_length=MAX_PROMPT_LENGTH,

examples/math_gsm8k/gemma_grpo_demo_nb.py

@@ -291,7 +291,7 @@ def extract_hash_answer(text: str) -> str | None:
 def get_ref_model():
   """Loads the reference model, from CNS in g3 or Kaggle in OSS."""
   mesh = jax.make_mesh(
-      *MESH, axis_types=(jax.sharding.AxisType.Auto,) * len(MESH[0])
+      *MESH, axis_types=(jax.sharding.AxisType.Auto,) * len(MESH[0])  # pyrefly: ignore[bad-argument-type]
   )
 
   if ENV == 'g3':

examples/math_gsm8k/qwen3_grpo_demo.py

@@ -625,7 +625,7 @@ def main() -> None:
       trust_remote_code=True,
   )
   chat_parser = VTCRawTextParser()
-  qwen_eos_tokens = tokenizer.encode("<|im_end|>", add_special_tokens=False)
+  qwen_eos_tokens = tokenizer.encode("<|im_end|>", add_special_tokens=False)  # pyrefly: ignore[missing-attribute]
 
   reference, actor = create_reference_and_actor(shared_mesh)
   show_hbm_usage("after loading qwen_ref / qwen_actor")

tests/experimental/train/peft_trainer_v2_test.py

@@ -149,19 +149,19 @@ def test_basic_training(self, cache_nnx_graph: bool):
     jax.tree.map_with_path(tc.assert_not_equal, original_variables, variables)
 
     self.assertGreater(
-        trainer.metrics_logger.get_metric('', 'perplexity', 'train'), 0
+        trainer.metrics_logger.get_metric('', 'perplexity', 'train'), 0  # pyrefly: ignore[missing-attribute]
     )
     self.assertEqual(
-        trainer.metrics_logger.get_metric('', 'learning_rate', 'train'),
+        trainer.metrics_logger.get_metric('', 'learning_rate', 'train'),  # pyrefly: ignore[missing-attribute]
         TEST_LEARNING_RATE,
     )
     self.assertGreater(
-        trainer.metrics_logger.get_metric('', 'perplexity', 'eval'), 0
+        trainer.metrics_logger.get_metric('', 'perplexity', 'eval'), 0  # pyrefly: ignore[missing-attribute]
     )
     self.assertGreater(trainer._train_steps, 0)
 
     self.assertLen(
-        trainer.metrics_logger.get_metric_history('', 'perplexity', 'train'),
+        trainer.metrics_logger.get_metric_history('', 'perplexity', 'train'),  # pyrefly: ignore[missing-attribute]
         trainer._train_steps,
     )
 
@@ -416,7 +416,7 @@ def test_lora_training(self, learning_rate_scheduler):
         tc.assert_not_equal, original_lora_params, lora_params
     )
     self.assertEqual(
-        trainer.metrics_logger.get_metric('', 'learning_rate', 'train'),
+        trainer.metrics_logger.get_metric('', 'learning_rate', 'train'),  # pyrefly: ignore[missing-attribute]
         TEST_LEARNING_RATE,
     )
 
@@ -446,7 +446,7 @@ def train(
 
       trainer.train(train_ds, self.eval_ds)
       self.assertEqual(
-          trainer.metrics_logger.get_metric('', 'learning_rate', 'train'),
+          trainer.metrics_logger.get_metric('', 'learning_rate', 'train'),  # pyrefly: ignore[missing-attribute]
           TEST_LEARNING_RATE,
       )
       return nnx.state(model, nnx.Param), trainer
@@ -675,7 +675,7 @@ def _post_process_train_step(self, aux):
         if self._buffered_train_metrics is not None:
           self._buffered_train_metrics.additional_metrics['foo'] = (
               [aux['foo']],
-              lambda xs: xs[-1],
+              lambda xs: xs[-1],  # pyrefly: ignore[bad-index]
           )
 
       def _post_process_eval_step(self, aux):
@@ -686,7 +686,7 @@ def _post_process_eval_step(self, aux):
         if self._buffered_eval_metrics is not None:
           self._buffered_eval_metrics.additional_metrics['foo'] = (
               [aux['foo']],
-              lambda xs: xs[-1],
+              lambda xs: xs[-1],  # pyrefly: ignore[bad-index]
           )
 
     config = peft_trainer_v2.TrainingConfig(eval_every_n_steps=2, max_steps=100)
@@ -732,7 +732,7 @@ def test_get_metrics(self):
     self.assertEqual(metrics.id, 0)
     self.assertEqual(metrics.mode, 'eval')
     self.assertIn('loss', metrics.scalar_metrics)
-    self.assertGreater(metrics.scalar_metrics['loss'], 0)
+    self.assertGreater(metrics.scalar_metrics['loss'], 0)  # pyrefly: ignore[no-matching-overload]
 
     # After calling get_metrics, the buffer should be cleared
     self.assertEqual(trainer.get_metrics().id, -1)
@@ -744,7 +744,7 @@ def test_get_metrics(self):
     self.assertEqual(train_metrics.mode, 'train')
     self.assertIn('loss', train_metrics.scalar_metrics)
     self.assertIn('grad_norm', train_metrics.scalar_metrics)
-    self.assertGreater(train_metrics.scalar_metrics['loss'], 0)
+    self.assertGreater(train_metrics.scalar_metrics['loss'], 0)  # pyrefly: ignore[no-matching-overload]
 
   def test_injected_params(self):
     config = peft_trainer_v2.TrainingConfig(eval_every_n_steps=2, max_steps=100)
@@ -760,7 +760,7 @@ def test_injected_params(self):
     trainer = trainer.with_gen_model_input_fn(dummy_gen_model_input_fn)
     trainer.train(self.train_ds, self.eval_ds)
     self.assertEqual(
-        trainer.metrics_logger.get_metric('', 'learning_rate', 'train'),
+        trainer.metrics_logger.get_metric('', 'learning_rate', 'train'),  # pyrefly: ignore[missing-attribute]
         TEST_LEARNING_RATE,
     )
 

tests/generate/sampler_test.py

@@ -431,7 +431,7 @@ def test_decode_stops_after_prefill_for_single_generation_step(self):
     prompt_tokens = sampler.tokenize('input string')
     all_input_ids = jnp.array([
         utils.pad_to_length(
-            prompt_tokens,
+            prompt_tokens,  # pyrefly: ignore[bad-argument-type]
             target_length=max_prompt_length,
             pad_value=vocab.pad_id(),
             left=True,

tests/sft/otel_wandb_test.py

@@ -35,7 +35,7 @@ class WandbMetricsExporterTest(absltest.TestCase):
 
   def setUp(self):
     super().setUp()
-    self.run = _FakeWandbRun()
+    self.run = _FakeWandbRun()  # pyrefly: ignore[bad-assignment]
     self.exporter = otel_wandb.WandbMetricsExporter(self.run)
     self.reader = otel_sdk_export.PeriodicExportingMetricReader(
         self.exporter, export_interval_millis=3_600_000
@@ -76,8 +76,8 @@ def test_double_write_reaches_wandb_run(self):
 
     self.meter_provider.force_flush()
 
-    self.assertLen(self.run.calls, 1)
-    values, step = self.run.calls[0]
+    self.assertLen(self.run.calls, 1)  # pyrefly: ignore[missing-attribute]
+    values, step = self.run.calls[0]  # pyrefly: ignore[missing-attribute]
     self.assertEqual(step, 3)
     self.assertAlmostEqual(values["actor/train/tunix.training.loss"], 0.5)
     self.assertAlmostEqual(
@@ -91,10 +91,10 @@ def test_groups_are_logged_in_step_order(self):
 
     self.meter_provider.force_flush()
 
-    steps = [step for _, step in self.run.calls]
+    steps = [step for _, step in self.run.calls]  # pyrefly: ignore[missing-attribute]
     self.assertEqual(steps, sorted(steps))
     logged_keys = set()
-    for values, _ in self.run.calls:
+    for values, _ in self.run.calls:  # pyrefly: ignore[missing-attribute]
       logged_keys.update(values)
     self.assertIn("actor/train/tunix.training.loss", logged_keys)
     self.assertIn("critic/eval/tunix.training.loss", logged_keys)

tests/utils/mesh_test.py

@@ -499,7 +499,7 @@ def __init__(self, device_id, coords):
 
     allocated, _ = mesh._allocate_devices_by_coords(fake_devices, 2)
 
-    self.assertEqual([device.id for device in allocated], [0, 1])
+    self.assertEqual([device.id for device in allocated], [0, 1])  # pyrefly: ignore[not-iterable]
 
   def test_allocate_named_mesh_device_slices_prefers_coord_boxes(self):
     class FakeDevice:
@@ -701,7 +701,7 @@ def __init__(self, device_id, coords):
         allocation_policy="COMPACT",
     )
 
-    allocated_coords = [device.coords for device in allocated]
+    allocated_coords = [device.coords for device in allocated]  # pyrefly: ignore[not-iterable]
     mins = tuple(
         min(coords[dim] for coords in allocated_coords) for dim in range(3)
     )

tunix/cli/base_rl_pipeline.py

@@ -589,7 +589,7 @@ def create_rl_engine(self, tokenizer) -> rl_engine_lib.RLEngine:
           params=jax.random.key(critic_model_config.get("rng_seed", 0))
       )
 
-      if hasattr(critic_model.config.shd_config, "score_weight_d1"):
+      if hasattr(critic_model.config.shd_config, "score_weight_d1"):  # pyrefly: ignore[missing-attribute]
         critic_model = rl_utils.TransformerWithScoreHead(
             critic_model, rngs=rngs
         )
@@ -886,7 +886,7 @@ def _run(self, mode: str):
   def run_trainer(self):
     """Dispatch to standard or agentic trainer based on training_mode."""
     mode = self.config.get("training_mode", self._default_training_mode)
-    self._run(mode=mode)
+    self._run(mode=mode)  # pyrefly: ignore[bad-argument-type]
 
 
 def setup_jax_pathways(pathways_bns: str):

tunix/cli/config.py

@@ -283,7 +283,7 @@ def _config_mapping(self, key: str) -> dict[str, Any]:
           f" {type(value).__name__}."
       )
     if isinstance(value, omegaconf.DictConfig):
-      return omegaconf.OmegaConf.to_container(value, resolve=True)
+      return omegaconf.OmegaConf.to_container(value, resolve=True)  # pyrefly: ignore[bad-return]
     return dict(value)
 
   def _mutable_config_mapping(self, key: str) -> MutableMapping[str, Any]:

tunix/examples/data/translation_dataset.py

@@ -69,7 +69,7 @@ def create_datasets(
         dataset_name, split=("train", "valid"), download=tfds_download
     )
   elif dataset_name == "Helsinki-NLP/opus-100":  # Hugging Face dataloader
-    train_ds, eval_ds = datasets.load_dataset(
+    train_ds, eval_ds = datasets.load_dataset(  # pyrefly: ignore[no-matching-overload]
         dataset_name, data_dir="en-fr", split=("train", "validation")  # pyrefly: ignore[bad-argument-type]
     )
   else:

tunix/experimental/distributed/examples/basics/door.py

@@ -17,7 +17,7 @@ def main(argv, context: ProcessContext | None) -> None:
   knock_future = futures.Future()
   assert context is not None
   context.ipc.discovery.on_register(
-      callback=lambda hostname, _, metadata: (
+      callback=lambda hostname, _, metadata: (  # pyrefly: ignore[bad-argument-type]
           logging.info(
               f"{hostname} knocked and said: {pickle.loads(metadata)}"
           ),

tunix/experimental/distributed/examples/vllm_rollout/rollout.py

@@ -64,7 +64,7 @@ def main(argv: Sequence[str], context: ProcessContext | None) -> None:
 
   tokenizer = AutoTokenizer.from_pretrained(args.model_name)
   config = vllm_sampler.VllmConfig(engine_kwargs={"model": args.model_name})
-  sampler_server = legacy_sampler_lib.LegacyVllmSamplerAdapter(
+  sampler_server = legacy_sampler_lib.LegacyVllmSamplerAdapter(  # pyrefly: ignore[bad-instantiation]
       server_id="vllm-0",
       tokenizer=tokenizer,
       config=config,

tunix/experimental/orchestrator/async_rl_program.py

@@ -165,9 +165,9 @@ async def critique_stage(
       for idx, payload in enumerate(trainer_payloads):
         adv = payload.advantages
         reward_val = (
-            float(adv[0])
-            if hasattr(adv, "__len__") and len(adv) > 0
-            else float(adv)
+            float(adv[0])  # pyrefly: ignore[bad-index]
+            if hasattr(adv, "__len__") and len(adv) > 0  # pyrefly: ignore[bad-argument-type]
+            else float(adv)  # pyrefly: ignore[bad-argument-type]
         )
         item = datatypes.TrajectoryItem(
             pair_index=idx,
@@ -176,7 +176,7 @@ async def critique_stage(
             traj=datatypes.Trajectory(reward=reward_val),
             # TODO: Stream RLTrainerPayload directly instead of re-wrapping in TrajectoryItem.
         )
-        item.payload = payload
+        item.payload = payload  # pyrefly: ignore[missing-attribute]
         await self.scored_q.put(item)
 
   async def train_stage(
@@ -199,7 +199,7 @@ async def train_stage(
 
         payloads = [getattr(item, "payload", None) for item in scored_items]
         # TODO: Implement streaming microbatch assembly to overlap packing with trainer execution.
-        microbatches = self.assembler.pack(payloads)
+        microbatches = self.assembler.pack(payloads)  # pyrefly: ignore[bad-argument-type]
 
         is_final = group_idx == self.mini_batch_size - 1
         for batch in microbatches:
@@ -243,9 +243,9 @@ async def run_async(
         )
         for task in done:
           if task.exception():
-            raise task.exception()
+            raise task.exception()  # pyrefly: ignore[bad-raise]
       if train_task.exception():
-        raise train_task.exception()
+        raise train_task.exception()  # pyrefly: ignore[bad-raise]
     except Exception as exc:
       logging.error("Exception in StandardRLProgram execution: %s", exc)
       await self.raw_q.abort(exc)

tunix/experimental/orchestrator/batch_assembly.py

@@ -52,7 +52,7 @@ def pack(self, items: Sequence[datatypes.RLTrainerPayload]) -> list[datatypes.RL
     # Calculate token lengths from explicit fields
     item_lengths = []
     for it in items:
-      item_lengths.append(len(it.token_ids) if it.token_ids is not None else 0)
+      item_lengths.append(len(it.token_ids) if it.token_ids is not None else 0)  # pyrefly: ignore[bad-argument-type]
 
     item_list = sorted(zip(items, item_lengths), key=lambda x: x[1], reverse=True)
 

tunix/experimental/orchestrator/distributed_rl_engine.py

@@ -290,7 +290,7 @@ async def train_step(
         **kwargs,
     )
 
-  async def sync_weights(
+  async def sync_weights(  # pyrefly: ignore[bad-override]
       self,
       role: datatypes.Role = datatypes.Role.ACTOR,
       target_roles: Sequence[datatypes.Role] | None = None,

tunix/experimental/orchestrator/health_monitor.py

@@ -151,7 +151,7 @@ def overdue(self) -> list[OverdueWorker]:
     now = self._clock()
     result: list[OverdueWorker] = []
     for worker_id, (state, since) in sorted(self._state_since.items()):
-      deadline = self._deadlines.get(state)
+      deadline = self._deadlines.get(state)  # pyrefly: ignore[bad-argument-type]
       if deadline is None:
         continue
       elapsed = now - since

tunix/experimental/orchestrator/orchestrator.py

@@ -165,7 +165,7 @@ def run(
         assembler=active_assembler,
     )
     self.run_program(
-        program=active_program,
+        program=active_program,  # pyrefly: ignore[bad-argument-type]
         train_dataset=dataset,
         num_steps=num_steps,
         bring_up=False,