Repository: google/flax

Issue:
Accuracy doesn't work for simple binary classifier
Hello,

This is quite a simple issue but the function `nnx.metric.Accuracy`  only works for a network that returns more than one output, i.e. 
`super().update(values=(logits.argmax(axis=-1) == labels))`. Not sure if this is by design. 

I was playing around with a simple binary classifier than only returned a single logit per observation. The Accuracy function gives nonsense for this case.

Anyway, might be worth noting in the doc string what the function expects.

Hints from issue comments:
@cgarciae: Yeah I think current implementation is for multi class. We could add support for binary classification.@mattbahr: I'm new to the project and am interested in getting involved. I can take a stab at this if no one else is working on it.

Solution patch:
diff --git a/docs_nnx/api_reference/flax.nnx/variables.rst b/docs_nnx/api_reference/flax.nnx/variables.rst
index 02fd3048c..0bc903db0 100644
--- a/docs_nnx/api_reference/flax.nnx/variables.rst
+++ b/docs_nnx/api_reference/flax.nnx/variables.rst
@@ -21,4 +21,8 @@ variables
 .. autoclass:: VariableState
    :members:
 
-.. autofunction:: with_metadata
\ No newline at end of file
+.. autofunction:: with_metadata
+
+.. autofunction:: variable_name_from_type
+.. autofunction:: variable_type_from_name
+.. autofunction:: set_variable_name
diff --git a/docs_nnx/guides/bridge_guide.ipynb b/docs_nnx/guides/bridge_guide.ipynb
index c24b76a57..0f0775084 100644
--- a/docs_nnx/guides/bridge_guide.ipynb
+++ b/docs_nnx/guides/bridge_guide.ipynb
@@ -632,7 +632,7 @@
    "source": [
     "### NNX to Linen\n",
     "\n",
-    "If you define custom NNX variable types, you should register their names with `nnx.register_variable_name_type_pair` so that they go to the desired collections."
+    "If you define custom NNX variable types, you should register their names with `nnx.register_variable_name` so that they go to the desired collections."
    ]
   },
   {
@@ -653,8 +653,8 @@
     }
    ],
    "source": [
+    "@nnx.register_variable_name('counts', overwrite=True)\n",
     "class Count(nnx.Variable): pass\n",
-    "nnx.register_variable_name_type_pair('counts', Count, overwrite=True)\n",
     "\n",
     "class NNXMultiCollections(nnx.Module):\n",
     "  def __init__(self, din, dout, rngs):\n",
diff --git a/docs_nnx/guides/bridge_guide.md b/docs_nnx/guides/bridge_guide.md
index 3e2c9b4aa..a5abbb0db 100644
--- a/docs_nnx/guides/bridge_guide.md
+++ b/docs_nnx/guides/bridge_guide.md
@@ -389,12 +389,13 @@ y = model(x, mutable=True)  # still works!
 
 ### NNX to Linen
 
-If you define custom NNX variable types, you should register their names with `nnx.register_variable_name_type_pair` so that they go to the desired collections.
+If you define custom NNX variable types, you should register their names with `nnx.register_variable_name` so that they go to the desired collections.
 
 
 ```python
+@nnx.register_variable_name('counts', overwrite=True)
 class Count(nnx.Variable): pass
-nnx.register_variable_name_type_pair('counts', Count, overwrite=True)
+
 
 class NNXMultiCollections(nnx.Module):
   def __init__(self, din, dout, rngs):
diff --git a/flax/nnx/__init__.py b/flax/nnx/__init__.py
index 5c24d6375..7c4077a11 100644
--- a/flax/nnx/__init__.py
+++ b/flax/nnx/__init__.py
@@ -121,9 +121,7 @@
 from .spmd import with_sharding_constraint as with_sharding_constraint
 from .statelib import State as State
 from .training import metrics as metrics
-from .variablelib import (
-  Param as Param,
-)
+from .variablelib import Param as Param
 # this needs to be imported before optimizer to prevent circular import
 from .training import optimizer as optimizer
 from .training.metrics import Metric as Metric
@@ -163,7 +161,7 @@
 from .variablelib import with_metadata as with_metadata
 from .variablelib import variable_type_from_name as variable_type_from_name
 from .variablelib import variable_name_from_type as variable_name_from_type
-from .variablelib import register_variable_name_type_pair as register_variable_name_type_pair
+from .variablelib import register_variable_name as register_variable_name
 from .visualization import display as display
 from .extract import to_tree as to_tree
 from .extract import from_tree as from_tree
diff --git a/flax/nnx/module.py b/flax/nnx/module.py
index 39ee7e3b7..14679ff52 100644
--- a/flax/nnx/module.py
+++ b/flax/nnx/module.py
@@ -249,7 +249,7 @@ def perturb(
     if not hasattr(self, name):
       zeros = jax.tree.map(jnp.zeros_like, value)
       setattr(self, name, variable_type(zeros))
-    old_value = getattr(self, name)
+    old_value: variableslib.Variable[tp.Any] = getattr(self, name)
     if not isinstance(old_value, variable_type):
       raise ValueError(
         f"Expected '{name}' to be of type '{variable_type.__name__}', "
diff --git a/flax/nnx/training/metrics.py b/flax/nnx/training/metrics.py
index eca41bd4f..8a0ffabee 100644
--- a/flax/nnx/training/metrics.py
+++ b/flax/nnx/training/metrics.py
@@ -240,26 +240,64 @@ class Accuracy(Average):
     >>> metrics.reset()
     >>> metrics.compute()
     Array(nan, dtype=float32)
+
+    >>> logits3 = jax.random.normal(jax.random.key(2), (5,))
+    >>> labels3 = jnp.array([0, 1, 0, 1, 1])
+    >>> accuracy = nnx.metrics.Accuracy(threshold=0.5)
+    >>> accuracy.update(logits=logits3, labels=labels3)
+    >>> accuracy.compute()
+    Array(0.8, dtype=float32)
   """
 
+  def __init__(self, threshold: float | None = None, *args, **kwargs):
+    """For binary classification, pass in a float denoting a threshold to determine if a
+    prediction is positive. For example, constructing the metric as
+    ``acc = Accuracy(threshold=0.5)`` would cause any logit greater than or equal to 0.5
+    to be interpreted as a positive classification. For multi-class classification, do
+    not pass in a threshold.
+
+    Args:
+      threshold: for binary classification, determines if a prediction is
+        positive. Defaults to None.
+    """
+    if (threshold is not None) and (not isinstance(threshold, float)):
+      raise TypeError(f'Expected threshold to be a float, got {type(threshold)}')
+
+    self.threshold = threshold
+    super().__init__(*args, **kwargs)
+
   def update(self, *, logits: jax.Array, labels: jax.Array, **_) -> None:  # type: ignore[override]
     """In-place update this ``Metric``.
 
     Args:
-      logits: the outputted predicted activations. These values are
-        argmax-ed (on the trailing dimension), before comparing them
-        to the labels.
+      logits: the outputted predicted activations. For multi-class
+        classification, these values are argmax-ed (on the trailing
+        dimension), before comparing them to the labels. For binary
+        classification, these values are compared to the labels directly.
       labels: the ground truth integer labels.
     """
-    if logits.ndim != labels.ndim + 1:
+    if self.threshold is not None:  # Binary classification case
+      if logits.ndim != labels.ndim:
+        raise ValueError(
+          'For binary classification, expected logits.ndim==labels.ndim, got '
+          f'{logits.ndim} and {labels.ndim}'
+        )
+    elif logits.ndim != labels.ndim + 1:  # Multi-class classification case
       raise ValueError(
-        f'Expected logits.ndim==labels.ndim+1, got {logits.ndim} and {labels.ndim}'
+        'For multi-class classification, expected logits.ndim==labels.ndim+1, '
+        f'got {logits.ndim} and {labels.ndim}'
       )
-    elif labels.dtype in (jnp.int64, np.int32, np.int64):
+
+    if labels.dtype in (jnp.int64, np.int32, np.int64):
       labels = jnp.astype(labels, jnp.int32)
     elif labels.dtype != jnp.int32:
       raise ValueError(f'Expected labels.dtype==jnp.int32, got {labels.dtype}')
 
+    if self.threshold is not None:  # Binary classification case
+      super().update(values=((logits >= self.threshold) == (labels > 0)))
+      return
+
+    # Multi-class classification case
     super().update(values=(logits.argmax(axis=-1) == labels))
 
 
@@ -278,6 +316,7 @@ class MultiMetric(Metric):
     >>> metrics
     MultiMetric( # MetricState: 4 (16 B)
       accuracy=Accuracy( # MetricState: 2 (8 B)
+        threshold=None,
         argname='values',
         total=MetricState( # 1 (4 B)
           value=Array(0., dtype=float32)
@@ -299,6 +338,7 @@ class MultiMetric(Metric):
 
     >>> metrics.accuracy
     Accuracy( # MetricState: 2 (8 B)
+      threshold=None,
       argname='values',
       total=MetricState( # 1 (4 B)
         value=Array(0., dtype=float32)
diff --git a/flax/nnx/variablelib.py b/flax/nnx/variablelib.py
index bc0b42141..7e5048091 100644
--- a/flax/nnx/variablelib.py
+++ b/flax/nnx/variablelib.py
@@ -41,7 +41,6 @@
 RemoveAxisHook = tp.Callable[[V, AxisIndex, AxisName | None], None]
 
 
-
 @dataclasses.dataclass
 class VariableMetadata(tp.Generic[A]):
   raw_value: A
@@ -139,19 +138,19 @@ def __init__(
     object.__setattr__(self, 'raw_value', value)
 
     if 'on_get_value' in type_vars and 'on_get_value' not in metadata:
-      metadata['get_value'] = getattr(type(self), 'on_get_value')
+      metadata['on_get_value'] = getattr(type(self), 'on_get_value')
 
     if 'on_set_value' in type_vars and 'on_set_value' not in metadata:
-      metadata['set_value'] = getattr(type(self), 'on_set_value')
+      metadata['on_set_value'] = getattr(type(self), 'on_set_value')
 
     if 'on_create_value' in type_vars and 'on_create_value' not in metadata:
-      metadata['create_value'] = getattr(type(self), 'on_create_value')
+      metadata['on_create_value'] = getattr(type(self), 'on_create_value')
 
     if 'on_add_axis' in type_vars and 'on_add_axis' not in metadata:
-      metadata['add_axis'] = getattr(type(self), 'on_add_axis')
+      metadata['on_add_axis'] = getattr(type(self), 'on_add_axis')
 
     if 'on_remove_axis' in type_vars and 'on_remove_axis' not in metadata:
-      metadata['remove_axis'] = getattr(type(self), 'on_remove_axis')
+      metadata['on_remove_axis'] = getattr(type(self), 'on_remove_axis')
 
     vars_self['_var_metadata'] = metadata
     # run create_value hooks
@@ -252,10 +251,12 @@ def __eq__(self, other: object) -> bool:
     return type(self) is type(other) and vars(other) == vars(self)
 
   @tp.overload
-  def replace(self, value: B, **kwargs) -> Variable[B]: ...
+  def replace(self, value: B, **kwargs) -> Variable[B]:
+    ...
 
   @tp.overload
-  def replace(self, **kwargs) -> Variable[A]: ...
+  def replace(self, **kwargs) -> Variable[A]:
+    ...
 
   def replace(self, value: tp.Any = Missing, **kwargs) -> Variable[tp.Any]:
     if value is not Missing:
@@ -986,15 +987,14 @@ def split_flat_state(
         break
     else:
       raise ValueError(
-        'Non-exhaustive filters, got a non-empty remainder: '
-        f'{path} -> {value}.'
-        '\nUse `...` to match all remaining elements.'
+          'Non-exhaustive filters, got a non-empty remainder: '
+          f'{path} -> {value}.'
+          '\nUse `...` to match all remaining elements.'
       )
 
   return flat_states
 
 
-
 ###################################################
 ### Variable type/class <-> string name mapping ###
 ###################################################
@@ -1003,10 +1003,14 @@ def split_flat_state(
 VariableTypeCache: dict[str, tp.Type[Variable[tp.Any]]] = {}
 
 
-def variable_type_from_name(name: str) -> tp.Type[Variable[tp.Any]]:
+def variable_type_from_name(
+    name: str,
+    *,
+    base: type[Variable[tp.Any]] = Variable,
+) -> tp.Type[Variable[tp.Any]]:
   """Given a Linen-style collection name, get or create its NNX Variable class."""
   if name not in VariableTypeCache:
-    VariableTypeCache[name] = type(name, (Variable,), {})
+    VariableTypeCache[name] = type(name, (base,), {})
   return VariableTypeCache[name]
 
 
@@ -1023,21 +1027,56 @@ def variable_name_from_type(typ: tp.Type[Variable[tp.Any]]) -> str:
       'Name {name} is already registered in the registry as {VariableTypeCache[name]}. '
       'It cannot be linked with this type {typ}.'
     )
-  register_variable_name_type_pair(name, typ)
+  register_variable_name(name, typ)
   return name
 
 
-def register_variable_name_type_pair(name, typ, overwrite = False):
+class _Missing:
+  pass
+
+
+_MISSING = _Missing()
+
+
+@tp.overload
+def register_variable_name(
+    name: str,
+    typ: type[Variable[tp.Any]],
+    *,
+    overwrite: bool = False,
+) -> type[Variable[tp.Any]]:
+  ...
+
+
+@tp.overload
+def register_variable_name(
+    name: str,
+    *,
+    overwrite: bool = False,
+) -> tp.Callable[[type[Variable[tp.Any]]], type[Variable[tp.Any]]]:
+  ...
+
+
+def register_variable_name(
+    name: str,
+    typ: type[Variable[A]] | _Missing = _MISSING,
+    *,
+    overwrite=False,
+) -> type[Variable[A]] | tp.Callable[[type[Variable[A]]], type[Variable[A]]]:
   """Register a pair of Linen collection name and its NNX type."""
+  if typ is _MISSING:
+    return partial(register_variable_name, name, overwrite=overwrite)
+  typ = tp.cast(type[Variable[A]], typ)
   if not overwrite and name in VariableTypeCache:
     raise ValueError(f'Name {name} already mapped to type {VariableTypeCache[name]}. '
-                     'To overwrite, call register_variable_name_type_pair() with `overwrite=True`.')
+                     'To overwrite, call set_variable_name() with `overwrite=True`.')
   VariableTypeCache[name] = typ
+  return typ
 
 
 # add known variable type names
-register_variable_name_type_pair('params', Param)
-register_variable_name_type_pair('batch_stats', BatchStat)
-register_variable_name_type_pair('cache', Cache)
-register_variable_name_type_pair('intermediates', Intermediate)
-register_variable_name_type_pair('perturbations', Perturbation)
\ No newline at end of file
+register_variable_name('params', Param)
+register_variable_name('batch_stats', BatchStat)
+register_variable_name('cache', Cache)
+register_variable_name('intermediates', Intermediate)
+register_variable_name('perturbations', Perturbation)


Test patch:
diff --git a/tests/nnx/bridge/wrappers_test.py b/tests/nnx/bridge/wrappers_test.py
index 8c5849126..bb73f0b5c 100644
--- a/tests/nnx/bridge/wrappers_test.py
+++ b/tests/nnx/bridge/wrappers_test.py
@@ -174,7 +174,6 @@ def create_sharded_nnx_module(x):
     assert nnx_model.bias.value.sharding.is_equivalent_to(
       jax.sharding.NamedSharding(self.mesh, jax.sharding.PartitionSpec('out',)), ndim=1)
 
-
   def test_linen_to_nnx_state_structure_consistency(self):
     class LinenInner(nn.Module):
       dout: int
@@ -191,8 +190,10 @@ def __call__(self, x):
         b = self.variable('bias', 'b', nn.initializers.zeros_init(), None, (1, self.dout))
         return dot(x) + b.value
 
-    class Bias(nnx.Variable): pass
-    nnx.register_variable_name_type_pair('bias', Bias)
+    @nnx.register_variable_name('bias')
+    class Bias(nnx.Variable):
+      pass
+
     class NNXMiddle(nnx.Module):
       def __init__(self, dout: int, *, rngs: nnx.Rngs):
         self.dot = bridge.ToNNX(LinenInner(dout), rngs=rngs)
@@ -292,8 +293,10 @@ def __call__(self, x):
     assert y.shape == (2, 3)
 
   def test_nnx_to_linen_mutable(self):
-    class Count(nnx.Variable): pass
-    nnx.register_variable_name_type_pair('Count', Count, overwrite=True)
+
+    @nnx.register_variable_name('Count', overwrite=True)
+    class Count(nnx.Variable):
+      pass
 
     class Counter(nnx.Module):
       def __init__(self):
@@ -310,8 +313,10 @@ def __call__(self):
     _ = model.apply(variables | updates)
 
   def test_nnx_to_linen_mutated_static_data(self):
-    class Count(nnx.Variable): pass
-    nnx.register_variable_name_type_pair('Count', Count, overwrite=True)
+
+    @nnx.register_variable_name('Count', overwrite=True)
+    class Count(nnx.Variable):
+      pass
 
     class Counter(nnx.Module):
       def __init__(self):
@@ -383,8 +388,10 @@ def __init__(self, din: int, dout: int, *, rngs: nnx.Rngs):
       def __call__(self, x):
         return self.dropout(x @ self.w)
 
-    class Bias(nnx.Variable): pass
-    nnx.register_variable_name_type_pair('bias', Bias, overwrite=True)
+    @nnx.register_variable_name('bias', overwrite=True)
+    class Bias(nnx.Variable):
+      pass
+
     class NNXMiddle(nnx.Module):
       def __init__(self, din: int, dout: int, *, rngs: nnx.Rngs):
         self.dot = NNXInner(din, dout, rngs=rngs)
@@ -420,7 +427,6 @@ def get_weights(variables):
     self.assertEqual(jax.tree.structure(from_top_weights),
                      jax.tree.structure(from_middle_weights))
 
-
   ############################
   ### Hybrid mix-and-match ###
   ############################
diff --git a/tests/nnx/metrics_test.py b/tests/nnx/metrics_test.py
index 951734f93..5d4a1229a 100644
--- a/tests/nnx/metrics_test.py
+++ b/tests/nnx/metrics_test.py
@@ -17,6 +17,7 @@
 from flax import nnx
 import jax
 import jax.numpy as jnp
+import numpy as np
 
 
 class TestMetrics(parameterized.TestCase):
@@ -124,6 +125,36 @@ def test_multimetric(self):
     self.assertTrue(jnp.isnan(values['accuracy']))
     self.assertTrue(jnp.isnan(values['loss']))
 
+  def test_binary_classification_accuracy(self):
+    logits = jnp.array([0.4, 0.7, 0.2, 0.6])
+    labels = jnp.array([0, 1, 1, 1])
+    logits2 = jnp.array([0.1, 0.9, 0.8, 0.3])
+    labels2 = jnp.array([0, 1, 1, 0])
+    accuracy = nnx.metrics.Accuracy(threshold=0.5)
+    accuracy.update(logits=logits, labels=labels)
+    self.assertEqual(accuracy.compute(), 0.75)
+    accuracy.update(logits=logits2, labels=labels2)
+    self.assertEqual(accuracy.compute(), 0.875)
+
+  @parameterized.parameters(
+    {
+      'logits': np.array([[[0.0, 0.0], [0.0, 0.0]], [[0.0, 0.0], [0.0, 0.0]]]),
+      'labels': np.array([0, 0, 0, 0]),
+      'threshold': None,
+      'error_msg': 'For multi-class classification'
+    },
+    {
+      'logits': np.array([0.0, 0.0, 0.0, 0.0]),
+      'labels': np.array([[0, 0], [0, 0]]),
+      'threshold': 0.5,
+      'error_msg': 'For binary classification'
+    }
+  )
+  def test_accuracy_dims(self, logits, labels, threshold, error_msg):
+    accuracy = nnx.metrics.Accuracy(threshold=threshold)
+    with self.assertRaisesRegex(ValueError, error_msg):
+      accuracy.update(logits=logits, labels=labels)
+
 
 if __name__ == '__main__':
   absltest.main()