Added 3 new classification metrics.

src/metrax/__init__.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+"""Metrax metrics collection."""
+
 from metrax import audio_metrics
 from metrax import base
 from metrax import classification_metrics
@@ -26,6 +28,8 @@
 Average = base.Average
 AveragePrecisionAtK = ranking_metrics.AveragePrecisionAtK
 BLEU = nlp_metrics.BLEU
+BinaryAccuracy = classification_metrics.BinaryAccuracy
+CategoricalAccuracy = classification_metrics.CategoricalAccuracy
 CosineSimilarity = image_metrics.CosineSimilarity
 DCGAtK = ranking_metrics.DCGAtK
 Dice = image_metrics.Dice
@@ -48,6 +52,7 @@
 RougeL = nlp_metrics.RougeL
 RougeN = nlp_metrics.RougeN
 SNR = audio_metrics.SNR
+SparseCategoricalAccuracy = classification_metrics.SparseCategoricalAccuracy
 SpearmanRankCorrelation = regression_metrics.SpearmanRankCorrelation
 SSIM = image_metrics.SSIM
 WER = nlp_metrics.WER
@@ -60,6 +65,8 @@
     "Average",
     "AveragePrecisionAtK",
     "BLEU",
+    "BinaryAccuracy",
+    "CategoricalAccuracy",
     "CosineSimilarity",
     "DCGAtK",
     "Dice",
@@ -77,6 +84,7 @@
     "RMSE",
     "RMSLE",
     "RSQUARED",
+    "SparseCategoricalAccuracy",
     "SpearmanRankCorrelation",
     "Recall",
     "RecallAtK",

src/metrax/classification_metrics.py

@@ -110,6 +110,125 @@ def from_model_output(
     )
 
 
