Clarification: new mean option is synonym with sum_over_batch_size in loss function base class (#20352)

* mean as a synonym of sum_over_batch_size

* mean synonym with sum_over_batch_size

* typo fix

* Update loss.py

Undid ruff format while keeping typo fixes

* reduction in ("mean", "sum_over_batch_size") loss.py

Author: jm-willy

keras/src/losses/loss.py

@@ -17,7 +17,8 @@ class Loss(KerasSaveable):
     Args:
         reduction: Type of reduction to apply to the loss. In almost all cases
             this should be `"sum_over_batch_size"`.
-            Supported options are `"sum"`, `"sum_over_batch_size"` or `None`.
+            Supported options are `"sum"`, `"sum_over_batch_size"`, `"mean"`
+            or `None`.
         name: Optional name for the loss instance.
         dtype: The dtype of the loss's computations. Defaults to `None`, which
             means using `keras.backend.floatx()`. `keras.backend.floatx()` is a
@@ -95,7 +96,7 @@ def _obj_type(self):
 
 
 def standardize_reduction(reduction):
-    allowed = {"sum_over_batch_size", "sum", None, "none"}
+    allowed = {"sum_over_batch_size", "sum", None, "none", "mean"}
     if reduction not in allowed:
         raise ValueError(
             "Invalid value for argument `reduction`. "
@@ -135,7 +136,7 @@ def reduce_values(values, reduction="sum_over_batch_size"):
     ):
         return values
     loss = ops.sum(values)
-    if reduction == "sum_over_batch_size":
+    if reduction in ("mean", "sum_over_batch_size"):
         loss /= ops.cast(
             ops.prod(ops.convert_to_tensor(ops.shape(values), dtype="int32")),
             loss.dtype,
@@ -180,7 +181,7 @@ def apply_mask(sample_weight, mask, dtype, reduction):
     """Applies any mask on predictions to sample weights."""
     if mask is not None:
         mask = ops.cast(mask, dtype=dtype)
-        if reduction == "sum_over_batch_size":
+        if reduction in ("mean", "sum_over_batch_size"):
             # Valid entries have weight `total/valid`, while invalid ones
             # have 0. When summed over batch, they will be reduced to:
             #

keras/src/losses/loss_test.py

@@ -69,7 +69,7 @@ def test_reduction(self):
         self.assertEqual(backend.standardize_dtype(loss.dtype), "float32")
         self.assertAllClose(np.sum((y_true - y_pred) ** 2), loss)
 
-        # sum_over_batch_size
+        # sum_over_batch_size or mean
         loss_fn = ExampleLoss(reduction="sum_over_batch_size")
         loss = loss_fn(y_true, y_pred)
         self.assertEqual(backend.standardize_dtype(loss.dtype), "float32")