Repository: keras-team/keras (python)

Issue:
CompileLoss should use non loss weighted metric.update_state
Proposing that we should refactor `CompileLoss` to report non loss weighted metric as was done in `keras 2`

here is the suggested refactor:

```
diff --git a/keras/src/trainers/compile_utils.py b/keras/src/trainers/compile_utils.py
index 410a782db..71d37bde3 100644
--- a/keras/src/trainers/compile_utils.py
+++ b/keras/src/trainers/compile_utils.py
@@ -658,14 +658,15 @@ class CompileLoss(losses_module.Loss):
                 value = ops.cast(
                     loss_fn(y_t, y_p, sample_weight), dtype=self.dtype
                 )
-                if loss_weight is not None:
-                    value = ops.multiply(value, loss_weight)
-                loss_values.append(value)
                 # Record individual losses.
                 if metric:
                     metric.update_state(
                         value, sample_weight=tree.flatten(y_p)[0].shape[0]
                     )
+                if loss_weight is not None:
+                    value = ops.multiply(value, loss_weight)
+                loss_values.append(value)
```                     


Pull request description:
Update CompileLoss to report unweighted metric values (breaking change)
Fixes #20343. Thanks to rivershah@ for pointing this out.

This changes CompileLoss metrics to report the values before weights are applied, which brings it in line with Keras 2 behavior.

Patch:
diff --git a/keras/src/trainers/compile_utils.py b/keras/src/trainers/compile_utils.py
index 0847df08b..f643a3bfe 100644
--- a/keras/src/trainers/compile_utils.py
+++ b/keras/src/trainers/compile_utils.py
@@ -447,6 +447,7 @@ class CompileLoss(losses_module.Loss):
         )
         self._flat_losses = None
         self._y_pred_build_structure = None
+        self._y_true_build_structure = None
 
     @property
     def metrics(self):
@@ -512,10 +513,6 @@ class CompileLoss(losses_module.Loss):
                 y_true, y_pred, loss, output_names, current_path
             )
 
