From 3384a822f45d0927e612ebc66f852b782755dc30 Mon Sep 17 00:00:00 2001
From: Grain Team <grain-dev@google.com>
Date: Wed, 7 Oct 2026 04:39:04 -0700
Subject: [PATCH] Replace `# pytype: disable` suppressions with `# pyrefly:
 ignore`

PiperOrigin-RevId: 995022497
---
 grain/_src/core/profiler.py                   |  4 +--
 grain/_src/core/prometheus_monitoring.py      |  4 +--
 grain/_src/core/sharding.py                   |  2 +-
 grain/_src/core/traceback_util.py             |  2 +-
 grain/_src/core/tree_lib.py                   | 10 +++----
 grain/_src/python/checkpoint/handler.py       | 10 +++----
 grain/_src/python/data_loader.py              |  2 +-
 grain/_src/python/data_loader_test.py         |  2 +-
 grain/_src/python/dataset/dataset.py          |  8 ++++--
 grain/_src/python/dataset/dataset_test.py     |  6 ++---
 .../python/dataset/sources/hf_source_test.py  |  2 +-
 .../python/dataset/sources/parquet_dataset.py |  2 +-
 grain/_src/python/dataset/stats.py            | 20 +++++++-------
 grain/_src/python/dataset/stats_test.py       |  6 ++---
 .../experimental/device_put/device_put.py     |  2 +-
 .../example_packing/pack_and_batch_test.py    |  2 +-
 .../experimental/example_packing/packing.py   |  2 +-
 grain/_src/python/ipc/queue_test.py           |  8 +++---
 grain/_src/python/ipc/variable_size_queue.py  | 26 ++++++++++++-------
 grain/_src/python/operations.py               |  4 +--
 grain/_src/python/samplers_test.py            |  2 +-
 21 files changed, 69 insertions(+), 57 deletions(-)

diff --git a/grain/_src/core/profiler.py b/grain/_src/core/profiler.py
index 7e4ad8052..13bdf90cc 100644
--- a/grain/_src/core/profiler.py
+++ b/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
diff --git a/grain/_src/core/prometheus_monitoring.py b/grain/_src/core/prometheus_monitoring.py
index 6157bf963..529be2426 100644
--- a/grain/_src/core/prometheus_monitoring.py
+++ b/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)
diff --git a/grain/_src/core/sharding.py b/grain/_src/core/sharding.py
index 613a7d005..ae0c55c87 100644
--- a/grain/_src/core/sharding.py
+++ b/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:
diff --git a/grain/_src/core/traceback_util.py b/grain/_src/core/traceback_util.py
index 7dea345ac..d161bd403 100644
--- a/grain/_src/core/traceback_util.py
+++ b/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)
 
diff --git a/grain/_src/core/tree_lib.py b/grain/_src/core/tree_lib.py
index e0b715da6..3030fe33f 100644
--- a/grain/_src/core/tree_lib.py
+++ b/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)
 
diff --git a/grain/_src/python/checkpoint/handler.py b/grain/_src/python/checkpoint/handler.py
index 0765c876c..864277e5d 100644
--- a/grain/_src/python/checkpoint/handler.py
+++ b/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.
diff --git a/grain/_src/python/data_loader.py b/grain/_src/python/data_loader.py
index 07890a802..25f584731 100644
--- a/grain/_src/python/data_loader.py
+++ b/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."""
diff --git a/grain/_src/python/data_loader_test.py b/grain/_src/python/data_loader_test.py
index 4632d2676..f13222583 100644
--- a/grain/_src/python/data_loader_test.py
+++ b/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,
diff --git a/grain/_src/python/dataset/dataset.py b/grain/_src/python/dataset/dataset.py
index 1bf6ca2e1..1b574ad73 100644
--- a/grain/_src/python/dataset/dataset.py
+++ b/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):
diff --git a/grain/_src/python/dataset/dataset_test.py b/grain/_src/python/dataset/dataset_test.py
index 0929b59c7..e03ecd2a6 100644
--- a/grain/_src/python/dataset/dataset_test.py
+++ b/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):
diff --git a/grain/_src/python/dataset/sources/hf_source_test.py b/grain/_src/python/dataset/sources/hf_source_test.py
index d24352306..d23b31682 100644
--- a/grain/_src/python/dataset/sources/hf_source_test.py
+++ b/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)))
diff --git a/grain/_src/python/dataset/sources/parquet_dataset.py b/grain/_src/python/dataset/sources/parquet_dataset.py
index c073a2eb6..6a43e949b 100644
--- a/grain/_src/python/dataset/sources/parquet_dataset.py
+++ b/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")
diff --git a/grain/_src/python/dataset/stats.py b/grain/_src/python/dataset/stats.py
index de6c389de..264fc6a37 100644
--- a/grain/_src/python/dataset/stats.py
+++ b/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
diff --git a/grain/_src/python/dataset/stats_test.py b/grain/_src/python/dataset/stats_test.py
index a4f57ae99..7c0d36839 100644
--- a/grain/_src/python/dataset/stats_test.py
+++ b/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)
diff --git a/grain/_src/python/experimental/device_put/device_put.py b/grain/_src/python/experimental/device_put/device_put.py
index 53c6a89d2..6872f7e53 100644
--- a/grain/_src/python/experimental/device_put/device_put.py
+++ b/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)
diff --git a/grain/_src/python/experimental/example_packing/pack_and_batch_test.py b/grain/_src/python/experimental/example_packing/pack_and_batch_test.py
index bac6a00c6..59374fc40 100644
--- a/grain/_src/python/experimental/example_packing/pack_and_batch_test.py
+++ b/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)
diff --git a/grain/_src/python/experimental/example_packing/packing.py b/grain/_src/python/experimental/example_packing/packing.py
index 10cdff819..5a303b87c 100644
--- a/grain/_src/python/experimental/example_packing/packing.py
+++ b/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,
diff --git a/grain/_src/python/ipc/queue_test.py b/grain/_src/python/ipc/queue_test.py
index 1e10c1831..107e1c93c 100644
--- a/grain/_src/python/ipc/queue_test.py
+++ b/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,
     )
diff --git a/grain/_src/python/ipc/variable_size_queue.py b/grain/_src/python/ipc/variable_size_queue.py
index df84fd603..1bfa7ae4b 100644
--- a/grain/_src/python/ipc/variable_size_queue.py
+++ b/grain/_src/python/ipc/variable_size_queue.py
@@ -37,16 +37,20 @@ def __init__(
     self._pending_shrink = ctx.Value("i", 0, lock=False)
 
   def __getstate__(self):
+    # pytype: disable=attribute-error
     context.assert_spawning(self)
     return cast(tuple, queues.Queue.__getstate__(self)) + (  # pylint: disable=g-bare-generic
         self._resize_lock,
         self._max_size_val,
         self._pending_shrink,
     )
+    # pytype: enable=attribute-error
 
   def __setstate__(self, state):
+    # pytype: disable=attribute-error
     queues.Queue.__setstate__(self, state[:-3])
     self._resize_lock, self._max_size_val, self._pending_shrink = state[-3:]
+    # pytype: enable=attribute-error
 
   def set_max_size(self, max_size: int):
     """Sets the maximum size of the queue.
