google / google/metrax

`FBetaScore` missing `merge` method causes NotImplementedError with `metrax.nnx` wrapper

Open
#131 2 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
64
Forks
18
Avg merge
4h 46m
Merged PRs (30d)
1

Description

## Description
The `FBetaScore` metric in `metrax.classification_metrics` does not implement the `merge` method. This causes a `NotImplementedError` when using the metric with the `metrax.nnx` wrapper (e.g., `metrax.nnx.FBetaScore`), as the wrapper's `update` method relies on `merge` to combine metrics.

It appears the `merge` method is commented out in the source code of `metrax/classification_metrics.py`.
https://github.com/google/metrax/blob/4ff6ccfc72bcd220be2dbbf8bd92a92f609ef499/src/metrax/classification_metrics.py#L677-L695

## Minimal Reproduction
```python
import metrax.nnx
import jax.numpy as jnp
import jax.random

# Setup dummy data
predictions = jax.random.normal(jax.random.PRNGKey(0), (3,))
labels = jnp.arange(3) % 2

# Initialize and update metric
f1_metric = metrax.nnx.FBetaScore()
f1_metric.update(predictions=predictions, labels=labels) # Raises NotImplementedError
```

## Traceback
```text
Traceback (most recent call last):
File "repro.py", line 10, in
f1_metric.update(predictions=predictions, labels=labels)
File ".../site-packages/metrax/nnx/nnx_wrapper.py", line 31, in update
self.clu_metric = self.clu_metric.merge(other_clu_metric)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File ".../site-packages/clu/metrics.py", line 148, in merge
raise NotImplementedError("Must override merge()")
NotImplementedError: Must override merge()
```

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.