Incorrect type annotation on Multimetric.compute()

[gabbard]

The type signature of Multimetric.compute seems to be incorrect. It is currently:

def compute(self) -> dict[str, Metric]:

Its implementation is

return {
        f'{metric_name}': getattr(self, metric_name).compute()
        for metric_name in self._metric_names
    }

meaning that the second type parameter of the return type should be the return type of Metric.compute(). Metric.compute itself does not specify a return type, or even if it returns anything. The provided implementations (besides Multimetric itself!) return jax.Array and nnx.training.metrics.Statistics objects, suggesting (reasonably) that no restriction is placed on the return.

It looks like:

• the return type of Metric.compute should be either Any or Optional[Any] depending on whether Metric.compute() has to return something (I would assume it does?).
• the return type of MultiMetric.compute should be dict[str, Any] or dict[str, Optional[Any]).

System information

• Flax, jax, jaxlib versions (obtain with pip show flax jax jaxlib: 0.10.2

[cgarciae]

Thanks @gabbard, I think it should be dict[str, Any]. Created #4485 to fix this.