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

[steventango]

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.

metrax/src/metrax/classification_metrics.py

Lines 677 to 695 in 4ff6ccf

	"""
	This function is currently unused as the 'from_model_output' function can handle the whole
	dataset without needing to split and merge them. I'm leaving this here for now incase we want to
	repurpose this or need to change something that requires this function's use again. This function would need
	to be reworked for it to work with the current implementation of this class.
	"""
	# # Merge datasets together
	# def merge(self, other: 'FBetaScore') -> 'FBetaScore':
	#
	# # Check if the incoming beta is the same value as the current beta
	# if other.beta == self.beta:
	# return type(self)(
	# true_positives = self.true_positives + other.true_positives,
	# false_positives = self.false_positives + other.false_positives,
	# false_negatives = self.false_negatives + other.false_negatives,
	# beta=self.beta,
	# )
	# else:
	# raise ValueError('The "Beta" values between the two are not equal.')

Minimal Reproduction

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

Traceback (most recent call last):
  File "repro.py", line 10, in <module>
    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()

[2468cool]

For some frustrating reason my environment won't recognize the .nnx for metrax so it won't let me replicate the error. Idk if just uncommenting it would fix it but I would think not since the class isn't built with the merge function in mind at this time.

I remember why I did what I did though. For large batch sizes (greater than 1), keras's FBeta gives an array output based on the size. Also, it required a 2D array or would throw an error. They way I originally had made the metrax version only worked for batch size 1 and had to restructure the data to fit. During the process I asked if we should instead make it so that the metrax version gives a matching output (answer size dependent on entires in a batch) to the keras version. This changed how I approached calculating the score for metrax. I had to make it handle batches with more than one entry in a way that reflected keras's output. The way I did this made it so the merge function was not needed.

I'm doing some testing on my own to see if I can fix this.

[2468cool]

It looks like a possible solution would be to make sure the user gives a 2D array for it to process like how Keras checks for a 2D array. From what I can see on my end this allows for the merge function to be used again.