change base_metric to base

Author: jshin1394

src/metrax/__init__.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from metrax.base_metrics import (
+from metrax.base import (
     Average,
 )
 from metrax.classification_metrics import (

src/metrax/base_metrics.py src/metrax/base.pysrc/metrax/base_metrics.py renamed to src/metrax/base.py

@@ -18,7 +18,7 @@
 import flax
 import jax
 import jax.numpy as jnp
-from metrax import base_metrics
+from metrax import base
 
 
 def _default_threshold(num_thresholds: int) -> jax.Array:
@@ -112,7 +112,7 @@ def merge(self, other: 'Precision') -> 'Precision':
     )
 
   def compute(self) -> jax.Array:
-    return base_metrics.divide_no_nan(
+    return base.divide_no_nan(
         self.true_positives, (self.true_positives + self.false_positives)
     )
 
@@ -183,7 +183,7 @@ def merge(self, other: 'Recall') -> 'Recall':
     )
 
   def compute(self) -> jax.Array:
-    return base_metrics.divide_no_nan(
+    return base.divide_no_nan(
         self.true_positives, (self.true_positives + self.false_negatives)
     )
 
@@ -361,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 = base_metrics.divide_no_nan(dtp, jnp.maximum(dp, 0))
+    prec_slope = base.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),
-        base_metrics.divide_no_nan(
+        base.divide_no_nan(
             p[: self.num_thresholds - 1],
             jnp.maximum(p[1:], 0),
         ),
         jnp.ones_like(p[1:]),
     )
     # pr_auc_increment
-    pr_auc_increment = base_metrics.divide_no_nan(
+    pr_auc_increment = base.divide_no_nan(
         prec_slope * (dtp + intercept * jnp.log(safe_p_ratio)),
         jnp.maximum(self.true_positives[1:] + self.false_negatives[1:], 0),
     )
@@ -502,10 +502,10 @@ def merge(self, other: 'AUCROC') -> 'AUCROC':
     )
 
   def compute(self) -> jax.Array:
-    tp_rate = base_metrics.divide_no_nan(
+    tp_rate = base.divide_no_nan(
         self.true_positives, self.true_positives + self.false_negatives
     )
-    fp_rate = base_metrics.divide_no_nan(
+    fp_rate = base.divide_no_nan(
         self.false_positives, self.false_positives + self.true_negatives
     )
     # Threshold goes from 0 to 1, so trapezoid is negative.

src/metrax/base_metrics_test.py src/metrax/base_test.pysrc/metrax/base_metrics_test.py renamed to src/metrax/base_test.py

@@ -18,7 +18,7 @@
 import flax
 import jax
 import jax.numpy as jnp
-from metrax import base_metrics
+from metrax import base
 
 
 @flax.struct.dataclass
@@ -86,7 +86,7 @@ def from_model_output(
       ValueError: If type of `labels` is wrong or the shapes of `predictions`
       and `labels` are incompatible.
     """
-    predictions = base_metrics.divide_no_nan(
+    predictions = base.divide_no_nan(
         predictions, jnp.sum(predictions, axis=-1, keepdims=True)
     )
     labels_one_hot = jax.nn.one_hot(labels, predictions.shape[-1], axis=-1)
@@ -97,7 +97,7 @@ def from_model_output(
     if sample_weights is not None:
       crossentropy = crossentropy * sample_weights
       # Normalize by the sum of weights for each sequence.
-      crossentropy = base_metrics.divide_no_nan(
+      crossentropy = base.divide_no_nan(
           jnp.sum(crossentropy), jnp.sum(sample_weights)
       )
     else:
@@ -119,14 +119,12 @@ def merge(self, other: 'Perplexity') -> 'Perplexity':
 
   def compute(self) -> jax.Array:
     return jnp.exp(
-        base_metrics.divide_no_nan(
-            self.aggregate_crossentropy, self.num_samples
-        )
+        base.divide_no_nan(self.aggregate_crossentropy, self.num_samples)
     )
 
 
 @flax.struct.dataclass
-class WER(base_metrics.Average):
+class WER(base.Average):
   r"""Computes Word Error Rate (WER) for speech recognition or text generation tasks.
 
   Word Error Rate measures the edit distance between reference texts and

src/metrax/classification_metrics.py

@@ -17,11 +17,11 @@
 import flax
 import jax
 import jax.numpy as jnp
-from metrax import base_metrics
+from metrax import base
 
 
 @flax.struct.dataclass
-class AveragePrecisionAtK(base_metrics.Average):
+class AveragePrecisionAtK(base.Average):
   r"""Computes AP@k (average precision at k) metrics in JAX.
 
   Average precision at k (AP@k) is a metric used to evaluate the performance of
@@ -63,16 +63,14 @@ def average_precision_at_ks(
     def compute_ap_at_k_single(relevant_labels, total_relevant, ks):
       cumulative_precision = jnp.where(
           relevant_labels,
-          base_metrics.divide_no_nan(
+          base.divide_no_nan(
               jnp.cumsum(relevant_labels),
               jnp.arange(1, len(relevant_labels) + 1),
           ),
           0,
       )
       return jnp.array([
-          base_metrics.divide_no_nan(
-              jnp.sum(cumulative_precision[:k]), total_relevant
-          )
+          base.divide_no_nan(jnp.sum(cumulative_precision[:k]), total_relevant)
           for k in ks
       ])
 

src/metrax/nlp_metrics.py

@@ -18,11 +18,11 @@
 import flax
 import jax
 import jax.numpy as jnp
-from metrax import base_metrics
+from metrax import base
 
 
 @flax.struct.dataclass
-class MSE(base_metrics.Average):
+class MSE(base.Average):
   r"""Computes the mean squared error for regression problems given `predictions` and `labels`.
 
   The mean squared error without sample weights is defined as:
@@ -216,6 +216,6 @@ def compute(self) -> jax.Array:
     Returns:
       The r-squared score.
     """
-    mean = base_metrics.divide_no_nan(self.total, self.count)
+    mean = base.divide_no_nan(self.total, self.count)
     sst = self.sum_of_squared_label - self.count * jnp.power(mean, 2)
-    return 1 - base_metrics.divide_no_nan(self.sum_of_squared_error, sst)
+    return 1 - base.divide_no_nan(self.sum_of_squared_error, sst)