Replace # pytype: disable suppressions with # pyrefly: ignore

grain/_src/core/profiler.py

@@ -123,8 +123,8 @@ def _worker_init_fn() -> None:
 
 try:
   if _framework == _NO_FRAMEWORK:
-    from jax import profiler  # pylint: disable=g-import-not-at-top  # pyrefly: ignore[missing-import]
-    from jax._src.lib import jaxlib_extension_version  # pylint: disable=g-import-not-at-top  # pyrefly: ignore[missing-import]
+    from jax import profiler  # pylint: disable=g-import-not-at-top  # pytype: disable=import-error
+    from jax._src.lib import jaxlib_extension_version  # pylint: disable=g-import-not-at-top  # pytype: disable=import-error
 
     TraceAnnotation = profiler.TraceAnnotation
     is_enabled = profiler.TraceAnnotation.is_enabled

grain/_src/core/prometheus_monitoring.py

@@ -39,8 +39,8 @@
 
 try:
   # pylint: disable=g-import-not-at-top
-  import prometheus_client  # pyrefly: ignore[missing-import]
-  from prometheus_client import multiprocess  # pyrefly: ignore[missing-import]
+  import prometheus_client  # pytype: disable=import-error
+  from prometheus_client import multiprocess  # pytype: disable=import-error
 
   prometheus_client = cast(Any, prometheus_client)
   multiprocess = cast(Any, multiprocess)

grain/_src/core/sharding.py

@@ -105,7 +105,7 @@ def even_split(num_examples: int, options: ShardOptions) -> tuple[int, int]:
 
 def get_process_index_and_count():
   try:
-    import jax  # pylint:disable=g-import-not-at-top  # pyrefly: ignore[missing-import]
+    import jax  # pylint:disable=g-import-not-at-top  # pytype:disable=import-error
 
     return jax.process_index(), jax.process_count()
   except ImportError:

grain/_src/core/traceback_util.py

@@ -295,7 +295,7 @@ def _running_under_ipython() -> bool:
 
 def _ipython_supports_tracebackhide() -> bool:
   """Returns true if the IPython version supports __tracebackhide__."""
-  import IPython  # pylint: disable=g-import-not-at-top  # pyrefly: ignore[missing-import]
+  import IPython  # pylint: disable=g-import-not-at-top # pytype: disable=import-error
 
   return IPython.version_info[:2] >= (7, 17)
 

grain/_src/core/tree_lib.py

@@ -60,7 +60,7 @@ def _attrs_fields(unused_obj):
 
 
 try:
-  from jax import tree_util  # pylint: disable=g-import-not-at-top  # pyrefly: ignore[missing-import]
+  from jax import tree_util  # pylint: disable=g-import-not-at-top  # pytype: disable=import-error
 
   map_structure = tree_util.tree_map
   map_structure_with_path = tree_util.tree_map_with_path
@@ -107,7 +107,7 @@ def map_structure(f, *trees, is_leaf=None):
         new_tree.update(new_items)
         return new_tree
       else:
-        return type(a)(new_items)  # pyrefly: ignore[bad-argument-count]
+        return type(a)(new_items)  # pyrefly: ignore[bad-argument-count, bad-instantiation]
     else:  # leaf
       return f(*trees)
 
@@ -143,7 +143,7 @@ def _unflatten_iter_as(structure, flat_iter):
         new_tree.update(new_items)
         return new_tree
       else:
-        return type(structure)(new_items)  # pyrefly: ignore[bad-argument-count]
+        return type(structure)(new_items)  # pyrefly: ignore[bad-argument-count, bad-instantiation]
     else:  # leaf
       return next(flat_iter)
 
@@ -232,7 +232,7 @@ def _map_structure_with_path(prefix, f, *trees, is_leaf=None):
         new_tree.update(new_items)
         return new_tree
       else:
-        return type(a)(new_items)  # pyrefly: ignore[bad-argument-count]
+        return type(a)(new_items)  # pyrefly: ignore[bad-argument-count, bad-instantiation]
     else:  # leaf
       return f(prefix, *trees)
 
