Replace # pytype: disable suppressions with # pyrefly: ignore

recml/core/metrics/base_metrics.py

@@ -31,7 +31,7 @@
 
 # TODO(b/387463777): Consider removing the dependency on CLU metrics longer term
 # since it's just an interface.
-@dataclass_transform(field_specifiers=(struct.field,))  # pytype: disable=not-supported-yet
+@dataclass_transform(field_specifiers=(struct.field,))
 class Metric(abc.ABC, clu_metrics.Metric, struct.PyTreeNode):
   """PyTree node representing the state of a metric.
 
@@ -63,7 +63,7 @@ def from_model_output(cls, *args, **kwargs) -> Self:
     raise NotImplementedError()
 
   @abc.abstractmethod
-  def merge(self, other: Self) -> Self:  # pytype: disable=signature-mismatch
+  def merge(self, other: Self) -> Self:
     """Merges two metrics.
 
     Args:
@@ -79,7 +79,7 @@ def merge(self, other: Self) -> Self:  # pytype: disable=signature-mismatch
     """
 
   @abc.abstractmethod
-  def compute(self) -> Scalar:  # pytype: disable=signature-mismatch
+  def compute(self) -> Scalar:  # pyrefly: ignore[bad-override]
     """Computes the value of the metric.
 
     NOTE: This method is almost always called on the host which means that it

recml/core/metrics/confusion_metrics.py

@@ -67,8 +67,8 @@ def from_model_output(
         false_negatives=fn,
     )
 
