Introduce MinSummary and MaxSummary.

GitOrigin-RevId: 8d3511d

Author: ds-hwang

axlearn/common/metrics.py

@@ -50,6 +50,48 @@ def accumulate(self, other: Summary) -> Summary:
         return self + other
 
 
+class MinSummary(Summary):
+    """A summary that computes the minimum value across tensor elements."""
+
+    _value: Tensor
+
+    def value(self) -> Tensor:
+        return self._value
+
+    def validate(self):
+        val = self._value
+        if not isinstance(val, Tensor):
+            raise ValueError(f"MinSummary value must be a Tensor, but got {str(type(val))}.")
+        if val.ndim >= 1:
+            raise ValueError(f"MinSummary value must be a scalar, but got {val.ndim=}.")
+
+    def accumulate(self, other: Summary) -> Summary:
+        if not isinstance(other, MinSummary):
+            raise TypeError(f"Expected MinSummary, got {type(other)}.")
+        return MinSummary(jnp.minimum(self.value(), other.value()))
+
+
+class MaxSummary(Summary):
+    """A summary that computes the maximum value across tensor elements."""
+
+    _value: Tensor
+
+    def value(self) -> Tensor:
+        return self._value
+
+    def validate(self):
+        val = self._value
+        if not isinstance(val, Tensor):
+            raise ValueError(f"MaxSummary value must be a Tensor, but got {str(type(val))}.")
+        if val.ndim >= 1:
+            raise ValueError(f"MaxSummary value must be a scalar, but got {val.ndim=}.")
+
+    def accumulate(self, other: Summary) -> Summary:
+        if not isinstance(other, MaxSummary):
+            raise TypeError(f"Expected MaxSummary, got {type(other)}.")
+        return MaxSummary(jnp.maximum(self.value(), other.value()))
+
+
 class MetricAccumulator(Configurable):
     """A MetricAccumulator is used during evaluation to accumulate metrics across batches."""
 

axlearn/common/metrics_test.py

@@ -2,14 +2,16 @@
 
 """Tests for metrics.py."""
 
+import contextlib
 from typing import Optional
 
 import chex
 import jax
 import jax.numpy as jnp
-from absl.testing import parameterized
+from absl.testing import absltest, parameterized
 
 from axlearn.common import metrics, summary, test_utils, utils
+from axlearn.common.metrics import MaxSummary, MinSummary
 from axlearn.common.module import Summable
 
 
@@ -55,6 +57,29 @@ def test_metric_accumulator(self):
         expected = jax.tree_util.tree_leaves(expected)
         chex.assert_trees_all_close(result, expected)
 
+    @parameterized.parameters(
+        dict(cls=metrics.MinSummary, expected=-10),
+        dict(cls=metrics.MaxSummary, expected=10),
+    )
+    def test_metric_min_max_accumulator(self, cls, expected):
+        acc = metrics.MetricAccumulator.default_config().instantiate()
+        summaries = [
+            dict(foo=cls(jnp.array(5))),
+            dict(foo=cls(jnp.array(-10))),
+            dict(foo=cls(jnp.array(10))),
+        ]
+
+        summaries_copy = jax.tree.map(lambda x: x, summaries)
+        for s in summaries_copy:
+            acc.update(s)
+        result = acc.summaries()
+        expected = dict(foo=cls(jnp.array(expected)))
+
+        chex.assert_trees_all_equal_structs(result, expected)
+        result = jax.tree_util.tree_leaves(result)
+        expected = jax.tree_util.tree_leaves(expected)
+        self.assertEqual(result, expected)
+
     def test_flatten_unflatten_metric_accumulator(self):
         acc = metrics.MetricAccumulator.default_config().instantiate()
         summaries = [
@@ -128,3 +153,45 @@ def add(weight):
 
         # Test isinstance check.
         self.assertIsInstance(metrics.WeightedScalar(1.0, 1.0), Summable)
+
+
+class MinSummaryTest(test_utils.TestCase):
+    @parameterized.parameters(
+        (jnp.array(1), jnp.array(1)),
+        (jnp.array([1, 2]), ValueError("MinSummary value must be a scalar, but got val.ndim=1.")),
+        (jnp.array([[1, 2]]), ValueError("MinSummary value must be a scalar, but got val.ndim=2.")),
+        (1, ValueError("MinSummary value must be a Tensor, but got <class 'int'>.")),
+    )
+    def test_min_summary(self, value, expected):
+        min_summary = MinSummary(value)
+        if isinstance(expected, ValueError):
+            ctx = self.assertRaisesRegex(ValueError, expected.args[0])
+        else:
+            ctx = contextlib.nullcontext()
+        with ctx:
+            min_summary.validate()
+            new_summary = min_summary.accumulate(MinSummary(jnp.array(10)))
+            self.assertEqual(new_summary.value(), expected)
+
+
+class MaxSummaryTest(test_utils.TestCase):
+    @parameterized.parameters(
+        (jnp.array(1), jnp.array(1)),
+        (jnp.array([1, 2]), ValueError("MaxSummary value must be a scalar, but got val.ndim=1.")),
+        (jnp.array([[1, 2]]), ValueError("MaxSummary value must be a scalar, but got val.ndim=2.")),
+        (1, ValueError("MaxSummary value must be a Tensor, but got <class 'int'>.")),
+    )
+    def test_min_summary(self, value, expected):
+        max_summary = MaxSummary(value)
+        if isinstance(expected, ValueError):
+            ctx = self.assertRaisesRegex(ValueError, expected.args[0])
+        else:
+            ctx = contextlib.nullcontext()
+        with ctx:
+            max_summary.validate()
+            new_summary = max_summary.accumulate(MaxSummary(jnp.array(-10)))
+            self.assertEqual(new_summary.value(), expected)
+
+
+if __name__ == "__main__":
+    absltest.main()