-        # At this point loss should match the structure of y_pred and y_true
-        # We assume y_pred and y_true have the same structure and this
-        # have been already asserted.
-
         if not isinstance(loss, type(y_pred)):
             raise KeyError(
                 f"The path: {current_path} in "
@@ -528,17 +525,20 @@ class CompileLoss(losses_module.Loss):
             iterator = loss.items()
 
             def key_check_fn(key, objs):
-                return all([key in obj for obj in objs])
+                return all(
+                    [isinstance(obj, dict) and key in obj for obj in objs]
+                )
 
         elif issubclass(type(loss), (list, tuple)):
             iterator = enumerate(loss)
 
             def key_check_fn(key, objs):
-                try:
-                    [obj[key] for obj in objs]
-                except:
-                    return False
-                return True
+                return all(
+                    [
+                        issubclass(type(obj), (list, tuple)) and key < len(obj)
+                        for obj in objs
+                    ]
+                )
 
         else:
             raise TypeError(
@@ -547,6 +547,8 @@ class CompileLoss(losses_module.Loss):
             )
 
         for key, _loss in iterator:
+            if _loss is None:
+                continue
             if not key_check_fn(key, (y_true, y_pred)):
                 raise KeyError(
                     f"The path: {current_path + (key,)} in "
@@ -570,6 +572,11 @@ class CompileLoss(losses_module.Loss):
 
         # Pytree leaf container
         class WeightedLoss:
+            def __new__(cls, loss, weight):
+                if loss is None:
+                    return None
+                return object.__new__(cls)
+
             def __init__(self, loss, weight):
                 self.loss = loss
                 self.weight = weight
@@ -646,6 +653,9 @@ class CompileLoss(losses_module.Loss):
         self._y_pred_build_structure = tree.map_structure(
             lambda x: None, y_pred
         )
+        self._y_true_build_structure = tree.map_structure(
+            lambda x: None, y_true
+        )
         self.built = True
 
     def _get_y_pred_output_names(self, y_pred):
@@ -683,8 +693,22 @@ class CompileLoss(losses_module.Loss):
                     # 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))
+                    try:
+                        tree.assert_same_paths(y_true, y_pred)
+                        y_true = tree.pack_sequence_as(
+                            y_pred, tree.flatten(y_true)
+                        )
+                    except:
+                        # Check case where loss is partially defined over y_pred
+                        flat_y_true = tree.flatten(y_true)
+                        flat_loss = tree.flatten(self._user_loss)
+                        flat_loss_non_nones = [
+                            loss for loss in flat_loss if loss is not None
+                        ]
+                        assert len(flat_y_true) == len(flat_loss_non_nones)
+                        if not tree.is_nested(y_true):
+                            y_true = flat_y_true
+
             except:
                 y_true_struct = tree.map_structure(lambda _: "*", y_true)
                 y_pred_struct = tree.map_structure(lambda _: "*", y_pred)
@@ -705,8 +729,13 @@ class CompileLoss(losses_module.Loss):
             y_pred = tree.pack_sequence_as(
                 self._y_pred_build_structure, tree.flatten(y_pred)
             )
+        try:
+            tree.assert_same_structure(
+                self._y_true_build_structure, y_true, check_types=False
+            )
+        except ValueError:
             y_true = tree.pack_sequence_as(
-                self._y_pred_build_structure, tree.flatten(y_true)
+                self._y_true_build_structure, tree.flatten(y_true)
             )
 
         # We need to add a dummy `None` if the model has only a single output.
@@ -731,14 +760,14 @@ class CompileLoss(losses_module.Loss):
             value = ops.cast(
                 loss_fn(y_t, y_p, _sample_weight), dtype=self.dtype
             )
-            if loss_weight is not None:
-                value = ops.multiply(value, loss_weight)
-            loss_values.append(value)
-            # Record individual losses.
+            # Record *unweighted* individual losses.
             if metric:
                 metric.update_state(
                     value, sample_weight=tree.flatten(y_p)[0].shape[0]
                 )
+            if loss_weight is not None:
+                value = ops.multiply(value, loss_weight)
+            loss_values.append(value)
 
         if loss_values:
             total_loss = sum(loss_values)


Test patch:
diff --git a/keras/src/trainers/compile_utils_test.py b/keras/src/trainers/compile_utils_test.py
index 4bc2a2f4d..554cf3e2c 100644
--- a/keras/src/trainers/compile_utils_test.py
+++ b/keras/src/trainers/compile_utils_test.py
@@ -375,6 +375,33 @@ class TestCompileLoss(testing.TestCase):
         value = compile_loss(y_true, y_pred)
         self.assertAllClose(value, 1.07666, atol=1e-5)
 
+    def test_struct_loss_valid_weights(self):
+        y_true = {
+            "a": np.array([1, 2]),
+            "b": np.array([1, 2]),
+        }
+        y_pred = {
+            "a": np.array([3, 4]),
+            "b": np.array([3, 4]),
+        }
+        loss = {"a": "mse", "b": "mse"}
+        compile_loss = CompileLoss(
+            loss=loss,
+            output_names=["a", "b"],
+            loss_weights={
+                "a": np.ones((2,)),
+                "b": np.zeros((2,)),
+            },
+        )
+        compile_loss.build(y_true, y_pred)
+        value = compile_loss(y_true, y_pred)
+        self.assertAllClose(value, 4)
+
+        # Metrics still report unweighted loss.
+        a_loss_mean, b_loss_mean = compile_loss.metrics
+        self.assertEqual(a_loss_mean.result(), 4)
+        self.assertEqual(b_loss_mean.result(), 4)
+
     def test_struct_loss_invalid_weights(self):
         y_true = {
             "a": {
@@ -520,3 +547,18 @@ class TestCompileLoss(testing.TestCase):
         ):
             wrong_struc_y_true = [np.array([[1]])]
             compile_loss(wrong_struc_y_true, y_pred)
+
+    def test_y_true_partial_y_pred_span(self):
+        ones = np.ones((320, 3))
+        zeros = np.zeros((320, 3))
+        y_pred = [ones, zeros, zeros]
+        y_true = ones
+
+        compile_loss = CompileLoss(
+            loss=["mse", None, None], output_names=["a", "b", "c"]
+        )
+        # build call
+        compile_loss(y_true, y_pred)
+        # built call
+        loss = compile_loss(y_true, y_pred)
+        self.assertEqual(loss, 0.0)
diff --git a/keras/src/trainers/trainer_test.py b/keras/src/trainers/trainer_test.py
index edf1555be..434fe47c9 100644
--- a/keras/src/trainers/trainer_test.py
+++ b/keras/src/trainers/trainer_test.py
@@ -2544,6 +2544,23 @@ class TestTrainer(testing.TestCase):
             atol=1e-3,
         )
 
+    @pytest.mark.requires_trainable_backend
+    def test_partial_loss_partial_label(self):
+        inputs = keras.Input((2,))
+        x = keras.layers.Dense(3, kernel_initializer="ones")(inputs)
+        partial_model = keras.Model(inputs, [x, x, x])
+        partial_model.compile(loss=["mse", None, None])
+        full_model = keras.Model(inputs, [x, x, x])
+        full_model.compile(loss=["mse", "mse", "mse"])
+
+        eval_x = np.ones((32, 2))
+        eval_y = np.ones((32, 3))
+
+        partial_logs = partial_model.evaluate(eval_x, eval_y, return_dict=True)
+        logs = full_model.evaluate(eval_x, [eval_y] * 3, return_dict=True)
+
+        self.assertAlmostEqual(partial_logs["loss"] * 3, logs["loss"])
+
     def test_symbolic_build(self):
         class ExampleModelWithTrainingArgs(Trainer, layers.Layer):
             def __init__(self, units):