Update imports to be modules not classes

This aligns with the Google style guide:
https://google.github.io/styleguide/pyguide.html#22-imports

Also remove `__all__` from `__init__.py` files. This only applies when
using `import *`, which is not common or recommended, so it's not worth us
protecting against within the library.

Author: jeffcarp

src/metrax/__init__.py

@@ -12,38 +12,20 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from metrax.base import (
-    Average,
-)
-from metrax.classification_metrics import (
-    AUCPR,
-    AUCROC,
-    Precision,
-    Recall,
-)
-from metrax.nlp_metrics import (
-    Perplexity,
-    WER
-)
-from metrax.ranking_metrics import (
-    AveragePrecisionAtK,
-)
-from metrax.regression_metrics import (
-    MSE,
-    RMSE,
-    RSQUARED,
-)
+from metrax import base
+from metrax import classification_metrics
+from metrax import nlp_metrics
+from metrax import ranking_metrics
+from metrax import regression_metrics
 
-__all__ = [
-    "AUCPR",
-    "AUCROC",
-    "Average",
-    "AveragePrecisionAtK",
-    "MSE",
-    "Perplexity",
-    "Precision",
-    "Recall",
-    "RMSE",
-    "RSQUARED",
-    "WER",
-]
+Average = base.Average
+AUCPR = classification_metrics.AUCPR
+AUCROC = classification_metrics.AUCROC
+Precision = classification_metrics.Precision
+Recall = classification_metrics.Recall
+Perplexity = nlp_metrics.Perplexity
+WER = nlp_metrics.WER
+AveragePrecisionAtK = ranking_metrics.AveragePrecisionAtK
+MSE = regression_metrics.MSE
+RMSE = regression_metrics.RMSE
+RSQUARED = regression_metrics.RSQUARED

src/metrax/nnx/__init__.py

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

src/metrax/nnx/nnx_metrics_test.py

@@ -14,12 +14,14 @@
 
 """Tests for metrax NNX metrics."""
 
+import dataclasses
 import importlib
 import inspect
 import pkgutil
+
 from absl.testing import absltest
 from absl.testing import parameterized
-from clu import metrics as clu_metrics
+from flax import nnx
 import metrax
 import metrax.nnx
 
@@ -28,13 +30,17 @@ 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)
+    metrax_metric_keys = [
+        key for key, metric in inspect.getmembers(metrax)
+        if dataclasses.is_dataclass(metric)
+    ]
+    metrax_nnx_metric_keys = [
+        key for key, metric in inspect.getmembers(metrax.nnx)
+        if inspect.isclass(metric) and issubclass(metric, nnx.Metric)
+    ]
+    self.assertGreater(len(metrax_metric_keys), 0)
+    self.assertSameElements(metrax_metric_keys, metrax_nnx_metric_keys)
+    self.assertSameElements(metrax_metric_keys, metrax_nnx_metric_keys)
 
 
 if __name__ == '__main__':