Merge pull request #139 from eamag:main

PiperOrigin-RevId: 859822315

Author: copybara-github

.github/workflows/ci.yml

@@ -4,15 +4,15 @@ jobs:
   ruff:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v6
       - name: Lint
         uses: astral-sh/ruff-action@v2
   test:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v6
       - name: Set up Python 3.12
-        uses: actions/setup-python@v5
+        uses: actions/setup-python@v6
         with:
           python-version: 3.12
       - name: Install dependencies

src/metrax/__init__.py

@@ -45,6 +45,7 @@
 RougeL = nlp_metrics.RougeL
 RougeN = nlp_metrics.RougeN
 SNR = audio_metrics.SNR
+SpearmanRankCorrelation = regression_metrics.SpearmanRankCorrelation
 SSIM = image_metrics.SSIM
 WER = nlp_metrics.WER
 
@@ -70,6 +71,7 @@
     "PSNR",
     "RMSE",
     "RSQUARED",
+    "SpearmanRankCorrelation",
     "Recall",
     "RecallAtK",
     "RougeL",

src/metrax/nnx/__init__.py

@@ -39,6 +39,7 @@
 RougeL = nnx_metrics.RougeL
 RougeN = nnx_metrics.RougeN
 SNR = nnx_metrics.SNR
+SpearmanRankCorrelation = nnx_metrics.SpearmanRankCorrelation
 SSIM = nnx_metrics.SSIM
 WER = nnx_metrics.WER
 
@@ -68,6 +69,7 @@
     "RougeL",
     "RougeN",
     "SNR",
+    "SpearmanRankCorrelation",
     "SSIM",
     "WER",
 ]

src/metrax/nnx/nnx_metrics.py

@@ -191,6 +191,13 @@ def __init__(self):
     super().__init__(metrax.SNR)
 
 
+class SpearmanRankCorrelation(NnxWrapper):
+  """An NNX class for the Metrax metric SpearmanRankCorrelation."""
+
+  def __init__(self):
+    super().__init__(metrax.SpearmanRankCorrelation)
+
+
 class SSIM(NnxWrapper):
   """An NNX class for the Metrax metric SSIM."""
 

src/metrax/regression_metrics.py

@@ -275,3 +275,78 @@ def compute(self) -> jax.Array:
     mean = base.divide_no_nan(self.total, self.count)
     sst = self.sum_of_squared_label - self.count * jnp.power(mean, 2)
     return 1 - base.divide_no_nan(self.sum_of_squared_error, sst)
+
+
+@flax.struct.dataclass
+class SpearmanRankCorrelation(clu_metrics.Metric):
+  r"""Computes the Spearman rank correlation coefficient.
+
+  The Spearman rank correlation coefficient measures the monotonic relationship
+  between two variables. It is defined as the Pearson correlation coefficient
+  between the ranked variables.
+
+  .. math::
+      \rho = 1 - \frac{6 \sum d_i^2}{n(n^2 - 1)}
+
+  where:
+      - :math:`d_i` is the difference between the ranks of each observation
+      - :math:`n` is the number of observations
+
+  This implementation accumulates all `predictions` and `labels` to compute the
+  exact ranks upon calling `compute()`.
+
+  .. warning::
+      For very large datasets, this may lead to Out-of-Memory (OOM) errors.
+
+  Attributes:
+    predictions: Accumulated predictions.
+    labels: Accumulated labels.
+  """
+
+  predictions: jax.Array
+  labels: jax.Array
+
+  @classmethod
+  def empty(cls) -> 'SpearmanRankCorrelation':
+    return cls(
+        predictions=jnp.array([], jnp.float32),
+        labels=jnp.array([], jnp.float32),
+    )
+
+  @classmethod
+  def from_model_output(
+      cls,
+      predictions: jax.Array,
+      labels: jax.Array,
+      **kwargs,
+  ) -> 'SpearmanRankCorrelation':
+    del kwargs
+    return cls(
+        predictions=predictions.flatten(),
+        labels=labels.flatten(),
+    )
+
+  def merge(
+      self, other: 'SpearmanRankCorrelation'
+  ) -> 'SpearmanRankCorrelation':
+    return type(self)(
+        predictions=jnp.concatenate([self.predictions, other.predictions]),
+        labels=jnp.concatenate([self.labels, other.labels]),
+    )
+
+  def compute(self) -> jax.Array:
+    if self.predictions.size == 0:
+      return jnp.array(jnp.nan, jnp.float32)
+
+    rank_preds = jax.scipy.stats.rankdata(self.predictions)
+    rank_labels = jax.scipy.stats.rankdata(self.labels)
+
+    def pearson_correlation(x, y):
+      mu_x = jnp.mean(x)
+      mu_y = jnp.mean(y)
+      xm, ym = x - mu_x, y - mu_y
+      r_num = jnp.sum(xm * ym)
+      r_den = jnp.sqrt(jnp.sum(xm**2) * jnp.sum(ym**2))
+      return base.divide_no_nan(r_num, r_den)
+
+    return pearson_correlation(rank_preds, rank_labels)

src/metrax/regression_metrics_test.py

@@ -24,8 +24,11 @@
 import keras
 import metrax
 import numpy as np
+import scipy
 from sklearn import metrics as sklearn_metrics
 
+spearmanr = scipy.stats.spearmanr
+
 np.random.seed(42)
 BATCHES = 4
 BATCH_SIZE = 8
@@ -321,6 +324,39 @@ def test_rsquared(self, y_true, y_pred, sample_weights):
         atol=atol,
     )
 
+  @parameterized.named_parameters(
+      ('basic_f16', OUTPUT_LABELS, OUTPUT_PREDS_F16),
+      ('basic_f32', OUTPUT_LABELS, OUTPUT_PREDS_F32),
+      ('basic_bf16', OUTPUT_LABELS, OUTPUT_PREDS_BF16),
+      ('batch_size_one', OUTPUT_LABELS_BS1, OUTPUT_PREDS_BS1),
+  )
+  def test_spearman(self, y_true, y_pred):
+    """Test that `SpearmanRankCorrelation` Metric computes correct values."""
+    y_true = y_true.astype(y_pred.dtype)
+    y_pred = y_pred.astype(y_true.dtype)
+
+    metric = None
+    for labels, logits in zip(y_true, y_pred):
+      update = metrax.SpearmanRankCorrelation.from_model_output(
+          predictions=logits,
+          labels=labels,
+      )
+      metric = update if metric is None else metric.merge(update)
+
+    expected, _ = spearmanr(
+        y_true.astype('float32').flatten(),
+        y_pred.astype('float32').flatten(),
+    )
+    # Use lower tolerance for lower precision dtypes.
+    rtol = 1e-2 if y_true.dtype in (jnp.float16, jnp.bfloat16) else 1e-05
+    atol = 1e-2 if y_true.dtype in (jnp.float16, jnp.bfloat16) else 1e-05
+    np.testing.assert_allclose(
+        metric.compute(),
+        expected,
+        rtol=rtol,
+        atol=atol,
+    )
+
 
 if __name__ == '__main__':
   os.environ['XLA_FLAGS'] = (