def update_metrics_dtype(
    metrics: Dict[str, Dict[str, Any]]) -> Dict[str, Dict[str, Any]]:
  """Convert metrics to more precise types."""
  return jax.tree.map(update_value_dtype, metrics)
Convert metrics to more precise types.
language/mentionmemory/utils/metric_utils.py