clu.metrics: Fix type annotations for Metric

Similar to the previous change for `Collection`, this allows the
methods and classmethods on `Metric` and its subclasses to be properly
typed. This allows us to remove the `disable=g-bare-generic` lint.

This also enables typechecking for `metrics_test.py`, which failed to
typecheck before this change.

To get `Metrics.from_fun()` to properly typecheck, we introduce a `FromFunCallable`
[Callback Protocol](https://mypy.readthedocs.io/en/stable/protocols.html#callback-protocols)
to encode the type of "a function that takes in any number of `ArrayLike`
types and outputs `Array` or a dictionary with `Array` keys".

Finally, this change uses `Protocol` instead of `abc.ABC` for `Value`,
so that we can call `v.value` on abstract `Value` types.

PiperOrigin-RevId: 529495726

Author: josephlr

clu/metrics.py

@@ -55,8 +55,9 @@ def evaluate(model, p_variables, test_ds):
       ms = p_eval_step(ms, model, p_variables, inputs, labels)
     return ms.unreplicate().compute()
 """
-from collections.abc import Callable, Mapping, Sequence
-from typing import Any, Optional, TypeVar
+from __future__ import annotations
+from collections.abc import Mapping, Sequence
+from typing import Any, TypeVar, Protocol
 
 from absl import logging
 
@@ -67,6 +68,17 @@ def evaluate(model, p_variables, test_ds):
 import jax.numpy as jnp
 import numpy as np
 
+Array = jax.Array
+ArrayLike = jax.typing.ArrayLike
+
+
+class FromFunCallable(Protocol):
+  """The type of functions that can be passed to `Metrics.from_fun()`."""
+
+  def __call__(self, **kwargs: ArrayLike) -> Array | Mapping[str, Array]:
+    """Returns the argument/arguments passed to the base from_model_output()."""
+
+
 # TODO(b/200953513): Migrate away from logging imports (on module level)
 #                    to logging the actual usage. See b/200953513.
 
@@ -78,6 +90,9 @@ def _assert_same_shape(a: jnp.array, b: jnp.array):
     raise ValueError(f"Expected same shape: {a.shape} != {b.shape}")
 
 
+M = TypeVar("M", bound="Metric")
+
+
 class Metric:
   """Interface for computing metrics from intermediate values.
 
@@ -114,11 +129,11 @@ def compute(self):
   """
 
   @classmethod
-  def from_model_output(cls, *args, **kwargs) -> "Metric":
+  def from_model_output(cls: type[M], *args, **kwargs) -> M:
     """Creates a `Metric` from model outputs."""
     raise NotImplementedError("Must override from_model_output()")
 
-  def merge(self, other: "Metric") -> "Metric":
+  def merge(self: M, other: M) -> M:
     """Returns `Metric` that is the accumulation of `self` and `other`.
 
     Args:
@@ -141,23 +156,23 @@ def merge(self, other: "Metric") -> "Metric":
   # `_reduce_merge()` must be associative[1], otherwise we would get
   # different results when using different devices.
   # [1] https://en.wikipedia.org/wiki/Associative_property
-  def _reduce_merge(self, other: "Metric") -> "Metric":
+  def _reduce_merge(self: M, other: M) -> M:
     return self.merge(other)
 
   def compute(self) -> jnp.array:
     """Computes final metrics from intermediate values."""
     raise NotImplementedError("Must override compute()")
 
   @classmethod
-  def empty(cls) -> "Metric":
+  def empty(cls: type[M]) -> M:
     """Returns an empty instance (i.e. `.merge(Metric.empty())` is a no-op)."""
     raise NotImplementedError("Must override empty()")
 
   def compute_value(self) -> clu.values.Value:
     """Wraps compute() and returns a values.Value."""
     return clu.values.Scalar(self.compute())
 
-  def reduce(self) -> "Metric":
+  def reduce(self: M) -> M:
     """Reduces the metric along it first axis by calling `_reduce_merge()`.
 
     This function primary use case is to aggregate metrics collected across
@@ -173,7 +188,7 @@ def reduce(self) -> "Metric":
       reduced metric.
     """
 
-    def reduce_step(reduced: Metric, metric: Metric) -> tuple[Metric, None]:
+    def reduce_step(reduced: M, metric: M) -> tuple[M, None]:
       # pylint: disable-next=protected-access
       return reduced._reduce_merge(metric), None
 
@@ -183,7 +198,7 @@ def reduce_step(reduced: Metric, metric: Metric) -> tuple[Metric, None]:
     return jax.lax.scan(reduce_step, first, remainder)[0]
 
   @classmethod
-  def from_fun(cls, fun: Callable):  # pylint: disable=g-bare-generic
+  def from_fun(cls, fun: FromFunCallable):  # No way to annotate return type
     """Calls `cls.from_model_output` with the return value from `fun`.
 
     Returns a `Metric` derived from `cls` whose `.from_model_output` (1) calls
@@ -233,7 +248,7 @@ class FromFun(cls):
       """Wrapper Metric class that collects output after applying `fun`."""
 
       @classmethod
-      def from_model_output(cls, **model_output) -> Metric:
+      def from_model_output(cls: type[M], **model_output) -> M:
         mask = model_output.get("mask")
         output = fun(**model_output)
         if isinstance(output, Mapping) and "mask" in output:
@@ -266,7 +281,7 @@ def from_model_output(cls, **model_output) -> Metric:
     return FromFun
 
   @classmethod
-  def from_output(cls, name: str):  # pylint: disable=g-bare-generic
+  def from_output(cls, name: str):  # No way to annotate return type
     """Calls `cls.from_model_output` with model output named `name`.
 
     Synopsis:
@@ -295,7 +310,7 @@ class FromOutput(cls):
       """Wrapper Metric class that collects output named `name`."""
 
       @classmethod
-      def from_model_output(cls, **model_output) -> Metric:
+      def from_model_output(cls: type[M], **model_output) -> M:
         output = jnp.array(model_output[name])
         mask = model_output.get("mask")
         if mask is not None and (output.shape or [0])[0] != mask.shape[0]:
@@ -366,10 +381,10 @@ def merge(update):
   values: dict[str, tuple[np.ndarray, ...]]
 
   @classmethod
-  def empty(cls) -> "CollectingMetric":
+  def empty(cls) -> CollectingMetric:
     return cls(values={})
 
-  def merge(self, other: "CollectingMetric") -> "CollectingMetric":
+  def merge(self, other: CollectingMetric) -> CollectingMetric:
     values = {
         name: (*value, *other.values[name])
         for name, value in self.values.items()
@@ -384,25 +399,24 @@ def merge(self, other: "CollectingMetric") -> "CollectingMetric":
       return self
     return type(self)(jax.tree_map(np.asarray, values))
 
-  def reduce(self) -> "CollectingMetric":
+  def reduce(self) -> CollectingMetric:
     # Note that this is usually called from inside a `pmap()` via
     # `Collection.gather_from_model_output()` so we concatenate using jnp.
     return type(self)(
         {name: jnp.concatenate(values) for name, values in self.values.items()})
 
-  def compute(self) -> dict[str, np.ndarray]:
+  def compute(self):  # No return type annotation, so subclasses can override
     return {k: np.concatenate(v) for k, v in self.values.items()}
 
   @classmethod
-  def from_outputs(cls, names: Sequence[str]):
+  def from_outputs(cls, names: Sequence[str]) -> type[CollectingMetric]:
     """Returns a metric class that collects all model outputs named `names`."""
 
     @flax.struct.dataclass
     class FromOutputs(cls):  # pylint:disable=missing-class-docstring
 
       @classmethod
-      def from_model_output(cls, **model_output) -> Metric:
-
+      def from_model_output(cls: type[M], **model_output) -> M:
         def make_array(value):
           value = jnp.array(value)
           # Can't jnp.concatenate() scalars, promote to shape=(1,) in that case.
@@ -420,10 +434,10 @@ class _ReductionCounter(Metric):
   value: jnp.array
 
   @classmethod
-  def empty(cls):
+  def empty(cls) -> _ReductionCounter:
     return cls(value=jnp.array(1, jnp.int32))
 
-  def merge(self, other: "_ReductionCounter") -> "_ReductionCounter":
+  def merge(self, other: _ReductionCounter) -> _ReductionCounter:
     return _ReductionCounter(self.value + other.value)
 
 
@@ -461,7 +475,7 @@ class Metrics(Collection):
   _reduction_counter: _ReductionCounter
 
   @classmethod
-  def create(cls, **metrics: type[Metric]) -> type["Collection"]:
+  def create(cls, **metrics: type[Metric]) -> type[Collection]:
     """Handy short-cut to define a `Collection` inline.
 
     Instead declaring a `Collection` dataclass:
@@ -487,7 +501,7 @@ class MyMetrics(metrics.Collection):
         type("_InlineCollection", (Collection,), {"__annotations__": metrics}))
 
   @classmethod
-  def create_collection(cls, **metrics: Metric) -> "Collection":
+  def create_collection(cls, **metrics: Metric) -> Collection:
     """Creates a custom collection object with fields metrics.
 
     This object will be an instance of custom subclass of `Collection` with
@@ -650,10 +664,12 @@ class LastValue(Metric):
   total: jnp.array
   count: jnp.array
 
-  def __init__(self, total: Optional[jnp.array] = None,
-               count: Optional[jnp.array] = None,
-               value: Optional[jnp.array] = None,
-               ):
+  def __init__(
+      self,
+      total: jnp.array | None = None,
+      count: jnp.array | None = None,
+      value: jnp.array | None = None,
+  ):
     """Constructor which supports keyword argument value as initializer.
 
     If  "value" is provided, then  "total" should *not* be provided.
@@ -673,26 +689,25 @@ def __init__(self, total: Optional[jnp.array] = None,
     object.__setattr__(self, "count", count)
 
   @classmethod
-  def empty(cls):
+  def empty(cls) -> LastValue:
     return cls(jnp.array(0, jnp.float32), count=jnp.array(0, jnp.int32))
 
   @classmethod
-  def from_model_output(cls,
-                        value: jnp.array,
-                        mask: Optional[jnp.array] = None,
-                        **_) -> Metric:
+  def from_model_output(
+      cls, value: jnp.array, mask: jnp.array | None = None, **_
+  ) -> LastValue:
     if mask is None:
       mask = jnp.ones((value.shape or [()])[0])
     return cls(
         total=jnp.where(mask, value, jnp.zeros_like(value)).sum(),
         count=mask.sum().astype(jnp.int32),
     )
 
-  def merge(self, other: "LastValue") -> "LastValue":
+  def merge(self, other: LastValue) -> LastValue:
     _assert_same_shape(self.value, other.value)
     return other
 
-  def _reduce_merge(self, other: "LastValue") -> "LastValue":
+  def _reduce_merge(self, other: LastValue) -> LastValue:
     # We need to average during reduction.
     _assert_same_shape(self.total, other.total)
     return type(self)(
@@ -730,14 +745,13 @@ class Average(Metric):
   count: jnp.array
 
   @classmethod
-  def empty(cls) -> Metric:
+  def empty(cls) -> Average:
     return cls(total=jnp.array(0, jnp.float32), count=jnp.array(0, jnp.int32))
 
   @classmethod
-  def from_model_output(cls,
-                        values: jnp.array,
-                        mask: Optional[jnp.array] = None,
-                        **_) -> Metric:
+  def from_model_output(
+      cls, values: jnp.array, mask: jnp.array | None = None, **_
+  ) -> Average:
     if values.ndim == 0:
       values = values[None]
     if mask is None:
@@ -760,7 +774,7 @@ def from_model_output(cls,
                         jnp.zeros_like(values, dtype=jnp.int32)).sum(),
     )
 
-  def merge(self, other: "Average") -> "Average":
+  def merge(self, other: Average) -> Average:
     _assert_same_shape(self.total, other.total)
     return type(self)(
         total=self.total + other.total,
@@ -783,17 +797,16 @@ class Std(Metric):
   count: jnp.array
 
   @classmethod
-  def empty(cls):
+  def empty(cls) -> Std:
     return cls(
         total=jnp.array(0, jnp.float32),
         sum_of_squares=jnp.array(0, jnp.float32),
         count=jnp.array(0, jnp.int32))
 
   @classmethod
-  def from_model_output(cls,
-                        values: jnp.array,
-                        mask: Optional[jnp.array] = None,
-                        **_) -> Metric:
+  def from_model_output(
+      cls, values: jnp.array, mask: jnp.array | None = None, **_
+  ) -> Std:
     if values.ndim == 0:
       values = values[None]
     utils.check_param(values, ndim=1)
@@ -805,7 +818,7 @@ def from_model_output(cls,
         count=mask.sum(),
     )
 
-  def merge(self, other: "Std") -> "Std":
+  def merge(self, other: Std) -> Std:
     _assert_same_shape(self.total, other.total)
     return type(self)(
         total=self.total + other.total,

clu/values.py

@@ -17,9 +17,8 @@
 A Metric should return one of the following types when compute() is called.
 """
 
-import abc
 import dataclasses
-from typing import Any, Union
+from typing import Any, Union, Protocol, runtime_checkable
 
 import jax.numpy as jnp
 import numpy as np
@@ -28,13 +27,14 @@
 ScalarType = Union[int, float, np.number, np.ndarray, jnp.ndarray]
 
 
-class Value(abc.ABC):
+@runtime_checkable
+class Value(Protocol):
   """Class defining available metric computation return values.
 
   Types mirror those available in MetricWriter. See
   clu/metric_writers/interface.py
   """
-  pass
+  value: Any
 
 
 @dataclasses.dataclass