Unify behavior of self.losses in compute_loss across backends.

Author: fchollet

keras/backend/jax/trainer.py

@@ -51,6 +51,10 @@ def compute_loss_and_updates(
             return_losses=True,
             **kwargs,
         )
+        if losses:
+            # Make forward pass losses available to compute_loss.
+            self._losses_override.clear()
+            self._losses_override = losses
 
         loss, variables = self.stateless_compute_loss(
             trainable_variables,
@@ -61,14 +65,12 @@ def compute_loss_and_updates(
             y_pred=y_pred,
             sample_weight=sample_weight,
         )
+        if losses:
+            self._losses_override.clear()
         (trainable_variables, non_trainable_variables, metrics_variables) = (
             variables
         )
 
-        # Sum forward pass losses
-        if losses:
-            loss += ops.sum(losses)
-
         # Handle loss scaling
         unscaled_loss = loss
         if training and self.optimizer is not None:

keras/layers/layer.py

@@ -274,6 +274,7 @@ def __init__(
         self._trainable = trainable
         self._losses = []
         self._loss_ids = set()
+        self._losses_override = []
 
         self._call_signature = inspect.signature(self.call)
         call_signature_parameters = [
@@ -1091,6 +1092,8 @@ def _get_regularization_losses(self):
     @property
     def losses(self):
         """List of scalar losses from `add_loss`, regularizers and sublayers."""
+        if self._losses_override:
+            return self._losses_override
         losses = self._get_own_losses()
         for layer in self._flatten_layers(include_self=False):
             losses.extend(layer._get_own_losses())

keras/trainers/trainer_test.py

@@ -1319,3 +1319,29 @@ def compute_loss(
         self.assertAlmostEqual(
             history.history["custom"][0], history.history["loss"][0] * 4
         )
+
+    @pytest.mark.requires_trainable_backend
+    def test_fwd_pass_loss_presence_in_compute_loss(self):
+
+        class MyModel(keras.Model):
+            def __init__(self):
+                super().__init__()
+                self.custom_metric = keras.metrics.Mean(name="custom")
+                self.dense = keras.layers.Dense(2, activity_regularizer="l2")
+
+            def call(self, x):
+                return self.dense(x)
+
+            def compute_loss(
+                self, x=None, y=None, y_pred=None, sample_weight=None
+            ):
+                loss = super().compute_loss(x, y, y_pred, sample_weight)
+                self.custom_metric.update_state(sum(self.losses))
+                return loss
+
+        model = MyModel()
+        model.compile(optimizer="sgd", loss="mse")
+        x = np.ones((32, 4))
+        y = np.ones((32, 2)) * 2
+        history = model.fit(x, y)
+        self.assertGreater(history.history["custom"][0], 0.0)