-  def merge(self, other: Self) -> Self:  # pytype: disable=signature-mismatch
-    return type(self)(  # pytype: disable=not-instantiable
+  def merge(self, other: Self) -> Self:
+    return type(self)(
         true_positives=self.true_positives + other.true_positives,
         true_negatives=self.true_negatives + other.true_negatives,
         false_positives=self.false_positives + other.false_positives,
@@ -237,7 +237,7 @@ def compute(self) -> float:
         self.false_positives, self.false_positives + self.true_negatives
     )
     # We negate the integral because the thresholds are in ascending order.
-    return -np.trapezoid(tp_rate, fp_rate)  # pyrefly: ignore[bad-return]
+    return -np.trapezoid(tp_rate, fp_rate)
 
 
 class PrecisionAtRecall(ConfusionMetric):

recml/core/metrics/reduction_metrics.py

@@ -75,7 +75,7 @@ def from_model_output(
 
     return cls(total=total)
 
-  def merge(self, other: Self) -> Self:  # pytype: disable=signature-mismatch
+  def merge(self, other: Self) -> Self:
     return type(self)(total=self.total + other.total)
 
   def compute(self) -> base_metrics.Scalar:
@@ -111,7 +111,7 @@ def from_model_output(
 
     return cls(total=total, count=count)
 
-  def merge(self, other: Self) -> Self:  # pytype: disable=signature-mismatch
+  def merge(self, other: Self) -> Self:
     return type(self)(
         total=self.total + other.total,
         count=self.count + other.count,

recml/core/ops/binary_cross_entropy_ops.py

@@ -1610,7 +1610,7 @@ def _bce_bwd_sharded(
   vocab_axis_name = emb_spec[0] if emb_spec else None
 
   dp_axes = []
-  for axis in local_act_spec:  # pyrefly: ignore[not-iterable]
+  for axis in local_act_spec:
     if axis is not None and axis != vocab_axis_name:
       dp_axes.append(axis)
 
@@ -1646,13 +1646,13 @@ def _bwd_local_with_reduction(d_loss_, act_, emb_, tgt_):
   if targets.ndim == replicated_activations.ndim:
     local_tgt_spec = local_act_spec
   else:
-    local_tgt_spec = jax.sharding.PartitionSpec(*local_act_spec[1:])  # pyrefly: ignore[unsupported-operation]
+    local_tgt_spec = jax.sharding.PartitionSpec(*local_act_spec[1:])
 
   d_act_replicated, d_emb = jax.shard_map(
       _bwd_local_with_reduction,
       mesh=mesh,
       in_specs=(
-          jax.sharding.PartitionSpec(*local_act_spec[:-1]),  # pyrefly: ignore[unsupported-operation]
+          jax.sharding.PartitionSpec(*local_act_spec[:-1]),
           local_act_spec,
           emb_spec,
           local_tgt_spec,

recml/core/ops/binary_focal_cross_entropy_ops.py

@@ -1187,7 +1187,7 @@ def _focal_bce_bwd_sharded(
   vocab_axis_name = emb_spec[0] if emb_spec else None
 
   dp_axes = []
-  for axis in local_act_spec:  # pyrefly: ignore[not-iterable]
+  for axis in local_act_spec:
     if axis is not None and axis != vocab_axis_name:
       dp_axes.append(axis)
 
@@ -1223,13 +1223,13 @@ def _bwd_local_with_reduction(d_loss_, act_, emb_, tgt_):
   if targets.ndim == replicated_activations.ndim:
     local_tgt_spec = local_act_spec
   else:
-    local_tgt_spec = jax.sharding.PartitionSpec(*local_act_spec[1:])  # pyrefly: ignore[unsupported-operation]
+    local_tgt_spec = jax.sharding.PartitionSpec(*local_act_spec[1:])
 
   d_act_replicated, d_emb = jax.shard_map(
       _bwd_local_with_reduction,
       mesh=mesh,
       in_specs=(
-          jax.sharding.PartitionSpec(*local_act_spec[:-1]),  # pyrefly: ignore[unsupported-operation]
+          jax.sharding.PartitionSpec(*local_act_spec[:-1]),
           local_act_spec,
           emb_spec,
           local_tgt_spec,

recml/core/ops/hstu_ops.py

@@ -160,7 +160,7 @@ def _apply_mask(
       q_sequence = jnp.broadcast_to(q_sequence, (k_slice.size, bq))
 
     assert q_sequence.shape == k_sequence.shape
-    computed_mask = mask_function(q_sequence, k_sequence)  # pytype: disable=wrong-arg-count
+    computed_mask = mask_function(q_sequence, k_sequence)
     if computed_mask.dtype != jnp.dtype(jnp.bool_):
       raise ValueError(
           "Mask function must return a boolean-valued array, but got:"
@@ -244,7 +244,7 @@ def body(kv_compute_index, _):
         q_sequence_ref=q_sequence_ref,
         q_segment_ids_ref=q_segment_ids_ref,
         kv_segment_ids_ref=kv_segment_ids_ref,
-        k_slice=slice_k,  # pyrefly: ignore[bad-argument-type]
+        k_slice=slice_k,
         # When the iteration space is shrunk (for local attention for example),
         # the kv_index program_id does not correspond to the actual coordinates
         # of the KV data. Make sure to use the 'unshrunk' index (coming from the
@@ -479,7 +479,7 @@ def body(i, _):
         q_sequence_ref,
         q_segment_ids_ref,
         kv_segment_ids_ref,
-        k_slice=slice_k,  # pyrefly: ignore[bad-argument-type]
+        k_slice=slice_k,
         k_offset=j * bkv + i * bkv_compute,
         bq=bq,
         k_in_lanes=False,

recml/core/training/core.py

@@ -205,7 +205,7 @@ def get_iterators(
         f" but got {eval_datasets}. You cannot mix both."
     )
 
-  return train_dataset, eval_datasets  # pytype: disable=bad-return-type
+  return train_dataset, eval_datasets  # pyrefly: ignore[bad-return]
 
 
 def get_shape(

recml/examples/dlrm_experiment.py

@@ -215,7 +215,7 @@ def make(self) -> tf.data.Dataset:
 
     label = np.random.randint(0, 2, size=(batch_size,))
 
-    dataset = tf.data.Dataset.from_tensors((data, label))  # pyrefly: ignore[bad-argument-type]
+    dataset = tf.data.Dataset.from_tensors((data, label))
     dataset = dataset.take(1).repeat()
     dataset = dataset.prefetch(buffer_size=2048)
     options = tf.data.Options()