@@ -78,26 +82,28 @@ def set_max_size(self, max_size: int):
 
   def get(self, block: bool = True, timeout: float | None = None):
     """Gets an item from the queue, similar to `queue.Queue.get`."""
-    if self._closed:  # pyrefly: ignore[missing-attribute]
+    # pytype: disable=attribute-error
+    if self._closed:
       raise ValueError(f"Queue {self!r} is closed")
     if block and timeout is None:
-      with self._rlock:  # pyrefly: ignore[missing-attribute]
-        res = self._recv_bytes()  # pyrefly: ignore[missing-attribute]
+      with self._rlock:
+        res = self._recv_bytes()
     else:
       if block:
-        deadline = time.time() + timeout  # pyrefly: ignore[unsupported-operation]
-      if not self._rlock.acquire(block, timeout):  # pyrefly: ignore[missing-attribute]
+        deadline = time.time() + timeout
+      if not self._rlock.acquire(block, timeout):
         raise queue.Empty
       try:
         if block:
-          timeout = deadline - time.time()  # pyrefly: ignore[unbound-name]
-          if not self._poll(timeout):  # pyrefly: ignore[missing-attribute]
+          timeout = deadline - time.time()
+          if not self._poll(timeout):
             raise queue.Empty
-        elif not self._poll():  # pyrefly: ignore[missing-attribute]
+        elif not self._poll():
           raise queue.Empty
-        res = self._recv_bytes()  # pyrefly: ignore[missing-attribute]
+        res = self._recv_bytes()
       finally:
-        self._rlock.release()  # pyrefly: ignore[missing-attribute]
+        self._rlock.release()
+    # pytype: enable=attribute-error
     with self._resize_lock:
       if self._pending_shrink.value > 0:
         self._pending_shrink.value -= 1
diff --git a/grain/_src/python/operations.py b/grain/_src/python/operations.py
index 4addd1829..7d2f95c74 100644
--- a/grain/_src/python/operations.py
+++ b/grain/_src/python/operations.py
@@ -145,7 +145,7 @@ def __post_init__(self):
   def __call__(
       self, input_iterator: Iterator[record.Record[_IN]]
   ) -> Iterator[record.Record[_OUT]]:
-    if self._display_deprecation_message:
+    if self._display_deprecation_message:  # pytype: disable=attribute-error
       logging.error(
           "Applying deprecated PyGrain BatchOperation. Please use the"
           " grain.python.Batch transformation."
@@ -161,7 +161,7 @@ def __call__(
         yield record.Record(last_record_metadata.remove_record_key(), batch)
     if records_to_batch and not self.drop_remainder:
       yield record.Record(
-          last_record_metadata.remove_record_key(),  # pyrefly: ignore[missing-attribute]
+          last_record_metadata.remove_record_key(),  # pytype: disable=attribute-error
           self.batch_fn(records_to_batch),  # pyrefly: ignore[not-callable]
       )
 
diff --git a/grain/_src/python/samplers_test.py b/grain/_src/python/samplers_test.py
index b5af58c3c..47ec5dbee 100644
--- a/grain/_src/python/samplers_test.py
+++ b/grain/_src/python/samplers_test.py
@@ -179,7 +179,7 @@ def test_invalid_non_integer_seed(self):
           shard_options=sharding.NoSharding(),
           shuffle=True,
           num_epochs=2,
-          seed=(3, 4),  # pyrefly: ignore[bad-argument-type]
+          seed=(3, 4),  # pytype: disable=wrong-arg-types
       )
 
   def test_invalid_non_int32_seed(self):