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.

Discussion:
mattbahr: > Hey @mattbahr, thanks for doing this! Apart from the two other comments, I think `threshold` should be `None` by default and the user should explicitly opt-in to binary classification. This way we can have stronger shape checks for each case.

Good idea! Just made the requested changes. Let me know if there's anything else I can do to improve the implementation.

Patch:
diff --git a/.github/workflows/flax_publish.yml b/.github/workflows/flax_publish.yml
index 19aa2dd5..f688e8fc 100644
--- a/.github/workflows/flax_publish.yml
+++ b/.github/workflows/flax_publish.yml
@@ -8,37 +8,24 @@ on:
     types: [published]
 
 jobs:
-  build_package:
-    name: Build package
-    runs-on: ubuntu-latest
-    steps:
-      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683  # v4.2.2
-      - name: Build package
-        run: pipx run build
-      - name: List files
-        run: ls -l dist/
-      - uses: actions/upload-artifact@65c4c4a1ddee5b72f698fdd19549f0f0fb45cf08  # v4.6.0
-        with:
-          name: distribution
-          path: |
-            dist/*.tar.gz
-            dist/*.whl
+  deploy:
 
-  upload_pypi:
-    name: Release & Upload to PyPI
-    # Only publish release to PyPI when a github release is created.
-    if: github.event_name == 'release' && github.event.action == 'published'
-    needs: build_package
     runs-on: ubuntu-latest
-    environment: release
-    permissions:
-      id-token: write
+
     steps:
-      - uses: actions/download-artifact@fa0a91b85d4f404e444e00e005971372dc801d16  # v4.1.8
-        with:
-          name: distribution
-          path: dist
-      - name: List files
-        run: ls -l dist/
-      - name: Publish package distributions to PyPI
-        uses: pypa/gh-action-pypi-publish@76f52bc884231f62b9a034ebfe128415bbaabdfc  # v1.12.4
+    - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
+    - name: Set up Python
+      uses: actions/setup-python@0b93645e9fea7318ecaed2b359559ac225c90a2b # v5.3.0
+      with:
+        python-version: '3.x'
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install setuptools build wheel twine
+    - name: Build and publish
+      env:
+        TWINE_USERNAME: __token__
+        TWINE_PASSWORD: ${{ secrets.PYPI_PASSWORD }}
+      run: |
+        python -m build
+        twine upload dist/*
diff --git a/README.md b/README.md
index 42321ead..63b5592f 100644
--- a/README.md
+++ b/README.md
@@ -157,7 +157,7 @@ To cite this repository:
   author = {Jonathan Heek and Anselm Levskaya and Avital Oliver and Marvin Ritter and Bertrand Rondepierre and Andreas Steiner and Marc van {Z}ee},
   title = {{F}lax: A neural network library and ecosystem for {JAX}},
   url = {http://github.com/google/flax},
-  version = {0.10.3},
+  version = {0.10.2},
   year = {2024},
 }
 ```
diff --git a/docs/guides/training_techniques/batch_norm.rst b/docs/guides/training_techniques/batch_norm.rst
index 55613a95..65c3a67e 100644
--- a/docs/guides/training_techniques/batch_norm.rst
+++ b/docs/guides/training_techniques/batch_norm.rst
@@ -141,8 +141,8 @@ collection.
 Modifying ``flax.linen.apply``
 ******************************
 
-When using :meth:`flax.linen.apply <flax.linen.apply>` to run your model with the ``train=True``
-argument (that is, you have ``use_running_average=False`` in the call to ``BatchNorm``), you
+When using :meth:`flax.linen.apply <flax.linen.apply>` to run your model with the ``train==True``
+argument (that is, you have ``use_running_average==False`` in the call to ``BatchNorm``), you
 need to consider the following:
 
 * ``batch_stats`` must be passed as an input variable.
diff --git a/flax/nnx/training/metrics.py b/flax/nnx/training/metrics.py
index eca41bd4..8a0ffabe 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/version.py b/flax/version.py
index 2f12c4c3..bdba069c 100644
--- a/flax/version.py
+++ b/flax/version.py
@@ -13,4 +13,4 @@
 # limitations under the License.
 
 """Current Flax version at head on Github."""
-__version__ = '0.10.3'
+__version__ = '0.10.2'


Test patch:
diff --git a/tests/nnx/metrics_test.py b/tests/nnx/metrics_test.py
index 951734f9..3b3ee788 100644
--- a/tests/nnx/metrics_test.py
+++ b/tests/nnx/metrics_test.py
@@ -124,6 +124,36 @@ class TestMetrics(parameterized.TestCase):
     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': jnp.array([[[0.0, 0.0], [0.0, 0.0]], [[0.0, 0.0], [0.0, 0.0]]]),
+      'labels': jnp.array([0, 0, 0, 0]),
+      'threshold': None,
+      'error_msg': 'For multi-class classification'
+    },
+    {
+      'logits': jnp.array([0.0, 0.0, 0.0, 0.0]),
+      'labels': jnp.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()