Replace # pytype: disable suppressions with # pyrefly: ignore

ffn/inference/executor.py

@@ -195,7 +195,7 @@ def _run_executor_log_exceptions(self):
       logging.exception(e)
       # If the executor fails, the whole process becomes useless and we need
       # to make sure it gets terminated.
-      _thread.interrupt_main()  # pytype: disable=module-attr
+      _thread.interrupt_main()
       time.sleep(10)
       os._exit(1)  # pylint:disable=protected-access
 
@@ -325,7 +325,7 @@ def _schedule_batch(self, client_ids: Sequence[int], fetches: Sequence[str]):
         logging.exception(e)
         # If calling TF didn't work (faulty hardware, misconfiguration, etc),
         # we want to terminate the whole program.
-        _thread.interrupt_main()  # pytype: disable=module-attr
+        _thread.interrupt_main()
         raise e
 
     with timer_counter(self.counters, 'executor-output'):
@@ -379,7 +379,7 @@ def _dispatcher(self, data: jnp.ndarray, client_ids: Sequence[int]):
         logging.exception(e)
         self._interface.exit_request.set()
         # Terminate the whole program on failure.
-        _thread.interrupt_main()  # pytype: disable=module-attr
+        _thread.interrupt_main()
         raise e
 
       ret = np.array(ret)

ffn/inference/inference.py

@@ -54,7 +54,7 @@ class DynamicImage:
 
   def UpdateFromPIL(self, new_img):
     # pylint: disable=g-import-not-at-top
-    from IPython import display  # pytype:disable=import-error
+    from IPython import display
 
     display.clear_output(wait=True)
     image = BytesIO()

ffn/jax/input_pipeline.py

@@ -125,7 +125,7 @@ def _add_ffn_data(ex: volume.Example) -> volume.Example:
       if config.loss_mask_invert:
         loss_mask = tf.equal(ex['loss_mask'], 0)
       else:
-        loss_mask = ex['loss_mask'] > 0  # pyrefly: ignore[unsupported-operation]
+        loss_mask = ex['loss_mask'] > 0
 
       weights *= 1.0 - tf.cast(loss_mask, tf.float32)
 
@@ -137,7 +137,7 @@ def _add_ffn_data(ex: volume.Example) -> volume.Example:
         seg.shape[3] // 2,  #
         0,
     ]
-    lom = tf.logical_and(seg > 0, tf.equal(seg, center_val))  # pyrefly: ignore[unsupported-operation]
+    lom = tf.logical_and(seg > 0, tf.equal(seg, center_val))
     labels = inputs.soften_labels(lom)
 
     lx, ly, lz = load_shape
@@ -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[bad-argument-type, missing-argument]
+      self._fs.add(self._tpe.submit(lambda gen=gen, i=i: (i, next(gen))))  # pyrefly: ignore[bad-argument-type]
 
     # 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[bad-argument-type, missing-argument]
+              self._tpe.submit(lambda gen=gen, i=gen_idx: (i, next(gen)))  # pyrefly: ignore[bad-argument-type]
           )
 
     # Distribute data asynchronously.

ffn/jax/train.py

@@ -54,7 +54,7 @@
 from connectomics.jax.models import util as model_util
 
 
-class TrainState(flax.struct.PyTreeNode):  # pytype: disable=invalid-function-definition  # dataclass_transform
+class TrainState(flax.struct.PyTreeNode):
   step: int
   opt_state: optax.OptState
   params: flax.core.FrozenDict[str, Any]
@@ -235,15 +235,15 @@ def loss_fn(params):
           new_params, new_ema_params, step_size=1.0 - decay
       )
 
-  new_state = state.replace(  # pytype: disable=attribute-error
+  new_state = state.replace(
       step=step,
       params=new_params,
       opt_state=new_opt_state,
       batch_stats=new_batch_stats,
       ema_params=new_ema_params,
   )
 
-  lr = schedule(state.opt_state.count)  # pytype: disable=attribute-error
+  lr = schedule(state.opt_state.count)  # pyrefly: ignore[missing-attribute]
   metrics_update = TrainMetrics.single_from_model_output(
       loss=loss, learning_rate=lr
   )
@@ -407,9 +407,9 @@ def train_and_evaluate(
       load_shape=tuple(train_image_size(info, config)),
       data_service_address=data_service_address,
   )
-  train_iter = iter(train_ds)  # pytype: disable=wrong-arg-types
+  train_iter = iter(train_ds)
 
-  logging.info('train_elem_shape=%r', train_iter.element_spec['em'].shape)  # pytype:disable=attribute-error
+  logging.info('train_elem_shape=%r', train_iter.element_spec['em'].shape)
 
   # batch, z, y, x, (image, seed)
   input_shape = [1] + np.array(info.input_image_size).tolist()[::-1] + [2]
@@ -469,7 +469,7 @@ def train_and_evaluate(
   )
   checkpointed_state = {'train_state': state}
   if 'train_iter' in checkpoint_items:
-    checkpointed_state['train_iter'] = train_iter  # pyrefly: ignore[bad-assignment]
+    checkpointed_state['train_iter'] = train_iter
   latest_step = checkpoint_manager.latest_step()
   # If an initial checkpoint is provided and the checkpointing library does not
   # report a 'latest' checkpoint, then we are starting a new experiment.
@@ -498,7 +498,7 @@ def train_and_evaluate(
         'train_state': ocp.args.StandardRestore(state),
     }
     if 'train_iter' in checkpoint_items:
-      restore_args['train_iter'] = _get_ocp_args(train_iter)  # pyrefly: ignore[bad-assignment, bad-specialization]
+      restore_args['train_iter'] = _get_ocp_args(train_iter)  # pyrefly: ignore[bad-assignment]
     checkpointed_state = checkpoint_manager.restore(
         latest_step,
         args=ocp.args.Composite(**restore_args),

ffn/secgan/train.py

@@ -1314,7 +1314,7 @@ def create_input_pipeline(
       image_b_volinfo, bbox_b, patch_size_xyz, rng_b, config
   )
 
-  ds = tf.data.Dataset.zip((ds_a, ds_b))  # pyrefly: ignore[bad-argument-type]
+  ds = tf.data.Dataset.zip((ds_a, ds_b))
   ds = ds.prefetch(128)
   return ds
 

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, not-iterable]
+  for axis, ix in enumerate(center):  # pyrefly: ignore[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, not-iterable]
+      for ax, c in enumerate(center):  # pyrefly: ignore[not-iterable]
         if ax != axis:
           # Make axis i the 0-axis an work in-place.
           view = np.rollaxis(plane, i)