Adding type suppressions for pyrefly

src/metrax/audio_metrics.py

@@ -103,7 +103,7 @@ def _calculate_snr(
     return snr
 
   @classmethod
-  def from_model_output(
+  def from_model_output(  # pyrefly: ignore[bad-override]
       cls,
       predictions: jax.Array,
       targets: jax.Array,
@@ -126,4 +126,4 @@ def from_model_output(
         targets,
         zero_mean=zero_mean,
     )
-    return super().from_model_output(values=batch_snr_value)
+    return super().from_model_output(values=batch_snr_value)  # pyrefly: ignore[bad-return]

src/metrax/classification_metrics.py

@@ -64,7 +64,7 @@ class Accuracy(base.Average):
   """
 
   @classmethod
-  def from_model_output(
+  def from_model_output(  # pyrefly: ignore[bad-override]
       cls,
       predictions: jax.Array,
       labels: jax.Array,

src/metrax/image_metrics.py

@@ -315,7 +315,7 @@ def _calculate_ssim_for_channel(x_ch, y_ch, conv_kernel, c1, c2):
     return jnp.mean(ssim_scores_stacked, axis=-1)  # (batch,)
 
   @classmethod
-  def from_model_output(
+  def from_model_output(  # pyrefly: ignore[bad-override]
       cls,
       predictions: jax.Array,
       targets: jax.Array,
@@ -360,7 +360,7 @@ def from_model_output(
         k1=k1,
         k2=k2,
     )
-    return super().from_model_output(values=batch_ssim_values)
+    return super().from_model_output(values=batch_ssim_values)  # pyrefly: ignore[bad-return]
 
 
 @flax.struct.dataclass
@@ -438,7 +438,7 @@ def _calculate_iou_for_single_class(
     return jnp.mean(iou_scores_per_class)
 
   @classmethod
-  def from_model_output(
+  def from_model_output(  # pyrefly: ignore[bad-override]
       cls,
       predictions: jax.Array,
       targets: jax.Array,
@@ -508,7 +508,7 @@ def from_model_output(
         target_class_ids=target_class_ids,
         epsilon=epsilon,
     )
-    return super().from_model_output(values=iou_score)
+    return super().from_model_output(values=iou_score)  # pyrefly: ignore[bad-return]
 
 
 @flax.struct.dataclass
@@ -570,7 +570,7 @@ def _calculate_psnr(
     return psnr
 
   @classmethod
-  def from_model_output(
+  def from_model_output(  # pyrefly: ignore[bad-override]
       cls,
       predictions: jnp.ndarray,
       targets: jnp.ndarray,
@@ -590,7 +590,7 @@ def from_model_output(
         A ``PSNR`` instance containing per‑image PSNR values.
     """
     batch_psnr = cls._calculate_psnr(predictions, targets, max_val=max_val)
-    return super().from_model_output(values=batch_psnr)
+    return super().from_model_output(values=batch_psnr)  # pyrefly: ignore[bad-return]
 
 
 @flax.struct.dataclass
@@ -681,7 +681,7 @@ class CosineSimilarity(base.Average):
   """
 
   @classmethod
-  def from_model_output(
+  def from_model_output(  # pyrefly: ignore[bad-override]
       cls,
       predictions: jax.Array,
       targets: jax.Array,
@@ -705,4 +705,4 @@ def from_model_output(
 
     cosine_similarity = dot_product / (predictions_norm * targets_norm)
 
-    return super().from_model_output(values=cosine_similarity)
+    return super().from_model_output(values=cosine_similarity)  # pyrefly: ignore[bad-return]

src/metrax/nlp_metrics.py

@@ -186,7 +186,7 @@ def merge(self, other: 'BLEU') -> 'BLEU':
   def compute(self) -> jax.Array:
     precisions = [0] * self.max_order
     for i in range(0, self.max_order):
-      precisions[i] = base.divide_no_nan(
+      precisions[i] = base.divide_no_nan(  # pyrefly: ignore[unsupported-operation]
           self.matches_by_order[i], self.possible_matches_by_order[i]
       )
     geo_mean = (
@@ -402,7 +402,7 @@ def from_model_output(
         'total_f1': jnp.array(total_f1, dtype=jnp.float32),
         'num_examples': jnp.array(num_examples, dtype=jnp.float32),
     }
-    constructor_args.update(
+    constructor_args.update(  # pyrefly: ignore[no-matching-overload]
         cls._get_specific_constructor_args_for_class(**kwargs)
     )
     return cls(**constructor_args)
@@ -423,7 +423,7 @@ def merge(self, other: 'RougeBase') -> 'RougeBase':
         'total_f1': self.total_f1 + other.total_f1,
         'num_examples': self.num_examples + other.num_examples,
     }
-    merged_data.update(self._get_specific_fields_for_merge_constructor())
+    merged_data.update(self._get_specific_fields_for_merge_constructor())  # pyrefly: ignore[no-matching-overload]
     return type(self)(**merged_data)  # type: ignore[call-arg]
 
   def _validate_merge_specifics(self, other: 'RougeBase'):
@@ -569,7 +569,7 @@ class RougeN(RougeBase):
   order: int
 
   @classmethod
-  def empty(cls, order: int = 2) -> 'RougeN':
+  def empty(cls, order: int = 2) -> 'RougeN':  # pyrefly: ignore[bad-override]
     common_values = super()._get_common_initial_values()
     return cls(order=order, **common_values)
 

src/metrax/ranking_metrics.py

@@ -106,7 +106,7 @@ def _compute_dcg_at_k(k, current_item_contributions, current_score_ranks):
     return dcg_at_ks
 
   @classmethod
-  def from_model_output(
+  def from_model_output(  # pyrefly: ignore[bad-override]
       cls,
       predictions: jax.Array,
       labels: jax.Array,
@@ -245,7 +245,7 @@ def compute_ap_at_k_single(relevant_labels, total_relevant, ks):
     return ap_at_ks
 
   @classmethod
-  def from_model_output(
+  def from_model_output(  # pyrefly: ignore[bad-override]
       cls,
       predictions: jax.Array,
       labels: jax.Array,
@@ -344,7 +344,7 @@ def _compute_rr_at_k(k, reciprocal_rank, first_relevant_rank):
     return rr_at_ks
 
   @classmethod
-  def from_model_output(
+  def from_model_output(  # pyrefly: ignore[bad-override]
       cls,
       predictions: jax.Array,
       labels: jax.Array,
@@ -443,7 +443,7 @@ def _calculate_metric_at_ks(
     raise NotImplementedError('Subclasses must implement this method.')
 
   @classmethod
-  def from_model_output(
+  def from_model_output(  # pyrefly: ignore[bad-override]
       cls,
       predictions: jax.Array,
       labels: jax.Array,

src/metrax/regression_metrics.py

@@ -44,7 +44,7 @@ class MAE(base.Average):
   """
 
   @classmethod
-  def from_model_output(
+  def from_model_output(  # pyrefly: ignore[bad-override]
       cls,
       predictions: jax.Array,
       labels: jax.Array,
@@ -100,7 +100,7 @@ class MSE(base.Average):
   """
 
   @classmethod
-  def from_model_output(
+  def from_model_output(  # pyrefly: ignore[bad-override]
       cls,
       predictions: jax.Array,
       labels: jax.Array,
@@ -176,7 +176,7 @@ class MSLE(base.Average):
   """
 
   @classmethod
-  def from_model_output(
+  def from_model_output(  # pyrefly: ignore[bad-override]
       cls,
       predictions: jax.Array,
       labels: jax.Array,