Implement empty function (#14)

Author: jshin1394

src/metrax/metrics.py

@@ -165,6 +165,15 @@ class RSQUARED(clu_metrics.Metric):
   count: jax.Array
   sum_of_squared_error: jax.Array
   sum_of_squared_label: jax.Array
+  
+
+  @classmethod
+  def empty(cls) -> 'RSQUARED':
+    return cls(
+      total=jnp.array(0, jnp.float32),
+      count=jnp.array(0, jnp.float32),
+      sum_of_squared_error=jnp.array(0, jnp.float32),
+      sum_of_squared_label=jnp.array(0, jnp.float32))
 
   @classmethod
   def from_model_output(
@@ -259,6 +268,12 @@ class Precision(clu_metrics.Metric):
   true_positives: jax.Array
   false_positives: jax.Array
 
+  @classmethod
+  def empty(cls) -> 'Precision':
+    return cls(
+      true_positives=jnp.array(0, jnp.float32),
+      false_positives=jnp.array(0, jnp.float32))
+
   @classmethod
   def from_model_output(
       cls,
@@ -326,6 +341,12 @@ class Recall(clu_metrics.Metric):
   true_positives: jax.Array
   false_negatives: jax.Array
 
+  @classmethod
+  def empty(cls) -> 'Recall':
+    return cls(
+      true_positives=jnp.array(0, jnp.float32),
+      false_negatives=jnp.array(0, jnp.float32))
+
   @classmethod
   def from_model_output(
       cls, predictions: jax.Array, labels: jax.Array, threshold: float = 0.5
@@ -415,6 +436,14 @@ class AUCPR(clu_metrics.Metric):
   false_positives: jax.Array
   false_negatives: jax.Array
   num_thresholds: int
+  
+  @classmethod
+  def empty(cls) -> 'AUCPR':
+    return cls(
+      true_positives=jnp.array(0, jnp.float32),
+      false_positives=jnp.array(0, jnp.float32),
+      false_negatives=jnp.array(0, jnp.float32),
+      num_thresholds=0)
 
   @classmethod
   def from_model_output(
@@ -589,6 +618,15 @@ class AUCROC(clu_metrics.Metric):
   false_negatives: jax.Array
   num_thresholds: int
 
+  @classmethod
+  def empty(cls) -> 'AUCROC':
+    return cls(
+      true_positives=jnp.array(0, jnp.float32),
+      true_negatives=jnp.array(0, jnp.float32),
+      false_positives=jnp.array(0, jnp.float32),
+      false_negatives=jnp.array(0, jnp.float32),
+      num_thresholds=0)
+
   @classmethod
   def from_model_output(
       cls,
@@ -698,6 +736,12 @@ class Perplexity(clu_metrics.Metric):
   aggregate_crossentropy: jax.Array
   num_samples: jax.Array
 
+  @classmethod
+  def empty(cls) -> 'Perplexity':
+    return cls(
+      aggregate_crossentropy=jnp.array(0, jnp.float32),
+      num_samples=jnp.array(0, jnp.float32))
+
   @classmethod
   def from_model_output(
       cls,

src/metrax/metrics_test.py

@@ -117,6 +117,61 @@ def compute_aucroc(self, model_outputs, sample_weights=None):
       metric = update if metric is None else metric.merge(update)
     return metric.compute()
 
+  def test_mse_empty(self):
+    """Tests the `empty` method of the `MSE` class."""
+    m = metrax.MSE.empty()
+    self.assertEqual(m.total, jnp.array(0, jnp.float32))
+    self.assertEqual(m.count, jnp.array(0, jnp.int32))
+
+  def test_rmse_empty(self):
+    """Tests the `empty` method of the `RMSE` class."""
+    m = metrax.RMSE.empty()
+    self.assertEqual(m.total, jnp.array(0, jnp.float32))
+    self.assertEqual(m.count, jnp.array(0, jnp.int32))
+
+  def test_rsquared_empty(self):
+    """Tests the `empty` method of the `RSQUARED` class."""
+    m = metrax.RSQUARED.empty()
+    self.assertEqual(m.total, jnp.array(0, jnp.float32))
+    self.assertEqual(m.count, jnp.array(0, jnp.float32))
+    self.assertEqual(m.sum_of_squared_error, jnp.array(0, jnp.float32))
+    self.assertEqual(m.sum_of_squared_label, jnp.array(0, jnp.float32))
+
+  def test_precision_empty(self):
+    """Tests the `empty` method of the `Precision` class."""
+    m = metrax.Precision.empty()
+    self.assertEqual(m.true_positives, jnp.array(0, jnp.float32))
+    self.assertEqual(m.false_positives, jnp.array(0, jnp.float32))
+
+  def test_recall_empty(self):
+    """Tests the `empty` method of the `Recall` class."""
+    m = metrax.Recall.empty()
+    self.assertEqual(m.true_positives, jnp.array(0, jnp.float32))
+    self.assertEqual(m.false_negatives, jnp.array(0, jnp.float32))
+
+  def test_aucpr_empty(self):
+    """Tests the `empty` method of the `AUCPR` class."""
+    m = metrax.AUCPR.empty()
+    self.assertEqual(m.true_positives, jnp.array(0, jnp.float32))
+    self.assertEqual(m.false_positives, jnp.array(0, jnp.float32))
+    self.assertEqual(m.false_negatives, jnp.array(0, jnp.float32))
+    self.assertEqual(m.num_thresholds, 0)
+
+  def test_aucroc_empty(self):
+    """Tests the `empty` method of the `AUCROC` class."""
+    m = metrax.AUCROC.empty()
+    self.assertEqual(m.true_positives, jnp.array(0, jnp.float32))
+    self.assertEqual(m.true_negatives, jnp.array(0, jnp.float32))
+    self.assertEqual(m.false_positives, jnp.array(0, jnp.float32))
+    self.assertEqual(m.false_negatives, jnp.array(0, jnp.float32))
+    self.assertEqual(m.num_thresholds, 0)
+
+  def test_perplexity_empty(self):
+    """Tests the `empty` method of the `Perplexity` class."""
+    m = metrax.Perplexity.empty()
+    self.assertEqual(m.aggregate_crossentropy, jnp.array(0, jnp.float32))
+    self.assertEqual(m.num_samples, jnp.array(0, jnp.float32))
+
   @parameterized.named_parameters(
       ('basic', OUTPUT_LABELS, OUTPUT_PREDS, 0.5),
       ('high_threshold', OUTPUT_LABELS, OUTPUT_PREDS, 0.7),