+@flax.struct.dataclass
+class BinaryAccuracy(base.Average):
+  r"""Computes binary classification accuracy for predictions and labels.
+
+  This metric calculates the proportion of correct predictions by comparing
+  `predictions >= threshold` and `labels` element-wise. It is the ratio of the
+  sum of weighted correct predictions to the sum of all corresponding weights.
+  If no `sample_weights` are provided, weights default to 1 for each element.
+  """
+
+  @classmethod
+  def from_model_output(
+      cls,
+      predictions: jax.Array,
+      labels: jax.Array,
+      sample_weights: jax.Array | None = None,
+      threshold: float = 0.5,
+  ) -> 'BinaryAccuracy':
+    """Updates the metric state with new `predictions` and `labels`.
+
+    Args:
+      predictions: JAX array of predicted values.
+      labels: JAX array of true binary values (0 or 1).
+      sample_weights: Optional JAX array of sample weights.
+      threshold: The threshold parameter used to convert predicted probabilities
+        to binary decisions.
+
+    Returns:
+      An updated instance of `BinaryAccuracy` metric.
+    """
+    correct = (predictions >= threshold) == labels
+    count = jnp.ones_like(labels, dtype=jnp.int32)
+    if sample_weights is not None:
+      correct = correct * sample_weights
+      count = count * sample_weights
+    return cls(
+        total=correct.sum(),
+        count=count.sum(),
+    )
+
+
+@flax.struct.dataclass
+class CategoricalAccuracy(base.Average):
+  r"""Computes accuracy for one-hot categorical classification models.
+
+  This metric calculates the frequency with which the predicted class matches
+  the true class by comparing `argmax(predictions, axis=-1)` and
+  `argmax(labels, axis=-1)`. If no `sample_weights` are provided, weights
+  default to 1 for each element.
+  """
+
+  @classmethod
+  def from_model_output(
+      cls,
+      predictions: jax.Array,
+      labels: jax.Array,
+      sample_weights: jax.Array | None = None,
+  ) -> 'CategoricalAccuracy':
+    """Updates the metric state with new `predictions` and `labels`.
+
+    Args:
+      predictions: JAX array of predicted probability distributions or logits.
+      labels: JAX array of one-hot encoded true class labels.
+      sample_weights: Optional JAX array of sample weights.
+
+    Returns:
+      An updated instance of `CategoricalAccuracy` metric.
+    """
+    correct = jnp.argmax(predictions, axis=-1) == jnp.argmax(labels, axis=-1)
+    count = jnp.ones_like(correct, dtype=jnp.int32)
+    if sample_weights is not None:
+      correct = correct * sample_weights
+      count = count * sample_weights
+    return cls(
+        total=correct.sum(),
+        count=count.sum(),
+    )
+
+
+@flax.struct.dataclass
+class SparseCategoricalAccuracy(base.Average):
+  r"""Computes accuracy for sparse categorical classification models.
+
+  This metric calculates the frequency with which the predicted class matches
+  the integer target class by comparing `argmax(predictions, axis=-1)` and
+  `labels`. If no `sample_weights` are provided, weights default to 1 for
+  each element.
+  """
+
+  @classmethod
+  def from_model_output(
+      cls,
+      predictions: jax.Array,
+      labels: jax.Array,
+      sample_weights: jax.Array | None = None,
+  ) -> 'SparseCategoricalAccuracy':
+    """Updates the metric state with new `predictions` and `labels`.
+
+    Args:
+      predictions: JAX array of predicted probability distributions or logits.
+      labels: JAX array of ground truth integer class labels.
+      sample_weights: Optional JAX array of sample weights.
+
+    Returns:
+      An updated instance of `SparseCategoricalAccuracy` metric.
+    """
+    if labels.ndim == predictions.ndim:
+      labels = jnp.squeeze(labels, axis=-1)
+    correct = jnp.argmax(predictions, axis=-1) == labels
+    count = jnp.ones_like(labels, dtype=jnp.int32)
+    if sample_weights is not None:
+      correct = correct * sample_weights
+      count = count * sample_weights
+    return cls(
+        total=correct.sum(),
+        count=count.sum(),
+    )
+
+
 @flax.struct.dataclass
 class Precision(clu_metrics.Metric):
   r"""Computes precision for binary classification given `predictions` and `labels`.

src/metrax/classification_metrics_test.py

@@ -47,6 +47,16 @@
     [0.5, 1, 0, 0, 0, 0, 0, 0],
     (BATCHES, 1),
 ).astype(np.float32)
+NUM_CLASSES = 5
+MC_LABELS_INT = np.random.randint(
+    0, NUM_CLASSES, size=(BATCHES, BATCH_SIZE)
+).astype(np.int32)
+MC_LABELS_OH = np.eye(NUM_CLASSES)[MC_LABELS_INT].astype(np.float32)
+MC_PREDS = np.random.uniform(size=(BATCHES, BATCH_SIZE, NUM_CLASSES)).astype(
+    np.float32
+)
+MC_PREDS_F16 = MC_PREDS.astype(jnp.float16)
+MC_PREDS_BF16 = MC_PREDS.astype(jnp.bfloat16)
 
 
 class ClassificationMetricsTest(parameterized.TestCase):
@@ -89,6 +99,24 @@ def test_fbeta_empty(self):
     self.assertEqual(m.false_negatives, jnp.array(0, jnp.float32))
     self.assertEqual(m.beta, 1.0)
 
+  def test_binary_accuracy_empty(self):
+    """Tests the `empty` method of `BinaryAccuracy`."""
+    m = metrax.BinaryAccuracy.empty()
+    self.assertEqual(m.total, jnp.array(0, jnp.float32))
+    self.assertEqual(m.count, jnp.array(0, jnp.int32))
+
+  def test_categorical_accuracy_empty(self):
+    """Tests the `empty` method of `CategoricalAccuracy`."""
+    m = metrax.CategoricalAccuracy.empty()
+    self.assertEqual(m.total, jnp.array(0, jnp.float32))
+    self.assertEqual(m.count, jnp.array(0, jnp.int32))
+
+  def test_sparse_categorical_accuracy_empty(self):
+    """Tests the `empty` method of `SparseCategoricalAccuracy`."""
+    m = metrax.SparseCategoricalAccuracy.empty()
+    self.assertEqual(m.total, jnp.array(0, jnp.float32))
+    self.assertEqual(m.count, jnp.array(0, jnp.int32))
+
   @parameterized.named_parameters(
       ('basic_f16', OUTPUT_LABELS, OUTPUT_PREDS_F16, SAMPLE_WEIGHTS),
       ('basic_f32', OUTPUT_LABELS, OUTPUT_PREDS_F32, SAMPLE_WEIGHTS),
@@ -120,6 +148,99 @@ def test_accuracy(self, y_true, y_pred, sample_weights):
         atol=atol,
     )
 
+  @parameterized.named_parameters(
+      ('f16', OUTPUT_LABELS, OUTPUT_PREDS_F16, SAMPLE_WEIGHTS, 0.5),
+      ('high_f16', OUTPUT_LABELS, OUTPUT_PREDS_F16, SAMPLE_WEIGHTS, 0.7),
+      ('low_f16', OUTPUT_LABELS, OUTPUT_PREDS_F16, SAMPLE_WEIGHTS, 0.1),
+      ('f32', OUTPUT_LABELS, OUTPUT_PREDS_F32, SAMPLE_WEIGHTS, 0.5),
+      ('high_f32', OUTPUT_LABELS, OUTPUT_PREDS_F32, SAMPLE_WEIGHTS, 0.7),
+      ('low_f32', OUTPUT_LABELS, OUTPUT_PREDS_F32, SAMPLE_WEIGHTS, 0.1),
+      ('bf16', OUTPUT_LABELS, OUTPUT_PREDS_BF16, SAMPLE_WEIGHTS, 0.5),
+      ('bs_one', OUTPUT_LABELS_BS1, OUTPUT_PREDS_BS1, None, 0.5),
+  )
+  def test_binary_accuracy(self, y_true, y_pred, sample_weights, threshold):
+    """Test that `BinaryAccuracy` metric computes correct values."""
+    if sample_weights is None:
+      sample_weights = np.ones_like(y_true)
+    metrax_metric = metrax.BinaryAccuracy.empty()
+    keras_metric = keras.metrics.BinaryAccuracy(threshold=threshold)
+    for labels, logits, weights in zip(y_true, y_pred, sample_weights):
+      update = metrax.BinaryAccuracy.from_model_output(
+          predictions=logits,
+          labels=labels,
+          sample_weights=weights,
+          threshold=threshold,
+      )
+      metrax_metric = metrax_metric.merge(update)
+      keras_metric.update_state(labels, logits, weights)
+
+    rtol = 1e-2 if y_pred.dtype in (jnp.float16, jnp.bfloat16) else 1e-5
+    atol = 1e-2 if y_pred.dtype in (jnp.float16, jnp.bfloat16) else 1e-5
+    np.testing.assert_allclose(
+        metrax_metric.compute(),
+        keras_metric.result(),
+        rtol=rtol,
+        atol=atol,
+    )
+
+  @parameterized.named_parameters(
+      ('f16', MC_LABELS_OH, MC_PREDS_F16, SAMPLE_WEIGHTS),
+      ('f32', MC_LABELS_OH, MC_PREDS, SAMPLE_WEIGHTS),
+      ('bf16', MC_LABELS_OH, MC_PREDS_BF16, None),
+  )
+  def test_categorical_accuracy(self, y_true, y_pred, sample_weights):
+    """Test that `CategoricalAccuracy` metric computes correct values."""
+    if sample_weights is None:
+      sample_weights = np.ones((BATCHES, BATCH_SIZE), dtype=np.float32)
+    metrax_metric = metrax.CategoricalAccuracy.empty()
+    keras_metric = keras.metrics.CategoricalAccuracy()
+    for labels, logits, weights in zip(y_true, y_pred, sample_weights):
+      update = metrax.CategoricalAccuracy.from_model_output(
+          predictions=logits,
+          labels=labels,
+          sample_weights=weights,
+      )
+      metrax_metric = metrax_metric.merge(update)
+      keras_metric.update_state(labels, logits, weights)
+
+    rtol = 1e-2 if y_pred.dtype in (jnp.float16, jnp.bfloat16) else 1e-5
+    atol = 1e-2 if y_pred.dtype in (jnp.float16, jnp.bfloat16) else 1e-5
+    np.testing.assert_allclose(
+        metrax_metric.compute(),
+        keras_metric.result(),
+        rtol=rtol,
+        atol=atol,
+    )
+
+  @parameterized.named_parameters(
+      ('f16', MC_LABELS_INT, MC_PREDS_F16, SAMPLE_WEIGHTS),
+      ('f32', MC_LABELS_INT, MC_PREDS, SAMPLE_WEIGHTS),
+      ('bf16', MC_LABELS_INT, MC_PREDS_BF16, None),
+  )
+  def test_sparse_categorical_accuracy(self, y_true, y_pred, sample_weights):
+    """Test that `SparseCategoricalAccuracy` computes correct values."""
+    if sample_weights is None:
+      sample_weights = np.ones((BATCHES, BATCH_SIZE), dtype=np.float32)
+    metrax_metric = metrax.SparseCategoricalAccuracy.empty()
+    keras_metric = keras.metrics.SparseCategoricalAccuracy()
+    for labels, logits, weights in zip(y_true, y_pred, sample_weights):
+      update = metrax.SparseCategoricalAccuracy.from_model_output(
+          predictions=logits,
+          labels=labels,
+          sample_weights=weights,
+      )
+      metrax_metric = metrax_metric.merge(update)
+      keras_metric.update_state(labels, logits, weights)
+
+    rtol = 1e-2 if y_pred.dtype in (jnp.float16, jnp.bfloat16) else 1e-5
+    atol = 1e-2 if y_pred.dtype in (jnp.float16, jnp.bfloat16) else 1e-5
+    np.testing.assert_allclose(
+        metrax_metric.compute(),
+        keras_metric.result(),
+        rtol=rtol,
+        atol=atol,
+    )
+
   @parameterized.named_parameters(
       ('basic_f16', OUTPUT_LABELS, OUTPUT_PREDS_F16, 0.5),
       ('high_threshold_f16', OUTPUT_LABELS, OUTPUT_PREDS_F16, 0.7),

src/metrax/nnx/__init__.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+"""Metrax NNX metrics collection."""
+
 from metrax.nnx import nnx_metrics
 
 AUCPR = nnx_metrics.AUCPR
@@ -20,6 +22,8 @@
 Average = nnx_metrics.Average
 AveragePrecisionAtK = nnx_metrics.AveragePrecisionAtK
 BLEU = nnx_metrics.BLEU
+BinaryAccuracy = nnx_metrics.BinaryAccuracy
+CategoricalAccuracy = nnx_metrics.CategoricalAccuracy
 CosineSimilarity = nnx_metrics.CosineSimilarity
 DCGAtK = nnx_metrics.DCGAtK
 Dice = nnx_metrics.Dice
@@ -42,6 +46,7 @@
 RougeL = nnx_metrics.RougeL
 RougeN = nnx_metrics.RougeN
 SNR = nnx_metrics.SNR
+SparseCategoricalAccuracy = nnx_metrics.SparseCategoricalAccuracy
 SpearmanRankCorrelation = nnx_metrics.SpearmanRankCorrelation
 SSIM = nnx_metrics.SSIM
 WER = nnx_metrics.WER
@@ -53,6 +58,8 @@
     "Average",
     "AveragePrecisionAtK",
     "BLEU",
+    "BinaryAccuracy",
+    "CategoricalAccuracy",
     "CosineSimilarity",
     "DCGAtK",
     "Dice",
@@ -73,6 +80,7 @@
     "RougeL",
     "RougeN",
     "SNR",
+    "SparseCategoricalAccuracy",
     "SpearmanRankCorrelation",
     "SSIM",
     "WER",