Include POSITIONAL_OR_KEYWORD params in MultiMetric kwarg validation

Author: us

flax/nnx/training/metrics.py

@@ -418,24 +418,27 @@ def __init__(self, **metrics):
       if self._expected_kwargs is None:
         continue
       sig = inspect.signature(metric.update)
-      has_keyword_only = False
+      has_named_params = False
       has_var_keyword = False
-      keyword_only_names: set[str] = set()
+      named_param_names: set[str] = set()
       for pname, param in sig.parameters.items():
         if pname == 'self':
           continue
-        if param.kind == param.KEYWORD_ONLY:
-          keyword_only_names.add(pname)
-          has_keyword_only = True
+        if param.kind in (
+          param.POSITIONAL_OR_KEYWORD,
+          param.KEYWORD_ONLY,
+        ):
+          named_param_names.add(pname)
+          has_named_params = True
         elif param.kind == param.VAR_KEYWORD:
           has_var_keyword = True
-      if has_keyword_only and has_var_keyword:
+      if has_named_params and has_var_keyword:
         # Metric declares specific params but also absorbs
         # extras (e.g. Accuracy's **_); can't validate
         # without false positives.
         self._expected_kwargs = None
-      elif has_keyword_only:
-        self._expected_kwargs.update(keyword_only_names)
+      elif has_named_params:
+        self._expected_kwargs.update(named_param_names)
       elif hasattr(metric, 'argname'):
         # Use argname convention (e.g. Average, Welford).
         self._expected_kwargs.add(metric.argname)