Ensures compiled metrics resolve output names correctly (#21694)

* Ensures compiled metrics resolve output names correctly

* improves readability

Author: JyotinderSingh

keras/src/trainers/compile_utils.py

@@ -148,6 +148,7 @@ def __init__(
         self.built = False
         self.name = "compile_metrics"
         self.output_names = output_names
+        self._resolved_output_names = None
 
     @property
     def metrics(self):
@@ -175,10 +176,16 @@ def variables(self):
 
     def build(self, y_true, y_pred):
         num_outputs = 1  # default
-        if self.output_names:
+        # Resolve output names. If y_pred is a dict, prefer its keys.
+        if isinstance(y_pred, dict):
+            keys = sorted(list(y_pred.keys()))
+            if self.output_names and set(self.output_names) == set(keys):
+                # If there is a perfect match, use the user-provided order.
+                output_names = self.output_names
+            else:
+                output_names = keys
+        elif self.output_names:
             output_names = self.output_names
-        elif isinstance(y_pred, dict):
-            output_names = sorted(list(y_pred.keys()))
         elif isinstance(y_pred, (list, tuple)):
             num_outputs = len(y_pred)
             if all(hasattr(x, "_keras_history") for x in y_pred):
@@ -187,6 +194,7 @@ def build(self, y_true, y_pred):
                 output_names = None
         else:
             output_names = None
+        self._resolved_output_names = output_names
         if output_names:
             num_outputs = len(output_names)
 
@@ -316,9 +324,10 @@ def _build_metrics_set(
         return flat_metrics
 
     def _flatten_y(self, y):
-        if isinstance(y, dict) and self.output_names:
+        names = self._resolved_output_names
+        if isinstance(y, dict) and names:
             result = []
-            for name in self.output_names:
+            for name in names:
                 if name in y:
                     result.append(y[name])
             return result

keras/src/trainers/compile_utils_test.py

@@ -235,6 +235,57 @@ def my_custom_metric(y_true, y_pred):
         self.assertEqual(len(result), 1)
         self.assertTrue("my_custom_metric" in result)
 
+    def test_dict_outputs_ignore_mismatched_output_names(self):
+        """Tests that when output_names does not match dict keys, the correct
+        keys are used."""
+
+        # output_names represent internal op names that do not match dict keys.
+        compile_metrics = CompileMetrics(
+            metrics={
+                "a": metrics_module.MeanSquaredError(),
+                "b": metrics_module.MeanSquaredError(),
+            },
+            weighted_metrics=None,
+            output_names=["dense", "dense_1"],
+        )
+
+        # Symbolic build with dict outputs keyed by user-facing names.
+        y_true = {
+            "a": backend.KerasTensor((3, 2)),
+            "b": backend.KerasTensor((3, 2)),
+        }
+        y_pred = {
+            "a": backend.KerasTensor((3, 2)),
+            "b": backend.KerasTensor((3, 2)),
+        }
+
+        # The build method should correctly map metrics for outputs 'a' and 'b',
+        # even when the op names do not match.
+        compile_metrics.build(y_true, y_pred)
+
+        # Make the two outputs produce different MSEs to verify mapping.
+        y_true = {
+            "a": np.zeros((3, 2), dtype="float32"),
+            "b": np.zeros((3, 2), dtype="float32"),
+        }
+        y_pred = {
+            # MSE(a) = 0.0
+            "a": np.zeros((3, 2), dtype="float32"),
+            # MSE(b) = 1.0
+            "b": np.ones((3, 2), dtype="float32"),
+        }
+        compile_metrics.update_state(y_true, y_pred)
+
+        result = compile_metrics.result()
+        self.assertIsInstance(result, dict)
+
+        # Should expose metrics under the dict keys ('a', 'b'),
+        # and not the internal names.
+        self.assertIn("a_mean_squared_error", result)
+        self.assertIn("b_mean_squared_error", result)
+        self.assertAllClose(result["a_mean_squared_error"], 0.0)
+        self.assertAllClose(result["b_mean_squared_error"], 1.0, atol=1e-6)
+
 
 class TestCompileLoss(testing.TestCase):
     def test_single_output_case(self):