Pass keyword arguments to metric functions in `weighted_average_metric` wrapper.
Pass keyword arguments to metric functions in `weighted_average_metric` wrapper.

PiperOrigin-RevId: 444593232

google/init2winit