@@ -258,7 +258,7 @@ def map_structure_up_to(shallow_structure, f, structure):
       new_tree.update(new_items)
       return new_tree
     else:
-      return type(shallow_structure)(new_items)  # pyrefly: ignore[bad-argument-count]
+      return type(shallow_structure)(new_items)  # pyrefly: ignore[bad-argument-count, bad-instantiation]
   else:  # leaf
     return f(structure)
 

grain/_src/python/checkpoint/handler.py

@@ -78,7 +78,7 @@ def save(
         if `item` is None. Note that the newest Orbax V1 API does not use `args`
         and instead operates on checkpointables directly.
     """
-    item = item or args.item
+    item = item or args.item  # pytype:disable=attribute-error
     if isinstance(item, dataset.DatasetIterator):
       # ElasticIterDatasetIterator uses a custom checkpointing mechanism which
       # saves multiple files in the checkpoint directory. We should save the
@@ -137,7 +137,7 @@ def restore(
       ValueError: If the required process-specific checkpoint file does not
         exist.
     """
-    item = item or args.item
+    item = item or args.item  # pytype:disable=attribute-error
     process_index, process_count = sharding.get_process_index_and_count()
     if isinstance(
         item,
@@ -190,9 +190,9 @@ def typestr(cls):
 try:
   # Register the handler to be used with the Orbax v0 CheckpointManager API if
   # Orbax is present.
-  import orbax.checkpoint as ocp  # v0 API; pylint:disable=g-import-not-at-top  # pyrefly: ignore[missing-import]
+  import orbax.checkpoint as ocp  # v0 API; pylint:disable=g-import-not-at-top # pytype:disable=import-error
 
-  @ocp.args.register_with_handler(CheckpointHandler, for_save=True)
+  @ocp.args.register_with_handler(CheckpointHandler, for_save=True)  # pytype:disable=wrong-arg-types
   @dataclasses.dataclass
   class CheckpointSave(ocp.args.CheckpointArgs):
     """Arguments for saving a PyGrain iterator via Orbax.
@@ -236,7 +236,7 @@ class CheckpointSave(ocp.args.CheckpointArgs):
     """
     item: Any
 
-  @ocp.args.register_with_handler(CheckpointHandler, for_restore=True)
+  @ocp.args.register_with_handler(CheckpointHandler, for_restore=True)  # pytype:disable=wrong-arg-types
   @dataclasses.dataclass
   class CheckpointRestore(ocp.args.CheckpointArgs):
     """Arguments for restoring a PyGrain iterator via Orbax.

grain/_src/python/data_loader.py

@@ -451,7 +451,7 @@ def _local_num_workers(self):
   @functools.cached_property
   def _global_num_workers(self):
     """Returns the number of workers across all data shards."""
-    return self._local_num_workers * self._shard_options.shard_count  # pyrefly: ignore[missing-attribute]
+    return self._local_num_workers * self._shard_options.shard_count  # pytype: disable=attribute-error
 
   def _create_dataset(self) -> dataset.IterDataset:
     """Returns the dataset for this data loader."""

grain/_src/python/data_loader_test.py

@@ -423,7 +423,7 @@ def test_data_loader_remote_exception(self):
     )
     with self.assertRaises(Exception) as e:
       list(data_loader)
-      assert "ZeroDivisionError: division by zero" in e.__cause__._traceback  # pyrefly: ignore[missing-attribute]
+      assert "ZeroDivisionError: division by zero" in e.__cause__._traceback  # pytype: disable=attribute-error
 
   def test_data_loader_with_used_array_record_data_source(
       self,

grain/_src/python/dataset/dataset.py

@@ -914,6 +914,7 @@ def to_iter_dataset(
   def __iter__(self) -> DatasetIterator[T]:
     return self.to_iter_dataset().__iter__()
 
+  # pytype: disable=attribute-error
   # pylint: disable=protected-access
 
   def _initialize_stats(
@@ -939,7 +940,7 @@ def _initialize_stats(
     parents_stats = []
     if hasattr(self, "_parents"):
       for p in self._parents:
-        parents_stats.append(p._initialize_stats(execution_tracking_mode))  # pyrefly: ignore[missing-attribute]
+        parents_stats.append(p._initialize_stats(execution_tracking_mode))
     config = dataset_stats.StatsConfig(
         name=str(self),
         transform_mutates_spec=self._MUTATES_ELEMENT_SPEC,
@@ -966,6 +967,7 @@ def _stats(self) -> dataset_stats.Stats:
     """Returns the Stats object for recording statistics about this dataset."""
     return self._initialize_stats(base.ExecutionTrackingMode.DISABLED)
 
+  # pytype: enable=attribute-error
   # pylint: enable=protected-access
 
 
@@ -1660,6 +1662,7 @@ def _assert_not_closed(self) -> None:
     if self._closed:
       raise ValueError(f"Trying to advance a closed iterator: '{self}'.")
 
+  # pytype: disable=attribute-error
   # pylint: disable=protected-access
 
   def _initialize_stats(
@@ -1708,7 +1711,7 @@ def _initialize_stats(
   def _stats(self) -> dataset_stats.Stats:
     """Returns the Stats object for recording statistics about this dataset."""
     return self._initialize_stats(
-        self._ctx.dataset_options.execution_tracking_mode  # pyrefly: ignore[bad-argument-type]
+        self._ctx.dataset_options.execution_tracking_mode
     )
 
   def _set_next_index(self, index: int) -> None:
@@ -1737,6 +1740,7 @@ def _get_next_index(self) -> int:
 
     raise NotImplementedError
 
+  # pytype: enable=attribute-error
   # pylint: enable=protected-access
 
   def __del__(self):

grain/_src/python/dataset/dataset_test.py

@@ -265,7 +265,7 @@ def __getitem__(self, index: int) -> tuple[int, int]:
     self.assertEqual(list(ds2), [10, 11, 12])
     selected_ds = dataset.MapDataset.select_from_datasets(
         [ds1, ds2],
-        AlternatingMap(),
+        AlternatingMap(),  # pyrefly: ignore[bad-arg-type]
     )
     self.assertEqual(list(selected_ds), [0, 10, 1, 11, 2, 12])
 
@@ -932,9 +932,9 @@ def test_seed_is_different_with_chained_transforms(self):
         .seed(seed)
         .random_map(AddRandomInteger())
     )
-    map_seed1 = ds._seed  # pyrefly: ignore[missing-attribute]
+    map_seed1 = ds._seed  # pytype: disable=attribute-error
     ds = ds.random_map(AddRandomInteger())
-    map_seed2 = ds._seed  # pyrefly: ignore[missing-attribute]
+    map_seed2 = ds._seed  # pytype: disable=attribute-error
     self.assertNotEqual(map_seed1, map_seed2)
 
   def test_seed_with_shuffle_and_map(self):

grain/_src/python/dataset/sources/hf_source_test.py

@@ -182,7 +182,7 @@ def __iter__(self) -> Iterator[Any]:
         return iter(self.data)
 
     dummy_ds = UnshardableDataset(list(range(5)))
-    grain_ds = hf_source.HFIterDataset(dummy_ds)  # pyrefly: ignore[bad-argument-type]
+    grain_ds = hf_source.HFIterDataset(dummy_ds)  # pytype: disable=wrong-arg-types
     # Should not raise AttributeError when set_slice is called.
     grain_ds.set_slice(slice(1, None, 2))
     self.assertEqual(list(grain_ds), list(range(5)))

grain/_src/python/dataset/sources/parquet_dataset.py

@@ -22,7 +22,7 @@
 
 # lazy import for pyarrow
 with epy.lazy_imports():
-  import pyarrow.parquet as pq  # pylint: disable=g-import-not-at-top
+  import pyarrow.parquet as pq  # pytype: disable=import-error # pylint: disable=g-import-not-at-top
 
 
 T = TypeVar("T")

grain/_src/python/dataset/stats.py

@@ -129,11 +129,11 @@ def _get_nodes_before_prefetch(
   nodes_to_visit = [node]
   while nodes_to_visit:
     node_id = nodes_to_visit.pop()
-    node = summary.nodes[node_id]  # pyrefly: ignore[bad-assignment]
+    node = summary.nodes[node_id]
     child_nodes.append(node_id)
-    if node.is_prefetch:  # pyrefly: ignore[missing-attribute]
+    if node.is_prefetch:
       continue  # Skip adding inputs for the prefetch node
-    nodes_to_visit.extend(node.inputs)  # pyrefly: ignore[missing-attribute]
+    nodes_to_visit.extend(node.inputs)
   return child_nodes
 
 
@@ -175,7 +175,7 @@ def _compute_wait_time_ratio(
     # ratio attributed to the prefetch node is distributed among these nodes
     # proportionally to their total processing time.
     if node.is_prefetch:
-      prefetch_factor = node.wait_time_ratio  # pyrefly: ignore[bad-assignment]
+      prefetch_factor = node.wait_time_ratio
       prefetch_child_nodes = _get_nodes_before_prefetch(input_node_id, summary)
       aggregated_wait_time_ns = _find_aggregated_processing_time(
           summary, prefetch_child_nodes
@@ -764,7 +764,7 @@ def _running_in_colab() -> bool:
 class _DefaultStats(Stats):
   """Default implementation for statistics collection that does nothing."""
 
-  def record_self_time(self, *, num_elements: int = 1, offset_ns: int = 0):  # pyrefly: ignore[bad-override]
+  def record_self_time(self, *, num_elements: int = 1, offset_ns: int = 0):
     return contextlib.nullcontext()
 
   def record_output_spec(self, element: T) -> T:
@@ -792,7 +792,7 @@ def __reduce__(self):
     return _VisualizationStats, (self._config.name, self._parents)
 
   @contextlib.contextmanager
-  def record_self_time(self, *, num_elements: int = 1, offset_ns: int = 0):  # pyrefly: ignore[bad-override]
+  def record_self_time(self, *, num_elements: int = 1, offset_ns: int = 0):
     yield
 
   def record_output_spec(self, element: T) -> T:
@@ -894,7 +894,7 @@ def _logging_execution_summary_loop(self) -> None:
       if not self._is_output:
         return
       if self._last_update_time > self._last_report_time:
-        self._last_report_time = time.time()  # pyrefly: ignore[bad-assignment]
+        self._last_report_time = time.time()
         summary = self._get_execution_summary()
         summary = stats_utils.sort_nodes_by_wait_time_ratio(summary)
         msg = (
@@ -931,7 +931,9 @@ def _build_execution_summary(
     for p in self._parents:
       node_id += 1
       execution_summary.nodes[current_node_id].inputs.append(node_id)
-      _, node_id = p._build_execution_summary(execution_summary, node_id)  # pylint: disable=protected-access  # pyrefly: ignore[missing-attribute]
+      # pytype: disable=attribute-error
+      _, node_id = p._build_execution_summary(execution_summary, node_id)  # pylint: disable=protected-access
+      # pytype: enable=attribute-error
     return execution_summary, node_id
 
   def _get_execution_summary(self) -> execution_summary_pb2.ExecutionSummary:
@@ -1033,7 +1035,7 @@ def _get_merged_summary_from_workers(
     """Calculates the aggregated execution summary from all workers."""
     aggregated_summary_from_workers = execution_summary_pb2.ExecutionSummary()
     stats_in_queues = self._config.stats_in_queues
-    for worker_index, worker_queue in enumerate(stats_in_queues):  # pyrefly: ignore[not-iterable]
+    for worker_index, worker_queue in enumerate(stats_in_queues):  # pyrefly: ignore[bad-argument-type, not-iterable]
       try:
         summary_from_worker = worker_queue.get(
             timeout=_WORKER_QUEUE_TIMEOUT_SEC

grain/_src/python/dataset/stats_test.py

@@ -300,9 +300,9 @@ def test_record_stats_thread_safe(self):
     reported_self_time = 0
 
     def mock_report(node):
-      while node._self_times_buffer:
+      while node._self_times_buffer:  # pytype: disable=attribute-error
         nonlocal reported_self_time
-        reported_self_time += node._self_times_buffer.pop().duration_ns
+        reported_self_time += node._self_times_buffer.pop().duration_ns  # pytype: disable=attribute-error
       for p in node._parents:
         p.report()
 
@@ -543,7 +543,7 @@ def setUp(self):
     )
 
   def _assert_visualization(self, ds, expected):
-    result = ds._stats._visualize_dataset_graph()
+    result = ds._stats._visualize_dataset_graph()  # pytype: disable=attribute-error
     # Remove line number from the result to make test less brittle.
     result = re.sub(r".py:\d+", ".py:XXX", result)
     self.assertStartsWith(result, expected)

grain/_src/python/experimental/device_put/device_put.py

@@ -42,7 +42,7 @@ def device_put(
   """
   ds = ThreadPrefetchIterDataset(ds, prefetch_buffer_size=cpu_buffer_size)
   # May raise ImportError if jax is not linked.
-  import jax  # pylint:disable=g-import-not-at-top  # pyrefly: ignore[missing-import]
+  import jax  # pylint:disable=g-import-not-at-top  # pytype:disable=import-error
 
   ds = ds.map(lambda x: jax.device_put(x, device))
   ds = ThreadPrefetchIterDataset(ds, prefetch_buffer_size=device_buffer_size)

grain/_src/python/experimental/example_packing/pack_and_batch_test.py

@@ -45,7 +45,7 @@ def common_test_body(
       length_struct=length_struct,
       max_sequences_per_bin=max_sequences_per_bin,
   )
-  packed_dataset = pack_op(input_dataset)
+  packed_dataset = pack_op(input_dataset)  # pytype: disable=wrong-arg-types
   actual_packed_dataset = list(packed_dataset)
   np.testing.assert_equal(
       len(actual_packed_dataset), len(expected_packed_dataset)

grain/_src/python/experimental/example_packing/packing.py

@@ -216,7 +216,7 @@ def __call__(
   ) -> Iterator[record.Record[tuple[_T, _T, _T]]]:
     for element in input_iterator:
       # Use `element` to set dtypes + trailing dimensions.
-      if self._cur_batch is None:
+      if self._cur_batch is None:  # pytype: disable=attribute-error
         self._cur_batch = _PackedBatch(
             element.data,
             self.batch_size,

grain/_src/python/ipc/queue_test.py

@@ -28,7 +28,7 @@ def test_add_element_to_queue(self):
     element = 1
     termination_event = multiprocessing.Event()
     self.assertTrue(
-        multiprocessing_common.add_element_to_queue(
+        multiprocessing_common.add_element_to_queue(  # pytype: disable=wrong-arg-types
             element=element,
             elements_queue=test_queue,  # pyrefly: ignore[bad-argument-type]
             should_stop=termination_event.is_set,
@@ -42,7 +42,7 @@ def test_add_element_to_queue_already_terminated(self):
     termination_event = multiprocessing.Event()
     termination_event.set()
     self.assertFalse(
-        multiprocessing_common.add_element_to_queue(
+        multiprocessing_common.add_element_to_queue(  # pytype: disable=wrong-arg-types
             element=element,
             elements_queue=test_queue,  # pyrefly: ignore[bad-argument-type]
             should_stop=termination_event.is_set,
@@ -56,7 +56,7 @@ def test_get_element_from_queue(self):
     expected_element = 1
     test_queue.put(expected_element)
     termination_event = multiprocessing.Event()
-    actual_element = multiprocessing_common.get_element_from_queue(
+    actual_element = multiprocessing_common.get_element_from_queue(  # pytype: disable=wrong-arg-types
         elements_queue=test_queue,  # pyrefly: ignore[bad-argument-type]
         should_stop=termination_event.is_set,
     )
@@ -68,7 +68,7 @@ def test_get_element_from_queue_already_terminated(self):
     test_queue.put(expected_element)
     termination_event = multiprocessing.Event()
     termination_event.set()
-    actual_element = multiprocessing_common.get_element_from_queue(
+    actual_element = multiprocessing_common.get_element_from_queue(  # pytype: disable=wrong-arg-types
         elements_queue=test_queue,  # pyrefly: ignore[bad-argument-type]
         should_stop=termination_event.is_set,
     )