From 6c3cf1bf1da48f89aac633324a8e0a7d64d33fec Mon Sep 17 00:00:00 2001 From: Hana Joo Date: Wed, 8 Jul 2026 18:25:48 -0700 Subject: [PATCH] Adding type suppressions for pyrefly PiperOrigin-RevId: 944813888 --- src/metrax/audio_metrics.py | 4 ++-- src/metrax/classification_metrics.py | 2 +- src/metrax/image_metrics.py | 16 ++++++++-------- src/metrax/nlp_metrics.py | 8 ++++---- src/metrax/ranking_metrics.py | 8 ++++---- src/metrax/regression_metrics.py | 6 +++--- 6 files changed, 22 insertions(+), 22 deletions(-) diff --git a/src/metrax/audio_metrics.py b/src/metrax/audio_metrics.py index c8769c4..2bfeaa9 100644 --- a/src/metrax/audio_metrics.py +++ b/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] diff --git a/src/metrax/classification_metrics.py b/src/metrax/classification_metrics.py index cf3a070..4e8b4d4 100644 --- a/src/metrax/classification_metrics.py +++ b/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, diff --git a/src/metrax/image_metrics.py b/src/metrax/image_metrics.py index e6f3479..b08567e 100644 --- a/src/metrax/image_metrics.py +++ b/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] diff --git a/src/metrax/nlp_metrics.py b/src/metrax/nlp_metrics.py index 4ce504c..7a99b70 100644 --- a/src/metrax/nlp_metrics.py +++ b/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) diff --git a/src/metrax/ranking_metrics.py b/src/metrax/ranking_metrics.py index 4f37125..4e0c2df 100644 --- a/src/metrax/ranking_metrics.py +++ b/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, diff --git a/src/metrax/regression_metrics.py b/src/metrax/regression_metrics.py index a75e921..fff944e 100644 --- a/src/metrax/regression_metrics.py +++ b/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,