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

@@ -534,7 +534,11 @@ def key_check_fn(key, objs):
             iterator = enumerate(loss)
 
             def key_check_fn(key, objs):
-                return all([key < len(obj) for obj in objs])
+                try:
+                    [obj[key] for obj in objs]
+                except:
+                    return False
+                return True
 
         else:
             raise TypeError(
@@ -547,7 +551,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 +669,29 @@ 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]
             try:
-                y_true = tree.pack_sequence_as(y_pred, y_true)
+                # Check case where 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]
+                try:
+                    y_true = tree.pack_sequence_as(y_pred, y_true)
+                except:
+                    # Check case where y_true has the same structure but uses
+                    # different (but reconcilable) container types,
+                    # e.g `list` vs `tuple`.
+                    tree.assert_same_paths(y_true, y_pred)
+                    y_true = tree.pack_sequence_as(y_pred, tree.flatten(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,36 @@ 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):
+        y1, y2, y3 = np.array([[1]]), np.array([[2]]), np.array([[3]])
+        y_true = ([{"a": y1}, {"b": ([y2], y3)}],)
+        y_pred = [({"a": y1}, {"b": [(y2,), y3]})]
+        loss = "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)