add static_field support to AutoState

PiperOrigin-RevId: 688944515

kauldron/metrics/__init__.py

@@ -17,6 +17,7 @@
 # pylint: disable=g-importing-member
 from kauldron.metrics.auto_state import AutoState
 from kauldron.metrics.auto_state import concat_field
+from kauldron.metrics.auto_state import static_field
 from kauldron.metrics.auto_state import sum_field
 from kauldron.metrics.auto_state import truncate_field
 from kauldron.metrics.base import Metric

kauldron/metrics/auto_state.py

@@ -22,6 +22,7 @@
 from typing import Any, Literal, Self, TypeAlias, TypeVar
 
 import jax
+from kauldron import kontext
 from kauldron.metrics import base_state
 from kauldron.metrics.base_state import EMPTY  # pylint: disable=g-importing-member
 from kauldron.typing import Array  # pylint: disable=g-multiple-import
@@ -137,6 +138,32 @@ def compute(self: _SelfT) -> _SelfT:
 # TODO(klausg): overloaded type annotation similar to dataclasses.field?
 
 
+def static_field(default: Any = dataclasses.MISSING, **kwargs):
+  """Define an AutoState static field.
+
+  Static fields are not merged, and instead are checked for equality during
+  `merge`. They are also not pytree nodes, so they are not touched by jax
+  transforms (but can lead to recompilation if changed).
+  These can be useful to store some parameters of the metric, e.g. the number
+  of elements to keep. Note that static fields are rarely needed, since it is
+  usually better to define static params in the corresponding metric and access
+  them through the `parent` field.
+
+  Args:
+    default: The default value of the field.
+    **kwargs: Additional arguments to pass to the dataclasses.field.
+
+  Returns:
+    A dataclasses.Field instance with additional metadata that marks this field
+    as a static field.
+  """
+  metadata = kwargs.pop("metadata", {})
+  metadata = metadata | {
+      "pytree_node": False,
+  }
+  return dataclasses.field(default=default, metadata=metadata, **kwargs)
+
+
 def sum_field(
     *,
     default: Any = dataclasses.MISSING,
@@ -331,7 +358,7 @@ class _Truncate(_FieldMerger):
   Attributes:
     axis: The axis along which to concatenate the two arrays. Defaults to 0.
     num_field: The name of the field that contains the number of elements to
-      keep.
+      keep. Can be any valid kontext.Path (e.g. "parent.num_images").
   """
 
   axis: int | None = 0
@@ -343,7 +370,7 @@ def merge(
       v2: ArrayOrEmpty,
       state: base_state.State,
   ) -> ArrayOrEmpty:
-    num = getattr(state, self.num_field)
+    num = kontext.get_by_path(state, self.num_field)
     v1 = self._maybe_truncate(v1, num)
     v2 = self._maybe_truncate(v2, num)
     if v1 is EMPTY and v2 is EMPTY:

kauldron/metrics/base_state.py

@@ -83,7 +83,9 @@ class State(abc.ABC, Generic[_MetricT]):
   """
 
   _: dataclasses.KW_ONLY
-  parent: _MetricT = flax.struct.field(pytree_node=False, default=EMPTY)
+  parent: _MetricT = flax.struct.field(
+      pytree_node=False, default=EMPTY
+  )  # pytype: disable=annotation-type-mismatch
 
   def __init_subclass__(cls, **kwargs):
     super().__init_subclass__(**kwargs)