Internal change

PiperOrigin-RevId: 527429713

Author: josephlr

clu/metrics.py

@@ -56,7 +56,7 @@ def evaluate(model, p_variables, test_ds):
     return ms.unreplicate().compute()
 """
 from collections.abc import Callable, Mapping, Sequence
-from typing import Any, Optional
+from typing import Any, Optional, TypeVar
 
 from absl import logging
 
@@ -435,6 +435,9 @@ def _check_reduction_counter_ndim(reduction_counter: _ReductionCounter):
         f"call a flax.jax_utils.unreplicate() or a Collections.reduce()?")
 
 
+C = TypeVar("C", bound="Collection")
+
+
 @flax.struct.dataclass
 class Collection:
   """Updates a collection of `Metric` from model outputs.
@@ -512,7 +515,7 @@ class MyMetrics(metrics.Collection):
     return collection_class(_reduction_counter=counter, **metrics)
 
   @classmethod
-  def empty(cls) -> "Collection":
+  def empty(cls: type[C]) -> C:
     return cls(
         _reduction_counter=_ReductionCounter(jnp.array(1, dtype=jnp.int32)),
         **{
@@ -521,7 +524,7 @@ def empty(cls) -> "Collection":
         })
 
   @classmethod
-  def _from_model_output(cls, **kwargs) -> "Collection":
+  def _from_model_output(cls: type[C], **kwargs) -> C:
     """Creates a `Collection` from model outputs."""
     return cls(
         _reduction_counter=_ReductionCounter(jnp.array(1, dtype=jnp.int32)),
@@ -531,7 +534,7 @@ def _from_model_output(cls, **kwargs) -> "Collection":
         })
 
   @classmethod
-  def single_from_model_output(cls, **kwargs) -> "Collection":
+  def single_from_model_output(cls: type[C], **kwargs) -> C:
     """Creates a `Collection` from model outputs.
 
     Note: This function should only be called when metrics are collected in a
@@ -546,9 +549,7 @@ def single_from_model_output(cls, **kwargs) -> "Collection":
     return cls._from_model_output(**kwargs)
 
   @classmethod
-  def gather_from_model_output(cls,
-                               axis_name="batch",
-                               **kwargs) -> "Collection":
+  def gather_from_model_output(cls: type[C], axis_name="batch", **kwargs) -> C:
     """Creates a `Collection` from model outputs in a distributed setting.
 
     Args:
@@ -563,14 +564,14 @@ def gather_from_model_output(cls,
     return jax.lax.all_gather(
         cls._from_model_output(**kwargs), axis_name=axis_name).reduce()
 
-  def merge(self, other: "Collection") -> "Collection":
+  def merge(self: C, other: C) -> C:
     """Returns `Collection` that is the accumulation of `self` and `other`."""
     return type(self)(**{
         metric_name: metric.merge(getattr(other, metric_name))
         for metric_name, metric in vars(self).items()
     })
 
-  def reduce(self) -> "Collection":
+  def reduce(self: C) -> C:
     """Reduces the collection by calling `Metric.reduce()` on each metric.
 
     The primary use case is to reduce collection that was gathered
@@ -620,7 +621,7 @@ def compute_values(self) -> dict[str, clu.values.Value]:
         if metric_name != "_reduction_counter"
     }
 
-  def unreplicate(self) -> "Collection":
+  def unreplicate(self: C) -> C:
     """Short-hand for `flax.jax_utils.unreplicate(self)`.
 
     The collection should be gathered and `reduce`d inside pmap,

clu/metrics_test.py

@@ -86,6 +86,8 @@ def setUp(self):
         dict(mask=mask, **model_output)
         for mask, model_output in zip(masks, self.model_outputs))
 
+    self.count = 4
+    self.count_masked = 2
     self.results = {
         "train_accuracy": 0.75,
         "learning_rate": 0.01,
@@ -367,18 +369,23 @@ def test_collection_create_collection(self):
       ("_masked", True),
   )
   def test_collection_single(self, masked):
-
+    @jax.jit
     def compute_collection(model_outputs):
       collection = Collection.empty()
       for model_output in model_outputs:
         update = Collection.single_from_model_output(**model_output)
         collection = collection.merge(update)
-      return collection.compute()
+      return collection
 
+    model_outputs = self.model_outputs_masked if masked else self.model_outputs
+    collection = compute_collection(model_outputs)
     chex.assert_trees_all_close(
-        jax.jit(compute_collection)(
-            self.model_outputs_masked if masked else self.model_outputs),
-        self.results_masked if masked else self.results)
+        collection.compute(), self.results_masked if masked else self.results
+    )
+    self.assertEqual(
+        collection.train_accuracy.count,
+        self.count_masked if masked else self.count,
+    )
 
   @parameterized.named_parameters(
       ("", False),