CompileLoss: Allow different but reconcilable structures for y_true and y_pred

keras/src/models/model_test.py

@@ -727,7 +727,8 @@ def test_functional_dict_outputs_dict_losses_invalid_keys(self):
         # Fit the model to make sure compile_metrics are built
         with self.assertRaisesRegex(
             KeyError,
-            "in the `loss` argument, can't be found in the model's output",
+            "in the `loss` argument, can't be found "
+            "in either the model's output",
         ):
             model.fit(x, (y1, y2), batch_size=2, epochs=1, verbose=0)
 

keras/src/trainers/compile_utils.py

@@ -547,7 +547,8 @@ def key_check_fn(key, objs):
                 raise KeyError(
                     f"The path: {current_path + (key,)} in "
                     "the `loss` argument, can't be found in "
-                    "the model's output (`y_pred`)."
+                    "either the model's output (`y_pred`) or in the "
+                    "labels (`y_true`)."
                 )
 
             self._build_nested(
@@ -664,20 +665,20 @@ def call(self, y_true, y_pred, sample_weight=None):
         try:
             tree.assert_same_structure(y_pred, y_true, check_types=False)
         except ValueError:
-            # y_true is either flat or leaf
-            if (
-                not tree.is_nested(y_true)
-                and hasattr(y_pred, "__len__")
-                and len(y_pred) == 1
-            ):
-                y_true = [y_true]
+            # y_true is either:
+            #   - flat or leaf
+            #   - has the same structure but uses different (but reconcilable)
+            #     container types, e.g `list` vs `tuple`
+            flat_y_true = tree.flatten(y_true)
             try:
-                y_true = tree.pack_sequence_as(y_pred, y_true)
+                y_true = tree.pack_sequence_as(y_pred, flat_y_true)
             except:
+                y_true_struct = tree.map_structure(lambda _: "*", y_true)
+                y_pred_struct = tree.map_structure(lambda _: "*", y_pred)
                 raise ValueError(
                     "y_true and y_pred have different structures.\n"
-                    f"y_true: {y_true}\n"
-                    f"y_pred: {y_pred}\n"
+                    f"y_true: {y_true_struct}\n"
+                    f"y_pred: {y_pred_struct}\n"
                 )
 
         if not self.built:

keras/src/trainers/compile_utils_test.py

@@ -490,3 +490,35 @@ def test_struct_loss_invalid_path(self):
             KeyError, "can't be found in the model's output"
         ):
             compile_loss.build(y_true_symb, y_pred_symb)
+
+    def test_different_container_types(self):
+        y_true = [[np.array([[1]]), np.array([[2]]), np.array([[3]])]]
+        y_pred = [(np.array([[1]]), np.array([[2]]), np.array([[3]]))]
+        loss = [["mse", "mse", "mse"]]
+        compile_loss = CompileLoss(loss=loss, output_names=["a", "b", "c"])
+        y_true_symb = tree.map_structure(
+            lambda _: backend.KerasTensor((1, 1)), y_true
+        )
+        y_pred_symb = tree.map_structure(
+            lambda _: backend.KerasTensor((1, 1)), y_pred
+        )
+        compile_loss.build(y_true_symb, y_pred_symb)
+        compile_loss(y_true, y_pred)
+
+    def test_structure_mismatch(self):
+        y_true = [np.array([[1]]), np.array([[1]])]
+        y_pred = [np.array([[1]]), np.array([[1]])]
+        loss = ["mse", "mse"]
+        compile_loss = CompileLoss(loss=loss, output_names=["a", "b"])
+        y_true_symb = tree.map_structure(
+            lambda _: backend.KerasTensor((1, 1)), y_true
+        )
+        y_pred_symb = tree.map_structure(
+            lambda _: backend.KerasTensor((1, 1)), y_pred
+        )
+        compile_loss.build(y_true_symb, y_pred_symb)
+        with self.assertRaisesRegex(
+            ValueError, "y_true and y_pred have different structures."
+        ):
+            wrong_struc_y_true = [np.array([[1]])]
+            compile_loss(wrong_struc_y_true, y_pred)