Suppress new Pyrefly findings before upgrading to v1.2.0

PiperOrigin-RevId: 978603323

Author: oprypin

ffn/jax/input_pipeline.py

@@ -270,7 +270,7 @@ def __init__(
     self._fs_lock = threading.Lock()
     self._fs = set()
     for i, gen in enumerate(self._generators):
-      self._fs.add(self._tpe.submit(lambda gen=gen, i=i: (i, next(gen))))  # pyrefly: ignore[missing-argument]
+      self._fs.add(self._tpe.submit(lambda gen=gen, i=i: (i, next(gen))))  # pyrefly: ignore[bad-argument-type, missing-argument]
 
     # Prefetching of complete batches.
     self._batch_tpe = futures.ThreadPoolExecutor(max_workers=batch_prefetch)
@@ -378,7 +378,7 @@ def _update(
         for gen_idx in current:
           gen = self._generators[gen_idx]
           self._fs.add(
-              self._tpe.submit(lambda gen=gen, i=gen_idx: (i, next(gen)))  # pyrefly: ignore[missing-argument]
+              self._tpe.submit(lambda gen=gen, i=gen_idx: (i, next(gen)))  # pyrefly: ignore[bad-argument-type, missing-argument]
           )
 
     # Distribute data asynchronously.

ffn/utils/ortho_plane_visualization.py

@@ -49,15 +49,15 @@ def cut_ortho_planes(
 
   planes = []
   full_slice = [slice(None)] * 3
-  for axis, ix in enumerate(center):  # pyrefly: ignore[bad-argument-type]
+  for axis, ix in enumerate(center):  # pyrefly: ignore[bad-argument-type, not-iterable]
     cut_slice = list(full_slice)
     cut_slice[axis] = ix  # pyrefly: ignore[unsupported-operation]
     planes.append(vol[tuple(cut_slice)])
     if cross_hair:
       # Copy because cross hair is written into array data.
       plane = planes[-1].copy()
       i = 0
-      for ax, c in enumerate(center):  # pyrefly: ignore[bad-argument-type]
+      for ax, c in enumerate(center):  # pyrefly: ignore[bad-argument-type, not-iterable]
         if ax != axis:
           # Make axis i the 0-axis an work in-place.
           view = np.rollaxis(plane, i)