def process_metrics(
    metrics: Dict[str, Dict[str, Array]],
    prefix: Optional[str] = None,
) -> Dict[str, Union[float, int]]:
  """Flatten 2-level dictionary of metrics and divide values by denominator."""
  processed_metrics = {}
  for group_key, group_value in metrics.items():
    denom = group_value.pop('denominator')
    for metric_key, metric_value in group_value.items():
      flattened_key = group_key + '_' + metric_key
      if prefix is not None:
        flattened_key = prefix + '/' + flattened_key
      processed_metrics[flattened_key] = metric_value / denom
    denom_key = group_key + '_denom'
    if prefix is not None:
      denom_key = prefix + '/' + denom_key
    processed_metrics[denom_key] = denom

  return processed_metrics  # pytype: disable=bad-return-type  # jax-ndarray
Flatten 2-level dictionary of metrics and divide values by denominator.
language/mentionmemory/utils/metric_utils.py