Modified docstrings to match order of variable appearance. Updated docstring to give more description for the from_model_output method. Added beta value into the names for the test cases. Order test cases based on beta value for readability.

Author: 2468cool

src/metrax/classification_metrics.py

@@ -578,10 +578,8 @@ def compute(self) -> jax.Array:
 
 @flax.struct.dataclass
 class FBetaScore(clu_metrics.Metric):
-
     """
     F-Beta score Metric class
-
     Computes the F-Beta score for the binary classification given 'predictions' and 'labels'.
 
     Formula for F-Beta Score:
@@ -591,13 +589,13 @@ class FBetaScore(clu_metrics.Metric):
     F-Beta turns into the F1 Score when beta = 1.0
 
     Attributes:
-        beta: The beta value used in the F-Score metric
         true_positives: The count of true positive instances from the given data,
             label, and threshold.
         false_positives: The count of false positive instances from the given data,
             label, and threshold.
         false_negatives: The count of false negative instances from the given data,
             label, and threshold.
+        beta: The beta value used in the F-Score metric
     """
 
     true_positives: jax.Array
@@ -609,10 +607,10 @@ class FBetaScore(clu_metrics.Metric):
     @classmethod
     def empty(cls) -> 'FBetaScore':
         return cls(
-            beta = 1.0,
             true_positives = jnp.array(0, jnp.float32),
             false_positives = jnp.array(0, jnp.float32),
             false_negatives = jnp.array(0, jnp.float32),
+            beta=1.0,
         )
 
     @classmethod
@@ -623,17 +621,19 @@ def from_model_output(
             beta = beta,
             threshold = 0.5,) -> 'FBetaScore':
         """Updates the metric.
+            Note: When only predictions and labels are given, the score calculated
+            is the F1 score if the FBetaScore beta value has not been previously modified.
 
             Args:
-                threshold: threshold value to use in the F-Score metric a floating number.
-                beta: beta value to use in the F-Score metric. A floating number.
                 predictions: A floating point 1D vector whose values are in the range [0,
                   1]. The shape should be (batch_size,).
                 labels: True value. The value is expected to be 0 or 1. The shape should
                   be (batch_size,).
+                beta: beta value to use in the F-Score metric. A floating number.
+                threshold: threshold value to use in the F-Score metric. A floating number.
 
             Returns:
-                The true positives, false positives, and false negatives.
+                The updated FBetaScore object.
 
             Raises:
                 ValueError: If type of `labels` is wrong or the shapes of `predictions`

src/metrax/classification_metrics_test.py

@@ -17,9 +17,9 @@
 import os
 os.environ['KERAS_BACKEND'] = 'jax'
 
-from src.metrax.classification_metrics import FBetaScore
 from absl.testing import absltest
 from absl.testing import parameterized
+from metrax import FBetaScore
 import jax.numpy as jnp
 import keras
 import metrax
@@ -83,10 +83,10 @@ def test_aucroc_empty(self):
   def test_fbeta_empty(self):
     """Tests the `empty` method of the `FBetaScore` class."""
     m = metrax.FBetaScore.empty()
-    self.assertEqual(m.beta, 1.0)
     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.beta, 1.0)
 
   @parameterized.named_parameters(
       ('basic_f16', OUTPUT_LABELS, OUTPUT_PREDS_F16, SAMPLE_WEIGHTS),
@@ -277,19 +277,18 @@ def test_aucroc(self, inputs, dtype):
 # Testing function for F-Beta classification
   # name, output true, output prediction, threshold, beta
   @parameterized.named_parameters(
-      ('basic_f16', OUTPUT_LABELS, OUTPUT_PREDS_F16, 0.5, 1.0),
-      ('high_threshold_f16', OUTPUT_LABELS, OUTPUT_PREDS_F16, 0.7, 2.0),
-      ('low_threshold_f16', OUTPUT_LABELS, OUTPUT_PREDS_F16, 0.1, 3.0),
-      ('basic_f32', OUTPUT_LABELS, OUTPUT_PREDS_F32, 0.5, 1.0),
-      ('high_threshold_f32', OUTPUT_LABELS, OUTPUT_PREDS_F32, 0.7, 2.0),
-      ('low_threshold_f32', OUTPUT_LABELS, OUTPUT_PREDS_F32, 0.1, 1.0),
-      ('basic_bf16', OUTPUT_LABELS, OUTPUT_PREDS_BF16, 0.5, 3.0),
-      ('high_threshold_bf16', OUTPUT_LABELS, OUTPUT_PREDS_BF16, 0.7, 1.0),
-      ('low_threshold_bf16', OUTPUT_LABELS, OUTPUT_PREDS_BF16, 0.1, 2.0),
-      ('batch_size_one', OUTPUT_LABELS_BS1, OUTPUT_PREDS_BS1, 0.5, 1.0),
+      ('basic_f16_beta_1.0', OUTPUT_LABELS, OUTPUT_PREDS_F16, 0.5, 1.0),
+      ('basic_f32_beta_1.0', OUTPUT_LABELS, OUTPUT_PREDS_F32, 0.5, 1.0),
+      ('low_threshold_f32_beta_1.0', OUTPUT_LABELS, OUTPUT_PREDS_F32, 0.1, 1.0),
+      ('high_threshold_bf16_beta_1.0', OUTPUT_LABELS, OUTPUT_PREDS_BF16, 0.7, 1.0),
+      ('batch_size_one_beta_1.0', OUTPUT_LABELS_BS1, OUTPUT_PREDS_BS1, 0.5, 1.0),
+      ('high_threshold_f16_beta_2.0', OUTPUT_LABELS, OUTPUT_PREDS_F16, 0.7, 2.0),
+      ('high_threshold_f32_beta_2.0', OUTPUT_LABELS, OUTPUT_PREDS_F32, 0.7, 2.0),
+      ('low_threshold_bf16_beta_2.0', OUTPUT_LABELS, OUTPUT_PREDS_BF16, 0.1, 2.0),
+      ('low_threshold_f16_beta_3.0', OUTPUT_LABELS, OUTPUT_PREDS_F16, 0.1, 3.0),
+      ('basic_bf16_beta_3.0', OUTPUT_LABELS, OUTPUT_PREDS_BF16, 0.5, 3.0),
   )
   def test_fbetascore(self, y_true, y_pred, threshold, beta):
-
     # Define the Keras FBeta class to be tested against
     keras_fbeta = keras.metrics.FBetaScore(beta=beta, threshold=threshold)
     keras_fbeta.update_state(y_true, y_pred)