add nnx wrappers for metrax metrics

Author: jshin1394

src/metrax/nnx/__init__.py

@@ -12,10 +12,30 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from metrax.nnx.nnx_wrapper import (
-    NnxWrapper,
+from metrax.nnx.nnx_metrics import (
+    AUCPR,
+    AUCROC,
+    Average,
+    AveragePrecisionAtK,
+    MSE,
+    Perplexity,
+    Precision,
+    RMSE,
+    RSQUARED,
+    Recall,
+    WER,
 )
 
 __all__ = [
-    "NnxWrapper",
+    "AUCPR",
+    "AUCROC",
+    "Average",
+    "AveragePrecisionAtK",
+    "MSE",
+    "Perplexity",
+    "Precision",
+    "Recall",
+    "RMSE",
+    "RSQUARED",
+    "WER",
 ]

src/metrax/nnx/nnx_metric.py

@@ -0,0 +1,95 @@
+# 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.
+
+import metrax
+from metrax.nnx import nnx_wrapper
+
+NnxWrapper = nnx_wrapper.NnxWrapper
+
+
+class AUCPR(NnxWrapper):
+  """An NNX class for the Metrax metric AUCPR."""
+
+  def __init__(self):
+    super().__init__(metrax.AUCPR)
+
+
+class AUCROC(NnxWrapper):
+  """An NNX class for the Metrax metric AUCROC."""
+
+  def __init__(self):
+    super().__init__(metrax.AUCROC)
+
+
+class Average(NnxWrapper):
+  """An NNX class for the Metrax metric Average."""
+
+  def __init__(self):
+    super().__init__(metrax.Average)
+
+
+class AveragePrecisionAtK(NnxWrapper):
+  """An NNX class for the Metrax metric AveragePrecisionAtK."""
+
+  def __init__(self):
+    super().__init__(metrax.AveragePrecisionAtK)
+
+
+class MSE(NnxWrapper):
+  """An NNX class for the Metrax metric MSE."""
+
+  def __init__(self):
+    super().__init__(metrax.MSE)
+
+
+class Perplexity(NnxWrapper):
+  """An NNX class for the Metrax metric Perplexity."""
+
+  def __init__(self):
+    super().__init__(metrax.Perplexity)
+
+
+class Precision(NnxWrapper):
+  """An NNX class for the Metrax metric Precision."""
+
+  def __init__(self):
+    super().__init__(metrax.Precision)
+
+
+class Recall(NnxWrapper):
+  """An NNX class for the Metrax metric Recall."""
+
+  def __init__(self):
+    super().__init__(metrax.Recall)
+
+
+class RMSE(NnxWrapper):
+  """An NNX class for the Metrax metric RMSE."""
+
+  def __init__(self):
+    super().__init__(metrax.RMSE)
+
+
+class RSQUARED(NnxWrapper):
+  """An NNX class for the Metrax metric RSQUARED."""
+
+  def __init__(self):
+    super().__init__(metrax.RSQUARED)
+
+
+class WER(NnxWrapper):
+  """An NNX class for the Metrax metric WER."""
+
+  def __init__(self):
+    super().__init__(metrax.WER)

src/metrax/nnx/nnx_metric_test.py

@@ -0,0 +1,41 @@
+# 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 NNX metrics."""
+
+import importlib
+import inspect
+import pkgutil
+from absl.testing import absltest
+from absl.testing import parameterized
+from clu import metrics as clu_metrics
+import metrax
+import metrax.nnx
+
+
+class NnxMetricsTest(parameterized.TestCase):
+
+  def test_nnx_metrics_exists(self):
+    """Tests that every metrax CLU metric has an NNX counterpart."""
+    metrax_metrics = metrax.nnx.__all__
+    for _, module_name, _ in pkgutil.iter_modules(metrax.__path__):
+      full_module_name = f"{metrax.__name__}.{module_name}"
+      module = importlib.import_module(full_module_name)
+      for name, obj in inspect.getmembers(module):
+        if inspect.isclass(obj) and issubclass(obj, clu_metrics.Metric):
+          self.assertIn(name, metrax_metrics)
+
+
+if __name__ == "__main__":
+  absltest.main()

src/metrax/nnx/nnx_wrapper_test.py

@@ -50,7 +50,7 @@ class NnxWrapperTest(parameterized.TestCase):
 
   def test_reset(self):
     """Tests the `reset` method of the `NnxWrapper` class."""
-    nnx_metric = metrax.nnx.NnxWrapper(metrax.MSE)
+    nnx_metric = metrax.nnx.MSE()
     self.assertEqual(nnx_metric.clu_metric.total, jnp.array(0, jnp.float32))
     self.assertEqual(nnx_metric.clu_metric.count, jnp.array(0, jnp.int32))
     nnx_metric.update(
@@ -78,7 +78,7 @@ def test_metric_update_and_compute(self, y_true, y_pred, sample_weights):
     if sample_weights is None:
       sample_weights = np.ones_like(y_true)
 
-    nnx_metric = metrax.nnx.NnxWrapper(metrax.MSE)
+    nnx_metric = metrax.nnx.MSE()
     for labels, logits, weights in zip(y_true, y_pred, sample_weights):
       nnx_metric.update(
           predictions=logits,