add base_metrics to metrax

Author: jshin1394

src/metrax/__init__.py

@@ -12,6 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from metrax.base_metrics import (
+    Average,
+)
 from metrax.classification_metrics import (
     AUCPR,
     AUCROC,
@@ -34,6 +37,7 @@
 __all__ = [
     "AUCPR",
     "AUCROC",
+    "Average",
     "AveragePrecisionAtK",
     "MSE",
     "Perplexity",

src/metrax/base_metrics.py

@@ -0,0 +1,60 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""A collection of base metrics for metrax."""
+
+from clu import metrics as clu_metrics
+import flax
+import jax
+import jax.numpy as jnp
+
+
+def divide_no_nan(x: jax.Array, y: jax.Array) -> jax.Array:
+  """Computes a safe divide which returns 0 if the y is zero."""
+  return jnp.where(y != 0, jnp.divide(x, y), 0.0)
+
+
+@flax.struct.dataclass
+class Average(clu_metrics.Average):
+  r"""Average Metric inherits clu.metrics.Average and performs safe division."""
+
+  @classmethod
+  def from_model_output(
+      cls,
+      values: jax.Array,
+      sample_weights: jax.Array | None = None,
+  ) -> 'Average':
+    """Updates the metric.
+
+    Args:
+      values: A floating point 1D vector representing the values. The shape
+      should be (batch_size,).
+      sample_weights: An optional floating point 1D vector representing the
+        weight of each sample. The shape should be (batch_size,).
+
+    Returns:
+      Updated Average metric.
+    """
+    total = values
+    count = jnp.ones_like(values, dtype=values.dtype)
+    if sample_weights is not None:
+      total = values * sample_weights
+      count = count * sample_weights
+    return cls(
+        total=total.sum(),
+        count=count.sum(),
+    )
+
+  def compute(self) -> jax.Array:
+    return divide_no_nan(self.total, self.count)

src/metrax/base_metrics_test.py

@@ -0,0 +1,140 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for metrax base metrics."""
+
+from absl.testing import absltest
+from absl.testing import parameterized
+import jax.numpy as jnp
+import keras
+import metrax
+from metrax import base_metrics
+import numpy as np
+
+np.random.seed(42)
+BATCHES = 4
+BATCH_SIZE = 8
+OUTPUT = np.random.uniform(size=(BATCHES, BATCH_SIZE))
+OUTPUT_F16 = OUTPUT.astype(jnp.float16)
+OUTPUT_F32 = OUTPUT.astype(jnp.float32)
+OUTPUT_BF16 = OUTPUT.astype(jnp.bfloat16)
+OUTPUT_BS1 = np.random.uniform(size=(BATCHES, 1)).astype(jnp.float32)
+SAMPLE_WEIGHTS = np.tile(
+    [0.5, 1, 0, 0, 0, 0, 0, 0],
+    (BATCHES, 1),
+)
+
+
+class BaseMetricsTest(parameterized.TestCase):
+
+  def test_basic_division(self):
+    x = jnp.array([10.0, 20.0, 30.0])
+    y = jnp.array([2.0, 4.0, 5.0])
+    expected = jnp.array([5.0, 5.0, 6.0])
+    result = base_metrics.divide_no_nan(x, y)
+    self.assertTrue(jnp.array_equal(result, expected))
+
+  def test_division_by_zero(self):
+    x = jnp.array([10.0, 20.0, 30.0])
+    y = jnp.array([2.0, 0.0, 5.0])
+    expected = jnp.array([5.0, 0.0, 6.0])
+    result = base_metrics.divide_no_nan(x, y)
+    self.assertTrue(jnp.array_equal(result, expected))
+
+  def test_all_zeros_denominator(self):
+    x = jnp.array([10.0, 20.0, 30.0])
+    y = jnp.array([0.0, 0.0, 0.0])
+    expected = jnp.array([0.0, 0.0, 0.0])
+    result = base_metrics.divide_no_nan(x, y)
+    self.assertTrue(jnp.array_equal(result, expected))
+
+  def test_all_zeros_numerator(self):
+    x = jnp.array([0.0, 0.0, 0.0])
+    y = jnp.array([2.0, 4.0, 5.0])
+    expected = jnp.array([0.0, 0.0, 0.0])
+    result = base_metrics.divide_no_nan(x, y)
+    self.assertTrue(jnp.array_equal(result, expected))
+
+  def test_mixed_zeros(self):
+    x = jnp.array([10.0, 0.0, 30.0, 0.0])
+    y = jnp.array([2.0, 0.0, 5.0, 4.0])
+    expected = jnp.array([5.0, 0.0, 6.0, 0.0])
+    result = base_metrics.divide_no_nan(x, y)
+    self.assertTrue(jnp.array_equal(result, expected))
+
+  def test_scalar_inputs(self):
+    x = jnp.array(10.0)
+    y = jnp.array(2.0)
+    expected = jnp.array(5.0)
+    result = base_metrics.divide_no_nan(x, y)
+    self.assertTrue(jnp.array_equal(result, expected))
+
+  def test_scalar_denominator_zero(self):
+    x = jnp.array(10.0)
+    y = jnp.array(0.0)
+    expected = jnp.array(0.0)
+    result = base_metrics.divide_no_nan(x, y)
+    self.assertTrue(jnp.array_equal(result, expected))
+
+  def test_negative_values(self):
+    x = jnp.array([-10.0, 20.0, -30.0])
+    y = jnp.array([2.0, -4.0, 5.0])
+    expected = jnp.array([-5.0, -5.0, -6.0])
+    result = base_metrics.divide_no_nan(x, y)
+    self.assertTrue(jnp.array_equal(result, expected))
+
+  def test_negative_and_zero_values(self):
+    x = jnp.array([-10.0, 20.0, -30.0, 10.0])
+    y = jnp.array([2.0, -4.0, 0.0, 0.0])
+    expected = jnp.array([-5.0, -5.0, 0.0, 0.0])
+    result = base_metrics.divide_no_nan(x, y)
+    self.assertTrue(jnp.array_equal(result, expected))
+
+  @parameterized.named_parameters(
+      ('basic_f16', OUTPUT_F16, None),
+      ('basic_f32', OUTPUT_F32, None),
+      ('basic_bf16', OUTPUT_BF16, None),
+      ('batch_size_one', OUTPUT_BS1, None),
+      ('weighted_f16', OUTPUT_F16, SAMPLE_WEIGHTS),
+      ('weighted_f32', OUTPUT_F32, SAMPLE_WEIGHTS),
+      ('weighted_bf16', OUTPUT_BF16, SAMPLE_WEIGHTS),
+  )
+  def test_average(self, values, sample_weights):
+    """Test that `Average` metric computes correct values."""
+    if sample_weights is None:
+      sample_weights = jnp.ones_like(values)
+    sample_weights = jnp.array(sample_weights, dtype=values.dtype)
+    metric = metrax.Average.from_model_output(
+        values=values,
+        sample_weights=sample_weights,
+    )
+
+    keras_mean = keras.metrics.Mean(dtype=values.dtype)
+    keras_mean.update_state(values, sample_weights)
+    keras_metrics = keras_mean.result()
+    keras_metrics = jnp.array(keras_metrics, dtype=values.dtype)
+
+    # Use lower tolerance for lower precision dtypes.
+    rtol = 1e-2 if values.dtype in (jnp.float16, jnp.bfloat16) else 1e-05
+    atol = 1e-2 if values.dtype in (jnp.float16, jnp.bfloat16) else 1e-05
+    np.testing.assert_allclose(
+        metric.compute(),
+        keras_metrics,
+        rtol=rtol,
+        atol=atol,
+    )
+
+
+if __name__ == '__main__':
+  absltest.main()

src/metrax/classification_metrics.py

@@ -18,6 +18,7 @@
 import flax
 import jax
 import jax.numpy as jnp
+from metrax import base_metrics
 
 
 def _default_threshold(num_thresholds: int) -> jax.Array:
@@ -42,11 +43,6 @@ def _default_threshold(num_thresholds: int) -> jax.Array:
   return thresholds
 
 
-def _divide_no_nan(x: jax.Array, y: jax.Array) -> jax.Array:
-  """Computes a safe divide which returns 0 if the y is zero."""
-  return jnp.where(y != 0, jnp.divide(x, y), 0.0)
-
-
 @flax.struct.dataclass
 class Precision(clu_metrics.Metric):
   r"""Computes precision for binary classification given `predictions` and `labels`.
@@ -116,7 +112,7 @@ def merge(self, other: 'Precision') -> 'Precision':
     )
 
   def compute(self) -> jax.Array:
-    return _divide_no_nan(
+    return base_metrics.divide_no_nan(
         self.true_positives, (self.true_positives + self.false_positives)
     )
 
@@ -187,7 +183,7 @@ def merge(self, other: 'Recall') -> 'Recall':
     )
 
   def compute(self) -> jax.Array:
-    return _divide_no_nan(
+    return base_metrics.divide_no_nan(
         self.true_positives, (self.true_positives + self.false_negatives)
     )
 
@@ -365,20 +361,20 @@ def interpolate_pr_auc(self) -> jax.Array:
     )
     p = self.true_positives + self.false_positives
     dp = p[: self.num_thresholds - 1] - p[1:]
-    prec_slope = _divide_no_nan(dtp, jnp.maximum(dp, 0))
+    prec_slope = base_metrics.divide_no_nan(dtp, jnp.maximum(dp, 0))
     intercept = self.true_positives[1:] - prec_slope * p[1:]
 
     # recall_relative_ratio
     safe_p_ratio = jnp.where(
         jnp.multiply(p[: self.num_thresholds - 1] > 0, p[1:] > 0),
-        _divide_no_nan(
+        base_metrics.divide_no_nan(
             p[: self.num_thresholds - 1],
             jnp.maximum(p[1:], 0),
         ),
         jnp.ones_like(p[1:]),
     )
     # pr_auc_increment
-    pr_auc_increment = _divide_no_nan(
+    pr_auc_increment = base_metrics.divide_no_nan(
         prec_slope * (dtp + intercept * jnp.log(safe_p_ratio)),
         jnp.maximum(self.true_positives[1:] + self.false_negatives[1:], 0),
     )
@@ -506,10 +502,10 @@ def merge(self, other: 'AUCROC') -> 'AUCROC':
     )
 
   def compute(self) -> jax.Array:
-    tp_rate = _divide_no_nan(
+    tp_rate = base_metrics.divide_no_nan(
         self.true_positives, self.true_positives + self.false_negatives
     )
-    fp_rate = _divide_no_nan(
+    fp_rate = base_metrics.divide_no_nan(
         self.false_positives, self.false_positives + self.true_negatives
     )
     # Threshold goes from 0 to 1, so trapezoid is negative.